divergences

Module for f-divergence based loss functions.

Functions:

Name	Description
`fused_linear_masked_generalized_jsd`	Compute a masked generalized Jensen-Shannon divergence between student and teacher token distributions directly from hidden states.
`jefferys_divergence`	Compute the Jefferys divergence between two discrete probabilities using logits.
`jensen_shannon_divergence`	Compute the Jensen Shannon divergence between two discrete probabilities using logits.
`masked_cross_entropy`	Compute cross-entropy loss between input logits and target logits.
`masked_jefferys_divergence`	Compute Jefferys divergence between input logits and target logits.
`masked_kl_divergence`	Compute KL divergence between input logits and target logits.
`masked_unbiased_dcor`	Compute an unbiased approximation to the distance correlation between the tensors, representing samples of random variables.
`squared_hellinger_distance`	Compute the Squared Hellinger distance between two discrete probabilities using logits.
`temperature_scaled_masked_jefferys_divergence`	Compute symmetric Jefferys divergence over masked positions: `KL(t\|\|s) + KL(s\|\|t)`.
`temperature_scaled_masked_kl_divergence`	Compute temperature-scaled `KL(softmax(teacher / T) \|\| softmax(student / T))` over masked positions.
`total_variation`	Compute the total variation between two discrete probabilities using logits.

fused_linear_masked_generalized_jsd ¶

fused_linear_masked_generalized_jsd(student_hidden: Tensor, teacher_hidden: Tensor, lm_head_weight: Tensor, attention_mask: Tensor | None = None, *, jsd_beta: float, temperature: float = 1.0, reduction: Literal['sequence_mean', 'token_mean'] = 'sequence_mean', chunk_size: int = 1024) -> <class 'torch.Tensor'>

Compute a masked generalized Jensen-Shannon divergence between student and teacher token distributions directly from hidden states.

Instead of consuming pre-computed logits, this takes the student and teacher hidden states plus the shared lm_head weight and fuses the output projection with the divergence. The full (B, T, V) logits are therefore never materialized — at the jsd_beta endpoints on CUDA via Liger's LigerFusedLinearJSD kernel, and otherwise via an exact pure-PyTorch computation chunked over tokens.

jsd_beta selects a member of the divergence family, with the two endpoints special-cased to the KL divergences:

jsd_beta=0.0 — forward KL KL(teacher || student), equivalent to masked_kl_divergence on the projected logits.
jsd_beta=1.0 — reverse KL KL(student || teacher). Summing the two endpoints with equal weight recovers Jeffrey's divergence; with teacher clean and student noisy, they are the distillation forward and reverse KLD terms.
0 < jsd_beta < 1 — the generalized Jensen-Shannon divergence beta * KL(teacher || m) + (1 - beta) * KL(student || m) with m = beta * teacher + (1 - beta) * student. At jsd_beta=0.5 this is the true Jensen-Shannon divergence, equivalent to jensen_shannon_divergence on the projected logits (with no support_mask), and like it bounded above by log(2).

Note

The family is discontinuous at both endpoints, by design and in agreement with Liger. The endpoint values are not the limits of the interior expression: at jsd_beta=0.0 and jsd_beta=1.0 the mixture m equals one of the two distributions, so the generalized value collapses to 0 there. Both endpoints are special-cased to the KL divergences listed above instead.

The teacher side is always detached, so gradients flow only through student_hidden and lm_head_weight — the standard distillation routing with a frozen teacher.

Note

