reductions

Module for tail-sensitive tensor reductions.

A plain mean is the wrong reduction whenever a loss must react to a small subset of positions. Averaging over N positions divides any single position's contribution by N, so a handful of positions that behave badly are indistinguishable from none at all. tail_mean replaces the mean with the mean of the worst fraction of positions, which keeps the loss sensitive to that subset while staying smoother than a hard max.

Functions:

Name	Description
tail_mean	Average the largest tail_fraction of values along dim.
top_k_mean	Average the k largest entries of values along dim.

tail_mean

tail_mean(
    values: Tensor, tail_fraction: float, dim: int
) -> torch.Tensor

Average the largest tail_fraction of values along dim.

The average top-k reduction: the mean of the ceil(tail_fraction * N) largest entries (delegating to top_k_mean). tail_fraction interpolates between the two familiar reductions — 1.0 is the plain mean (short-circuited to values.mean(dim), so bit-for-bit identical rather than merely close), while a tail_fraction small enough that k == 1 is values.max(dim). A fraction rather than an explicit k is the right parameterization when the offending subset scales with N — "the worst 5%" then means the same thing at 196 or 1024 positions; prefer top_k_mean when the subset has a fixed absolute size instead.

Gradients reach only the selected positions, each with weight 1 / k; the rest receive exactly zero. The result is never below values.mean(dim), and how far above depends on how concentrated values is, so a loss weight calibrated at one tail_fraction does not necessarily transfer to another.

Padding: the reduction ranges over every position along dim — both the count k and the selected set include padding — so a caller with padded rows must exclude the padding upstream (e.g. flatten to the valid positions) before reducing.

Parameters:

Name	Type	Description	Default
values	Tensor	Tensor to reduce. Larger entries are treated as worse.	required
tail_fraction	float	Fraction of positions to retain, in (0, 1]. 1.0 reproduces values.mean(dim). The realized count is max(1, ceil(tail_fraction * values.shape[dim])) — rounded first, so an inexact fraction cannot select one position too many — and the effective fraction is therefore quantized to multiples of 1 / values.shape[dim].	required
dim	int	Dimension to reduce. Required rather than defaulted: which axis carries the positions is a property of the caller's layout, and a wrong-axis reduction returns a plausible number rather than an error.	required

Returns:

Type	Description
torch.Tensor	Tensor with dim reduced, of the same dtype and device as values.

Raises:

Type	Description
ValueError	If tail_fraction is not finite or lies outside (0, 1].

Example

values = torch.tensor([[0.0, 0.0, 0.0, 0.4, 0.8]])

1.0 is the plain mean.

tail_mean(values, 1.0, dim=-1) tensor([0.2400])

The worst 40% of five positions is the top two: (0.8 + 0.4) / 2.

tail_mean(values, 0.4, dim=-1) tensor([0.6000])

Small enough that k == 1, so the reduction is the maximum.

tail_mean(values, 0.1, dim=-1) tensor([0.8000])

Reduction is per-row, so one row cannot absorb another row's tail.

tail_mean(torch.tensor([[0.0, 1.0], [0.2, 0.2]]), 0.5, dim=-1) tensor([1.0000, 0.2000])

Added in version v3.71.0. Tail-sensitive replacement for `mean`: averages the worst `tail_fraction` of positions along a dimension.

top_k_mean

top_k_mean(
    values: Tensor, k: int, dim: int
) -> torch.Tensor

Average the k largest entries of values along dim.

The count-parameterized form of tail_mean and the primitive it delegates to. Prefer this when the offending subset has a fixed absolute size rather than a fixed share of the positions — the count is then the invariant, and going through tail_mean's k / N does not round-trip in binary floating point. k is clamped to the number of positions, so k >= N selects everything (the plain mean) and short rows need no special-casing. See tail_mean for the reduction's properties (the equal 1 / k gradient share) and the same padding caveat: the top-k ranges over every position along dim, so a caller with padded rows must exclude the padding upstream.

Parameters:

Name	Type	Description	Default
values	Tensor	Tensor to reduce. Larger entries are treated as worse.	required
k	int	Number of largest entries to average, clamped to values.shape[dim]. Must be positive.	required
dim	int	Dimension to reduce.	required

Returns:

Type	Description
torch.Tensor	Tensor with dim reduced, of the same dtype and device as values.

Raises:

Type	Description
ValueError	If k is not positive.

Example

values = torch.tensor([[0.0, 0.0, 0.0, 0.4, 0.8]]) top_k_mean(values, 2, dim=-1) tensor([0.6000])

k at or above the position count selects everything.

top_k_mean(values, 99, dim=-1) tensor([0.2400])

Reduction is per-row, so one row cannot absorb another row's tail.

top_k_mean(torch.tensor([[0.0, 1.0], [0.2, 0.2]]), 1, dim=-1) tensor([1.0000, 0.2000])

Added in version v3.71.0. Count-parameterized form of `tail_mean`, for callers whose offending subset has a fixed absolute size rather than a fixed share of the positions.