FLOPSRegularization

class lightning_ir.loss.regularization.FLOPSRegularization(query_weight: float = 0.0001, doc_weight: float = 0.0001)

Bases: RegularizationLossFunction

FLOPS Regularization loss function for query and document embeddings. Originally proposed in: Minimizing FLOPS to Learn Efficient Sparse Representations

Methods

compute_loss(output)

Compute the FLOPS regularization loss.

compute_loss(output: BiEncoderOutput) → torch.Tensor

Compute the FLOPS regularization loss.

Parameters:

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

Returns:

The computed loss.

Return type:

torch.Tensor