nemo_automodel.components.loss.listmle

ListMLE ranking loss over a teacher’s top-k tokens.

Distillation objectives such as :func:~nemo_automodel.components.loss.soft_ce .masked_soft_cross_entropy match the teacher’s probability mass. A drafter in speculative decoding is instead accepted or rejected on whether its ranking of the next few candidates agrees with the target’s, so ViSpec (arXiv:2509.15235) adds a Plackett-Luce ranking term next to the distribution term. It is shared here rather than kept beside one algorithm because both the EAGLE-1/2 and the ViSpec objectives use it.

Module Contents

Functions

Name	Description
`listmle_loss`	ListMLE ranking loss over the target’s top-k tokens.

API

nemo_automodel.components.loss.listmle.listmle_loss(
    logits: torch.Tensor,
    target_probs: torch.Tensor,
    topk: int
) -> torch.Tensor

ListMLE ranking loss over the target’s top-k tokens.

Scores the student on reproducing the target’s ordering of its topk most likely tokens: the Plackett-Luce likelihood of drawing those tokens, under the student’s logits, in the target’s own descending-probability order.

Parameters:

logits

torch.Tensor

Tensor of shape [tokens, vocab] — the student’s logits at the supervised positions.

target_probs

torch.Tensor

Tensor of shape [tokens, vocab] — the target’s probabilities at the same positions.

topk

int

Number of top target tokens to rank.

Returns: torch.Tensor

Scalar Tensor: the mean over tokens of the summed negative