MarginMSE

class lightning_ir.loss.pairwise.MarginMSE(margin: float | 'scores' = 1.0)

Bases: PairwiseLossFunction

Mean Squared Error loss with a margin for pairwise ranking tasks.

MarginMSE optimizes pairwise ranking by penalizing the squared difference between the predicted score margin of a positive and negative document and a target margin. This target margin can be a fixed constant or dynamically derived from the difference in ground truth or teacher scores, making it particularly effective for knowledge distillation tasks where a student model learns to replicate the score distances of a stronger teacher model.

Originally proposed in: Improving Efficient Neural Ranking Models with Cross-Architecture Knowledge Distillation

__init__(margin: float | 'scores' = 1.0)

Initialize the MarginMSE loss function.

Parameters:

margin (float | Literal["scores"]) – The margin value for the loss.

Methods

__init__([margin])	Initialize the MarginMSE loss function.
compute_loss(output, batch)	Compute the MarginMSE loss.

compute_loss(output: LightningIROutput, batch: TrainBatch) → torch.Tensor

Compute the MarginMSE loss.

Parameters:

• output (LightningIROutput) – The output from the model containing scores.

• batch (TrainBatch) – The training batch containing targets.

Returns:

The computed loss.

Return type:

torch.Tensor

Raises:

ValueError – If the margin type is invalid.