nemo_automodel.components.loss.infonce

Module Contents

Classes

Name	Description
`InfoNCEDistillLoss`	InfoNCE soft listwise distillation loss module.
`InfoNCELoss`	InfoNCE loss module with optional learnable temperature.

Functions

Name	Description
`infonce_distill_loss`	Soft listwise distillation on InfoNCE candidate sets.
`infonce_loss`	InfoNCE contrastive loss with optional hard negatives.

API

class nemo_automodel.components.loss.infonce.InfoNCEDistillLoss(
    temperature: float = 0.05,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True,
    divergence: str = 'kl',
    cross_device_negatives: bool = True
)

Bases: Module

InfoNCE soft listwise distillation loss module.

nemo_automodel.components.loss.infonce.InfoNCEDistillLoss.forward(
    student_queries: torch.Tensor,
    student_documents: torch.Tensor,
    teacher_queries: torch.Tensor,
    teacher_documents: torch.Tensor,
    student_hard_negatives: torch.Tensor | None = None,
    teacher_hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None
) -> torch.Tensor

class nemo_automodel.components.loss.infonce.InfoNCELoss(
    temperature: float = 0.05,
    learnable_temperature: bool = False,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True,
    cross_device_negatives: bool = True
)

Bases: Module

InfoNCE loss module with optional learnable temperature.

log_inv_tau

nemo_automodel.components.loss.infonce.InfoNCELoss.current_temperature() -> torch.Tensor

nemo_automodel.components.loss.infonce.InfoNCELoss.forward(
    queries: torch.Tensor,
    documents: torch.Tensor,
    hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None
) -> torch.Tensor

nemo_automodel.components.loss.infonce.infonce_distill_loss(
    student_queries: torch.Tensor,
    student_documents: torch.Tensor,
    teacher_queries: torch.Tensor,
    teacher_documents: torch.Tensor,
    student_hard_negatives: torch.Tensor | None = None,
    teacher_hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None,
    temperature: float | torch.Tensor = 0.05,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True,
    divergence: str = 'kl'
) -> torch.Tensor

Soft listwise distillation on InfoNCE candidate sets.

Builds the same candidate sets as infonce_loss for the student and the (detached) teacher, then matches the student’s distribution over those candidates to the teacher’s.

Parameters:

student_queries

torch.Tensor

Tensor of shape [batch, student_hidden].

student_documents

torch.Tensor

Tensor of shape [batch, student_hidden], row-aligned with student_queries.

teacher_queries

torch.Tensor

Tensor of shape [batch, teacher_hidden]. Detached.

teacher_documents

torch.Tensor

Tensor of shape [batch, teacher_hidden], row-aligned with teacher_queries. Detached.

student_hard_negatives

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, negatives, student_hidden] — mined negative documents, per query, so they are candidates for "q2d" only.

teacher_hard_negatives

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, negatives, teacher_hidden], aligned with student_hard_negatives. Required when it is given. Detached.

hard_negatives_mask

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, negatives] — 1 for a real negative, 0 for padding. Padded columns are scored as -inf and masked out of the divergence.

temperature

float | torch.TensorDefaults to 0.05

Logit scale divisor; a 0-dim Tensor when it is learned.

direction

strDefaults to 'q2d'

"q2d", "d2q", or "symmetric".

use_in_batch_negatives

boolDefaults to True

Score every query against every document in the batch instead of only its own positive.

normalize

boolDefaults to True

L2-normalize the embeddings before scoring.

divergence

strDefaults to 'kl'

"kl", "ce", or "mse".

Returns: torch.Tensor

Scalar Tensor: the mean divergence over the batch.

Raises:

ValueError: If the embedding shapes disagree with the layouts above, if direction or divergence is unknown, or if the requested directions have no candidates to score against — use_in_batch_negatives=False needs hard negatives and direction="q2d", because the document-to-query direction can only draw candidates from the batch.

nemo_automodel.components.loss.infonce.infonce_loss(
    queries: torch.Tensor,
    documents: torch.Tensor,
    hard_negatives: torch.Tensor | None = None,
    hard_negatives_mask: torch.Tensor | None = None,
    temperature: float | torch.Tensor = 0.05,
    direction: str = 'q2d',
    use_in_batch_negatives: bool = True,
    normalize: bool = True
) -> torch.Tensor

InfoNCE contrastive loss with optional hard negatives.

Parameters:

queries

torch.Tensor

Tensor of shape [batch, hidden] — pooled query embeddings.

documents

torch.Tensor

Tensor of shape [batch, hidden] — pooled embeddings of the positive document for each query, row-aligned with queries.

hard_negatives

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, negatives, hidden] — mined negative documents for each query. Rows are per-query, so these are candidates for the "q2d" direction only.

hard_negatives_mask

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, negatives] — 1 for a real negative, 0 for padding in a ragged batch. Padded columns are scored as -inf and contribute nothing.

temperature

float | torch.TensorDefaults to 0.05

Logit scale divisor; a 0-dim Tensor when it is learned.

direction

strDefaults to 'q2d'

"q2d", "d2q", or "symmetric".

use_in_batch_negatives

boolDefaults to True

Score every query against every document in the batch instead of only its own positive.

normalize

boolDefaults to True

L2-normalize the embeddings before scoring.

Returns: torch.Tensor

Scalar Tensor: the mean cross-entropy over the batch.

Raises:

ValueError: If the embedding shapes disagree with the layouts above, if direction is unknown, or if the requested directions have no candidates to score against — use_in_batch_negatives=False needs hard_negatives and direction="q2d", because the document-to-query direction can only draw candidates from the batch.