ApproxRankMSE

class lightning_ir.loss.approximate.ApproxRankMSE(temperature: float = 1, discount: Literal['log2', 'reciprocal'] | None = None)

Bases: ApproxLossFunction

Approximate Rank Mean Squared Error (RankMSE) loss function for ranking tasks. Originally proposed in: Rank-DistiLLM: Closing the Effectiveness Gap Between Cross-Encoders and LLMs for Passage Re-ranking

__init__(temperature: float = 1, discount: Literal['log2', 'reciprocal'] | None = None)

Initialize the ApproxRankMSE loss function.

Parameters:

• temperature (float) – Temperature parameter for scaling the scores. Defaults to 1.

• discount (Literal["log2", "reciprocal"] | None) – Discounting strategy for the loss. If None, no discounting is applied. Defaults to None.

Methods

__init__([temperature, discount])	Initialize the ApproxRankMSE loss function.
compute_loss(output, batch)	Compute the ApproxRankMSE loss.

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

Compute the ApproxRankMSE 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

static get_approx_ranks(scores: Tensor, temperature: float) → Tensor

Compute approximate ranks from scores.

Parameters:

• scores (torch.Tensor) – The input scores.

• temperature (float) – Temperature parameter for scaling the scores.

Returns:

The computed approximate ranks.

Return type:

torch.Tensor

process_scores(output: LightningIROutput) → torch.Tensor

Process the scores from the output.

Parameters:

output (LightningIROutput) – The output from the model.

Returns:

The scores tensor.

Return type:

torch.Tensor

process_targets(scores: torch.Tensor, batch: TrainBatch) → torch.Tensor

Process the targets from the batch.

Parameters:

• scores (torch.Tensor) – The scores tensor.

• batch (TrainBatch) – The training batch.

Returns:

The processed targets tensor.

Return type:

torch.Tensor