means

Module for mean-related loss functions.

Functions:

Name	Description
means_norm_penalty	Penalize the absolute average difference between the norms of the mean-shifted

means_norm_penalty

means_norm_penalty(
    clean_input_embeddings: Tensor,
    means: Tensor,
    noise_mask: Tensor,
    penalty_value: float,
) -> torch.Tensor

Penalize the absolute average difference between the norms of the mean-shifted input embeddings and a penalty value.

Batch elements whose noise_mask row is entirely False have no positions to penalize, so they are excluded from the batch average rather than contributing a value to it. A batch in which every row is empty reduces to 0.

Parameters:

Name	Type	Description	Default
clean_input_embeddings	Tensor	The clean input embeddings tensor. (B, S, E)	required
means	Tensor	The means tensor to shift the input embeddings by. (B, S, E)	required
noise_mask	Tensor	The mask to apply to the reduction. (B, S)	required
penalty_value	float	The value to penalize the norms against.	required

Returns:

Type	Description
torch.Tensor	The mean absolute difference between the norms of the mean-shifted input embeddings
torch.Tensor	and the penalty value.

Changed in version v3.62.0: Reduces NaN-safely so a batch element with a fully-`False` `noise_mask` row is excluded from the batch average instead of propagating `NaN` through the whole batch. Previously such a row produced a 0/0 masked mean, which made the composite distillation loss non-finite and caused the entire batch to be skipped.