> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.loss.infonce

## Module Contents

### Classes

| Name                                                                               | Description                                              |
| ---------------------------------------------------------------------------------- | -------------------------------------------------------- |
| [`InfoNCEDistillLoss`](#nemo_automodel-components-loss-infonce-InfoNCEDistillLoss) | InfoNCE soft listwise distillation loss module.          |
| [`InfoNCELoss`](#nemo_automodel-components-loss-infonce-InfoNCELoss)               | InfoNCE loss module with optional learnable temperature. |

### Functions

| Name                                                                                   | Description                                            |
| -------------------------------------------------------------------------------------- | ------------------------------------------------------ |
| [`infonce_distill_loss`](#nemo_automodel-components-loss-infonce-infonce_distill_loss) | Soft listwise distillation on InfoNCE candidate sets.  |
| [`infonce_loss`](#nemo_automodel-components-loss-infonce-infonce_loss)                 | InfoNCE contrastive loss with optional hard negatives. |

### API

```python
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.

```python
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
```

```python
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`**

---

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

```python
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
```

```python
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 | None` — default: 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 | None` — default: 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 | None` — default: 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.Tensor` — default: 0.05

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

---

**`direction`** `str` — default: 'q2d'

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

---

**`use_in_batch_negatives`** `bool` — default: True

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

---

**`normalize`** `bool` — default: True

L2-normalize the embeddings before scoring.

---

**`divergence`** `str` — default: '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.

```python
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 | None` — default: 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 | None` — default: 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.Tensor` — default: 0.05

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

---

**`direction`** `str` — default: 'q2d'

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

---

**`use_in_batch_negatives`** `bool` — default: True

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

---

**`normalize`** `bool` — default: 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.