The fused-linear memory saving on the backward pass requires the Liger kernel, which runs only on CUDA tensors and only at jsd_beta=0.0/1.0 — its interior-beta branch forms the mixture in probability space and returns NaN once an entry underflows under both distributions, so interior betas take the fallback on every device (see https://github.com/linkedin/Liger-Kernel/issues/1453). The fallback is exact and bounds the forward working set but, like any autograd graph, retains per-chunk activations for the backward pass.

Parameters:

Name	Type	Description	Default
`student_hidden` ¶	`Tensor`	Student (e.g. noisy) hidden states of shape `(batch_size, sequence_length, hidden_dim)`.	required
`teacher_hidden` ¶	`Tensor`	Teacher (e.g. clean) hidden states of shape `(batch_size, sequence_length, hidden_dim)`.	required
`lm_head_weight` ¶	`Tensor`	The output projection weight of shape `(vocab_size, hidden_dim)`. Assumed bias-free (as for Llama-family heads).	required
`attention_mask` ¶	`Tensor \| None`	Optional mask of shape `(batch_size, sequence_length)`; `True`/non-zero positions are included. When omitted, every position is treated as valid.	`None`
`jsd_beta` ¶	`float`	Mixture weight on the teacher distribution, in the closed interval `[0, 1]`. `0.0` is the forward KL `KL(teacher \|\| student)` and `1.0` is the reverse KL `KL(student \|\| teacher)`; any interior value is the generalized Jensen-Shannon divergence, with `0.5` the true Jensen-Shannon divergence.	required
`temperature` ¶	`float`	Softmax temperature applied to both sides. Defaults to `1.0`.	`1.0`
`reduction` ¶	`Literal['sequence_mean', 'token_mean']`	`"sequence_mean"` (per-sequence masked mean then batch mean, the default; matches the logit-level masked divergences) or `"token_mean"` (a single global mean over all valid tokens).	`'sequence_mean'`
`chunk_size` ¶	`int`	Token-block size for the pure-PyTorch fallback. Ignored by the Liger path, which chunks internally. Defaults to `1024`.	`1024`

Returns:

Type	Description
`<class 'torch.Tensor'>`	Scalar `float32` divergence.

Raises:

Type	Description
`ValueError`	If `jsd_beta` is outside `[0.0, 1.0]`, if `reduction` is unsupported, or if `temperature` is not strictly positive.

Example

student_hidden = torch.randn(2, 5, 4) teacher_hidden = torch.randn(2, 5, 4) lm_head_weight = torch.randn(10, 4) attention_mask = torch.ones(2, 5, dtype=torch.bool) forward_kld = fused_linear_masked_generalized_jsd( ... student_hidden, teacher_hidden, lm_head_weight, attention_mask, jsd_beta=0.0 ... ) print(forward_kld.shape) torch.Size([]) jsd = fused_linear_masked_generalized_jsd( ... student_hidden, teacher_hidden, lm_head_weight, attention_mask, jsd_beta=0.5 ... ) bool(jsd <= math.log(2)) True

Changed in version v3.73.1: `jsd_beta` now accepts the whole closed interval `[0, 1]` rather than only `0.0` and `1.0`. Interior values give the generalized Jensen-Shannon divergence (`0.5` is the true Jensen-Shannon divergence), matching the family Liger's `LigerFusedLinearJSD` kernel already implemented. Results for `jsd_beta=0.0` and `jsd_beta=1.0` are unchanged.

Changed in version v3.73.1: Renamed from `fused_linear_masked_kl_divergence`, which was misleading once `jsd_beta` was no longer restricted to the two endpoints: the function computes the generalized Jensen-Shannon divergence, of which the forward and reverse KL divergences are special cases. No alias is kept under the old name — update call sites to `fused_linear_masked_generalized_jsd`.

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain.

Added in version v3.52.0. Hidden-states-based KL divergence that fuses the lm_head projection with the divergence so the full `(B, T, V)` logits are never materialized — for large vocabularies this is a large peak-memory saving. Uses Liger's `LigerFusedLinearJSD` on CUDA and an exact pure-PyTorch chunked fallback elsewhere.

jefferys_divergence ¶

jefferys_divergence(noisy_logits: Tensor, clean_logits: Tensor, attention_mask: Tensor, support_mask: Tensor | None = None, reduction: Literal['mean', 'none'] = 'mean') -> <class 'torch.Tensor'>

Compute the Jefferys divergence between two discrete probabilities using logits.

Note

See https://en.wikipedia.org/wiki/F-divergence#Common_examples_of_f-divergences for the implementation formula.

Note

If support_mask has non-zero elements then the divergence is taken from a masking of the original-distributions, making it no longer a divergence of true distributions.

Parameters:

Name	Type	Description	Default
`noisy_logits` ¶	`Tensor`	The logits produced from SGT data.	required
`clean_logits` ¶	`Tensor`	The logits produced from data without the SGT.	required
`attention_mask` ¶	`Tensor`	The attention mask for the batch.	required
`support_mask` ¶	`Tensor \| None`	The mask placed on both noisy and clean logits. Useful for comparing top-k logits.	`None`
`reduction` ¶	`Literal['mean', 'none']`	Specifies how to reduce across the batch dimension after computing the per-sequence masked mean.	`'mean'`

Returns:

Type	Description
`<class 'torch.Tensor'>`	The Jefferys divergence between the noisy and clean logits, independently averaged over the batch and sequence lengths.

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain.

Added in version v3.37.0.

jensen_shannon_divergence ¶

jensen_shannon_divergence(noisy_logits: Tensor, clean_logits: Tensor, attention_mask: Tensor, support_mask: Tensor | None = None, reduction: Literal['mean', 'none'] = 'mean') -> <class 'torch.Tensor'>

Compute the Jensen Shannon divergence between two discrete probabilities using logits.

Note

See https://en.wikipedia.org/wiki/F-divergence#Common_examples_of_f-divergences for the implementation formula.

Note

If support_mask has non-zero elements then the divergence is taken from a masking of the original-distributions, making it no longer a divergence of true distributions.

Note

The divergence is symmetric and bounded in [0, log(2)]. A support_mask computes a partial sum over the retained vocabulary entries; this remains bounded by log(2), but is no longer a divergence between normalized distributions.

Note

The computation runs in float32 via compute_in_precision and the result is returned in float32 rather than cast back, so low-precision logits neither propagate their coarse resolution into the divergence nor produce a -inf log-probability that would make it NaN. The cast back is disabled because reduction="none" returns one value per sequence, which the decorator's scalar overflow clamp cannot handle — and a quantity bounded by log(2) cannot overflow anyway.

Parameters:

Name	Type	Description	Default
`noisy_logits` ¶	`Tensor`	The logits produced from SGT data.	required
`clean_logits` ¶	`Tensor`	The logits produced from data without the SGT.	required
`attention_mask` ¶	`Tensor`	The attention mask for the batch.	required
`support_mask` ¶	`Tensor \| None`	The mask placed on both noisy and clean logits. Useful for comparing top-k logits.	`None`
`reduction` ¶	`Literal['mean', 'none']`	Specifies how to reduce across the batch dimension after computing the per-sequence masked mean.	`'mean'`

Returns:

Type	Description
`<class 'torch.Tensor'>`	The Jensen Shannon divergence between the noisy and clean logits, independently averaged over the batch and sequence lengths.

Changed in version v3.73.1: Computed in `float32` via `compute_in_precision`, and returned in `float32`. In float16 the log-probability of a sufficiently separated vocabulary entry overflows to `-inf`, and the `0 * -inf` term that follows poisons the divergence with `NaN`; short of that, a bfloat16 result cannot resolve the `log(2)` bound to three significant digits.

Changed in version v3.73.1: The mixture distribution was previously passed to `torch.nn.functional.kl_div` as the target rather than the input, so the function computed `0.5 * [KL(m || p) + KL(m || q)]` instead of the Jensen-Shannon divergence `0.5 * [KL(p || m) + KL(q || m)]`. That quantity is unbounded and grows without limit as either branch becomes confident, whereas the Jensen-Shannon divergence is bounded above by `log(2)`. Values returned before and after this change are not comparable.

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain.

Added in version v3.37.0.

masked_cross_entropy ¶

masked_cross_entropy(input_logits: Tensor, target_logits: Tensor, attention_mask: Tensor | None = None, max_loss: float | None = None, scaling_factor: float = 1.0) -> <class 'torch.Tensor'>

Compute cross-entropy loss between input logits and target logits.

Applies softmax to the target logits. The computation is performed in float32 (the logits are upcast on entry) and a float32 scalar is returned, so low-precision (bfloat16/float16) logits do not propagate their coarse scalar resolution into the loss.

Parameters:

Name	Type	Description	Default
`input_logits` ¶	`Tensor`	Logits from the model, shape (batch_size, sequence_length, embedding_dim)	required
`target_logits` ¶	`Tensor`	Logits from the target, shape (batch_size, sequence_length, embedding_dim)	required
`attention_mask` ¶	`Tensor \| None`	Optional attention mask, shape (batch_size, sequence_length). When omitted, every position is treated as valid.	`None`
`max_loss` ¶	`float \| None`	The maximum value for the loss to cap the cross entropy loss.	`None`
`scaling_factor` ¶	`float`	A scaling factor for the sigmoid cross-entropy loss, lower values will slow the speed of optimization. Defaults to 1.0.	`1.0`

Returns:

Type	Description
`<class 'torch.Tensor'>`	Cross-entropy loss, averaged over the batch and sequence length.

Example

input_logits = torch.randn(2, 5, 10) target_logits = torch.randn(2, 5, 10) attention_mask = torch.ones(2, 5, dtype=torch.bool) cross_entropy_loss = masked_cross_entropy( ... input_logits, target_logits, attention_mask ... ) print(cross_entropy_loss.shape) torch.Size([])

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain.

Added in version v1.8.0.

Added in version v2.18.0.

Changed in version v3.52.0: Computes the cross entropy in `float32` (returning a `float32` scalar) so bfloat16/float16 logits no longer inherit a coarse low-precision result, and reduces NaN-safely so a fully-masked sequence contributes `0` instead of propagating `NaN`.

masked_jefferys_divergence ¶

masked_jefferys_divergence(input_logits: Tensor, target_logits: Tensor, attention_mask: Tensor | None = None, log_target: bool = False) -> <class 'torch.Tensor'>

Compute Jefferys divergence between input logits and target logits.

Parameters:

Name	Type	Description	Default
`input_logits` ¶	`Tensor`	Logits from the model, shape (batch_size, sequence_length, embedding_dim)	required
`target_logits` ¶	`Tensor`	Logits from the target, shape (batch_size, sequence_length, embedding_dim)	required
`attention_mask` ¶	`Tensor \| None`	Optional attention mask, shape (batch_size, sequence_length	`None`
`log_target` ¶	`bool`	Whether the target logits are already in log space.	`False`

Returns:

Type	Description
`<class 'torch.Tensor'>`	torch.Tensor: Jefferys divergence, averaged over the batch and sequence length.

Example

input_logits = torch.randn(2, 5, 10) target_logits = torch.randn(2, 5, 10) attention_mask = torch.ones(2, 5, dtype=torch.bool) jefferys_div = masked_jefferys_divergence( ... input_logits, target_logits, attention_mask, log_target=False ... ) print(jefferys_div.shape) torch.Size([])

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain. Both the forward and reverse [`masked_kl_divergence`][stainedglass_core.loss.divergences.masked_kl_divergence] terms summed here inherit the new reduction.

masked_kl_divergence ¶

masked_kl_divergence(input_logits: Tensor, target_logits: Tensor, attention_mask: Tensor | None = None, log_target: bool = True) -> <class 'torch.Tensor'>

Compute KL divergence between input logits and target logits.

The computation is performed in float32 (the logits are upcast on entry) and a float32 scalar is returned, so low-precision (bfloat16/float16) logits do not propagate their coarse scalar resolution into the loss.

Parameters:

Name	Type	Description	Default
`input_logits` ¶	`Tensor`	Logits from the model, shape (batch_size, sequence_length, embedding_dim)	required
`target_logits` ¶	`Tensor`	Logits from the target, shape (batch_size, sequence_length, embedding_dim)	required
`attention_mask` ¶	`Tensor \| None`	Optional attention mask, shape (batch_size, sequence_length). When omitted, every position is treated as valid.	`None`
`log_target` ¶	`bool`	Whether the target logits are already in log space.	`True`

Returns:

Type	Description
`<class 'torch.Tensor'>`	torch.Tensor: KL divergence, averaged over the batch and sequence length.

Example

input_logits = torch.randn(2, 5, 10) target_logits = torch.randn(2, 5, 10) attention_mask = torch.ones(2, 5, dtype=torch.bool) kl_div = masked_kl_divergence(input_logits, target_logits, attention_mask) print(kl_div.shape) torch.Size([])

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain.

Added in version v1.8.0.

Changed in version v3.52.0: Computes the KL divergence in `float32` (returning a `float32` scalar) so bfloat16/float16 logits no longer inherit a coarse low-precision result, and reduces NaN-safely so a fully-masked sequence contributes `0` instead of propagating `NaN`.

masked_unbiased_dcor ¶

masked_unbiased_dcor(samples_1: Tensor, samples_2: Tensor, attention_mask: Tensor, safety_factor: float = 100.0) -> <class 'torch.Tensor'>

Compute an unbiased approximation to the distance correlation between the tensors, representing samples of random variables.

Note

The approximation assumes the last last tensorial dimension is the random vector, and all preceding dimensions represent the different observations. In the case of text, this corresponds to a token level distance correlation calculation.

Note

The tensors must have the same number of rows, representing the number of samples.

See https://arxiv.org/pdf/1701.06054.pdf for more details.

Parameters:

Name	Type	Description	Default
`samples_1` ¶	`Tensor`	The first tensor of samples.	required
`samples_2` ¶	`Tensor`	The second tensor of samples.	required
`attention_mask` ¶	`Tensor`	An attention mask indicating valid samples.	required
`safety_factor` ¶	`float`	A factor to scale the added epsilon for numerical stability.	`100.0`

Returns:

Type	Description
`<class 'torch.Tensor'>`	An approximation to the distance correlation between 0 and 1.

Added in version v3.10.0.

squared_hellinger_distance ¶

squared_hellinger_distance(noisy_logits: Tensor, clean_logits: Tensor, attention_mask: Tensor, support_mask: Tensor | None = None, reduction: Literal['mean', 'none'] = 'mean') -> <class 'torch.Tensor'>

Compute the Squared Hellinger distance between two discrete probabilities using logits.

Computes the divergence between embeddings and averages across the sequence.

Note

See https://en.wikipedia.org/wiki/F-divergence#Common_examples_of_f-divergences for the implementation formula.

Parameters:

Name	Type	Description	Default
`noisy_logits` ¶	`Tensor`	The logits produced from SGT data.	required
`clean_logits` ¶	`Tensor`	The logits produced from data without the SGT.	required
`attention_mask` ¶	`Tensor`	The attention mask for the batch.	required
`support_mask` ¶	`Tensor \| None`	The mask placed on both noisy and clean logits. Useful for comparing top-k logits.	`None`
`reduction` ¶	`Literal['mean', 'none']`	Specifies how to reduce across the batch dimension after computing the per-sequence masked mean.	`'mean'`

Returns:

Type	Description
`<class 'torch.Tensor'>`	The squared Hellinger distance between the noisy and clean logits, independently averaged over the batch and sequence lengths.

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain.

Added in version v3.37.0.

temperature_scaled_masked_jefferys_divergence ¶

temperature_scaled_masked_jefferys_divergence(teacher_logits: Tensor, student_logits: Tensor, position_mask: Tensor, temperature: float = 1.0) -> <class 'torch.Tensor'>

Compute symmetric Jefferys divergence over masked positions: KL(t||s) + KL(s||t).

Mirrors temperature_scaled_masked_kl_divergence (mask-before-softmax — only active (M, V) rows are softmaxed at large vocab — and temperature scaling) but adds the reverse-KL term so the divergence catches both mass-covering failures (student missing teacher's modes) and mode-seeking failures (student putting mass where teacher has none — e.g. mode-collapse onto a single response template). The forward-KL only penalises the former; the reverse-KL only the latter; Jefferys penalises both.

Gradient routing is the caller's responsibility — this function does not detach either side. If you want the standard distillation semantics (gradients flow only through the student), pass teacher_logits.detach() at the call site.

For an empty mask returns a zero scalar on the teacher's device/dtype to keep composite losses safe to backpropagate.

Parameters:

Name	Type	Description	Default
`teacher_logits` ¶	`Tensor`	Logits from the clean / target model of shape `(B, T, V)`. Caller must detach if standard distillation semantics are desired.	required
`student_logits` ¶	`Tensor`	Logits from the student / noisy model of shape `(B, T, V)`.	required
`position_mask` ¶	`Tensor`	Boolean mask of shape `(B, T)` — `True` at positions to include in the divergence.	required
`temperature` ¶	`float`	Softmax temperature applied to both teacher and student logits. Defaults to `1.0`.	`1.0`

Returns:

Type	Description
`<class 'torch.Tensor'>`	Scalar Jefferys divergence averaged over active positions.

Raises:

Type	Description
`ValueError`	If `temperature` is not strictly positive.

Example

teacher = torch.randn(2, 4, 16) student = torch.randn(2, 4, 16) mask = torch.zeros(2, 4, dtype=torch.bool) mask[:, -2:] = True jd = temperature_scaled_masked_jefferys_divergence(teacher, student, mask) jd.shape torch.Size([])

Added in version v3.43.0. Symmetric (forward + reverse KL) counterpart to `temperature_scaled_masked_kl_divergence` — penalises both mass-covering and mode-seeking failures. Filters by mask BEFORE the softmax and supports temperature scaling.

temperature_scaled_masked_kl_divergence ¶

temperature_scaled_masked_kl_divergence(teacher_logits: Tensor, student_logits: Tensor, position_mask: Tensor, temperature: float = 1.0) -> <class 'torch.Tensor'>

Compute temperature-scaled KL(softmax(teacher / T) || softmax(student / T)) over masked positions.

Gradient routing is the caller's responsibility — this function does not detach either side. If you want the standard distillation semantics (gradients flow only through the student), pass teacher_logits.detach() at the call site. Unlike masked_kl_divergence, the position mask is applied before the softmax — only active (M, V) rows are softmaxed, not the full (B, T, V). For large vocabularies (e.g. V > 100_000) this can shrink the softmax working set by orders of magnitude.

For an empty mask (no active positions) the function returns a zero scalar on the teacher's device/dtype rather than NaN, so a composite loss is safe to backpropagate.

Parameters:

Name	Type	Description	Default
`teacher_logits` ¶	`Tensor`	Logits from the clean / target model of shape `(B, T, V)`. Caller must detach if standard distillation semantics are desired.	required
`student_logits` ¶	`Tensor`	Logits from the student / noisy model of shape `(B, T, V)`.	required
`position_mask` ¶	`Tensor`	Boolean mask of shape `(B, T)` — `True` at positions to include in the KL.	required
`temperature` ¶	`float`	Softmax temperature applied to both teacher and student logits. Defaults to `1.0`.	`1.0`

Returns:

Type	Description
`<class 'torch.Tensor'>`	Scalar KL divergence averaged over active positions.

Raises:

Type	Description
`ValueError`	If `temperature` is not strictly positive.

Example

teacher = torch.randn(2, 4, 16) student = torch.randn(2, 4, 16) mask = torch.zeros(2, 4, dtype=torch.bool) mask[:, -2:] = True kl = temperature_scaled_masked_kl_divergence(teacher, student, mask) kl.shape torch.Size([])

Added in version v3.41.0. Logits-based KL distillation that filters by mask BEFORE the softmax — cuts the working set dramatically at large vocab — and supports temperature scaling. Complements the existing `masked_kl_divergence`.

Changed in version v3.43.0: No longer detaches `teacher_logits` internally — gradient routing is now the caller's responsibility. Pass `teacher_logits.detach()` for the previous distillation semantics. Integer 0/1 masks are now cast to bool before indexing, fixing silently-wrong (often zero) results when a long/int mask was passed instead of a bool mask.

total_variation ¶

total_variation(noisy_logits: Tensor, clean_logits: Tensor, attention_mask: Tensor, support_mask: Tensor | None = None, reduction: Literal['mean', 'none'] = 'mean') -> <class 'torch.Tensor'>

Compute the total variation between two discrete probabilities using logits.

Computes the average tokenwise total variation between embeddings across a sequence.

Note

See https://en.wikipedia.org/wiki/F-divergence#Common_examples_of_f-divergences for the implementation formula.

Parameters:

Name	Type	Description	Default
`noisy_logits` ¶	`Tensor`	The logits produced from SGT data.	required
`clean_logits` ¶	`Tensor`	The logits produced from data without the SGT.	required
`attention_mask` ¶	`Tensor`	The attention mask for the batch.	required
`support_mask` ¶	`Tensor \| None`	The mask placed on both noisy and clean logits. Useful for comparing top-k logits.	`None`
`reduction` ¶	`Literal['mean', 'none']`	Specifies how to reduce across the batch dimension after computing the per-sequence masked mean.	`'mean'`

Returns:

Type	Description
`<class 'torch.Tensor'>`	The total variation between the noisy and clean logits, independently averaged over the batch and sequence lengths.

Changed in version v3.62.0: A sequence with no unmasked positions is now excluded from the batch average instead of contributing a zero that was still counted in the batch-size denominator. Such a sequence carries no signal, so counting it biased the divergence toward zero in proportion to how many fully-padded or fully-masked sequences a batch happened to contain.

Added in version v3.37.0.

divergences

fused_linear_masked_generalized_jsd ¶

student_hidden ¶

teacher_hidden ¶

lm_head_weight ¶

attention_mask ¶

jsd_beta ¶

temperature ¶

reduction ¶

chunk_size ¶

jefferys_divergence ¶

noisy_logits ¶

clean_logits ¶

attention_mask ¶

support_mask ¶

reduction ¶

jensen_shannon_divergence ¶

noisy_logits ¶

clean_logits ¶

attention_mask ¶

support_mask ¶

reduction ¶

masked_cross_entropy ¶

input_logits ¶

target_logits ¶

attention_mask ¶

max_loss ¶

scaling_factor ¶

masked_jefferys_divergence ¶

input_logits ¶

target_logits ¶

attention_mask ¶

log_target ¶

masked_kl_divergence ¶

input_logits ¶

target_logits ¶

attention_mask ¶

log_target ¶

masked_unbiased_dcor ¶

samples_1 ¶

samples_2 ¶

attention_mask ¶

safety_factor ¶

squared_hellinger_distance ¶

noisy_logits ¶

clean_logits ¶

attention_mask ¶

support_mask ¶

reduction ¶

temperature_scaled_masked_jefferys_divergence ¶

teacher_logits ¶

student_logits ¶

position_mask ¶

temperature ¶

temperature_scaled_masked_kl_divergence ¶

teacher_logits ¶

student_logits ¶

position_mask ¶

temperature ¶

total_variation ¶

noisy_logits ¶

clean_logits ¶

attention_mask ¶

support_mask ¶

reduction ¶

`student_hidden` ¶

`teacher_hidden` ¶

`lm_head_weight` ¶

`attention_mask` ¶

`jsd_beta` ¶

`temperature` ¶

`reduction` ¶

`chunk_size` ¶

`noisy_logits` ¶

`clean_logits` ¶

`attention_mask` ¶

`support_mask` ¶

`reduction` ¶

`noisy_logits` ¶

`clean_logits` ¶

`attention_mask` ¶

`support_mask` ¶

`reduction` ¶

`input_logits` ¶

`target_logits` ¶

`attention_mask` ¶

`max_loss` ¶

`scaling_factor` ¶

`input_logits` ¶

`target_logits` ¶

`attention_mask` ¶

`log_target` ¶

`input_logits` ¶

`target_logits` ¶

`attention_mask` ¶

`log_target` ¶

`samples_1` ¶

`samples_2` ¶

`attention_mask` ¶

`safety_factor` ¶

`noisy_logits` ¶

`clean_logits` ¶

`attention_mask` ¶

`support_mask` ¶

`reduction` ¶

`teacher_logits` ¶

`student_logits` ¶

`position_mask` ¶

`temperature` ¶

`teacher_logits` ¶

`student_logits` ¶

`position_mask` ¶

`temperature` ¶

`noisy_logits` ¶

`clean_logits` ¶

`attention_mask` ¶

`support_mask` ¶

`reduction` ¶