L2Regularization

class lightning_ir.loss.regularization.L2Regularization(query_weight: float = 0.0001, doc_weight: float = 0.0001)[source]

Bases: RegularizationLossFunction

L2 Regularization loss function for query and document embeddings. Originally proposed in: Ridge Regression: Biased Estimation for Nonorthogonal Problems

__init__(query_weight: float = 0.0001, doc_weight: float = 0.0001) → None

Initialize the RegularizationLossFunction.

Parameters:

  • query_weight (float) – Weight for the query embeddings regularization. Defaults to 1e-4.

  • doc_weight (float) – Weight for the document embeddings regularization. Defaults to 1e-4.

Methods

compute_loss(output)

Compute the L2 regularization loss.
compute_loss(output: BiEncoderOutput) → torch.Tensor[source]

Compute the L2 regularization loss.

Parameters:

output (BiEncoderOutput) – The output from the model containing query and document embeddings.

Returns:

The computed loss.

Return type:

torch.Tensor

process_embeddings(output: BiEncoderOutput) → Tuple[torch.Tensor, torch.Tensor]

Process the embeddings from the output.

Parameters:

output (BiEncoderOutput) – The output from the model containing query and document embeddings.

Returns:

The processed query and document embeddings.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Raises:

  • ValueError – If query_embeddings are not present in the output.

  • ValueError – If doc_embeddings are not present in the output.

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