nemo_automodel.components.loss.embedding_distill

Module Contents

Classes

Name	Description
`EmbeddingDistillLoss`	Cosine embedding-distillation loss module.
`EmbeddingMSELoss`	MSE embedding-distillation loss module.
`ScoreDistillLoss`	Score-matching distillation loss module.

Functions

Name	Description
`_pairwise_cosine_softmax`	-
`_score_loss_side`	-
`cosine_distance`	Per-sample `1 - cosine_similarity`.
`distill_loss_pair`	Embedding-alignment cosine loss over query and doc sides.
`mse_loss_pair`	Per-element MSE alignment between projected student and teacher.
`score_distill_loss`	Score-matching distillation (row-softmax pairwise cosine matrix MSE).

API

class nemo_automodel.components.loss.embedding_distill.EmbeddingDistillLoss(
    reduction: str = 'mean'
)

Bases: Module

Cosine embedding-distillation loss module.

nemo_automodel.components.loss.embedding_distill.EmbeddingDistillLoss.forward(
    s_q_proj: torch.Tensor,
    t_q: torch.Tensor,
    s_d_proj: torch.Tensor,
    t_d: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.loss.embedding_distill.EmbeddingMSELoss(
    normalize: bool = False,
    reduction: str = 'mean'
)

Bases: Module

MSE embedding-distillation loss module.

nemo_automodel.components.loss.embedding_distill.EmbeddingMSELoss.forward(
    s_q_proj: torch.Tensor,
    t_q: torch.Tensor,
    s_d_proj: torch.Tensor,
    t_d: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.loss.embedding_distill.ScoreDistillLoss(
    temperature: float = 0.02
)

Bases: Module

Score-matching distillation loss module.

nemo_automodel.components.loss.embedding_distill.ScoreDistillLoss.forward(
    s_q: torch.Tensor,
    t_q: torch.Tensor,
    s_d: torch.Tensor,
    t_d: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.loss.embedding_distill._pairwise_cosine_softmax(
    z: torch.Tensor,
    temperature: float
) -> torch.Tensor

nemo_automodel.components.loss.embedding_distill._score_loss_side(
    s: torch.Tensor,
    t: torch.Tensor,
    temperature: float
) -> torch.Tensor

nemo_automodel.components.loss.embedding_distill.cosine_distance(
    z_s_proj: torch.Tensor,
    z_t: torch.Tensor
) -> torch.Tensor

Per-sample 1 - cosine_similarity.

nemo_automodel.components.loss.embedding_distill.distill_loss_pair(
    s_q_proj: torch.Tensor,
    t_q: torch.Tensor,
    s_d_proj: torch.Tensor,
    t_d: torch.Tensor,
    reduction: str = 'mean'
) -> torch.Tensor

Embedding-alignment cosine loss over query and doc sides.

nemo_automodel.components.loss.embedding_distill.mse_loss_pair(
    s_q_proj: torch.Tensor,
    t_q: torch.Tensor,
    s_d_proj: torch.Tensor,
    t_d: torch.Tensor,
    normalize: bool = False,
    reduction: str = 'mean'
) -> torch.Tensor

Per-element MSE alignment between projected student and teacher.

nemo_automodel.components.loss.embedding_distill.score_distill_loss(
    s_q: torch.Tensor,
    t_q: torch.Tensor,
    s_d: torch.Tensor,
    t_d: torch.Tensor,
    temperature: float = 0.02
) -> torch.Tensor

Score-matching distillation (row-softmax pairwise cosine matrix MSE).