NeuralLossFunction

class lightning_ir.loss.neural.NeuralLossFunction(temperature: float = 1, tol: float = 1e-05, max_iter: int = 50)[source]

Bases: ListwiseLossFunction

Base class for neural loss functions that compute ranks from scores using neural sorting.

__init__(temperature: float = 1, tol: float = 1e-05, max_iter: int = 50) → None[source]

Initialize the NeuralLossFunction.

Parameters:

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

  • tol (float) – Tolerance for convergence. Defaults to 1e-5.

  • max_iter (int) – Maximum number of iterations for convergence. Defaults to 50.

Methods

__init__([temperature, tol, max_iter])

Initialize the NeuralLossFunction.

get_sorted_targets(scores, targets)

Get the sorted targets based on the neural sort permutation matrix.

neural_sort(scores)

Compute the neural sort permutation matrix from scores.

sinkhorn_scaling(mat)

Apply Sinkhorn scaling to the permutation matrix.
abstract compute_loss(output: LightningIROutput, batch: TrainBatch) → torch.Tensor

Compute the loss based on the scores and targets in the output and batch.

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

get_sorted_targets(scores: Tensor, targets: Tensor) → Tensor[source]

Get the sorted targets based on the neural sort permutation matrix.

Parameters:

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

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

Returns:

The sorted targets tensor.

Return type:

torch.Tensor

neural_sort(scores: Tensor) → Tensor[source]

Compute the neural sort permutation matrix from scores.

Parameters:

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

Returns:

The computed permutation matrix.

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

sinkhorn_scaling(mat: Tensor) → Tensor[source]

Apply Sinkhorn scaling to the permutation matrix.

Parameters:

mat (torch.Tensor) – The input permutation matrix.

Returns:

The scaled permutation matrix.

Return type:

torch.Tensor