Per-token cross-entropy loss helper for masked-LM / token-class heads.
Encoder text families (BERT, RoFormer, …) all reduce a (B, T, V)
logit tensor against (B, T) labels with the same recipe: flatten
the sequence axis into the batch axis, then run
lucid.nn.functional.cross_entropy with ignore_index=-100.
Centralising the call here removes a 3-line stanza repeated across
every ForMaskedLM / ForTokenClassification class.
Notes
The mixin is stateless and exposes a single @staticmethod —
inheriting from it is purely a documentation / discoverability
convenience; the same effect is obtained by calling
MaskedLMMixin.compute_lm_loss(...) directly.
Examples
>>> class BERTForMaskedLM(PretrainedModel, MaskedLMMixin):
... def forward(self, input_ids, labels=None):
... logits = self.head(self.bert(input_ids))
... loss = self.compute_lm_loss(logits, labels) if labels else None
... return MaskedLMOutput(logits=logits, loss=loss)Used by 4
Static methods
1compute_lm_loss(logits: Tensor, labels: Tensor, ignore_index: int = -100)Compute per-token cross-entropy loss for an LM head.
Parameters
logitsTensor(B, T, V) per-token logits.labelsTensor(B, T) int target token ids. Entries equal to
ignore_index contribute zero loss and zero gradient.ignore_index(int, optional, keyword - only)= -100-100 marks non-masked positions in MLM and
padding tokens in token classification.Returns
TensorScalar loss tensor (mean over the contributing positions).
Notes
Implementation: flatten logits to (B*T, V) and labels
to (B*T,), cast labels to long, then call
lucid.nn.functional.cross_entropy. No probabilistic
reweighting; the loss is a plain mean over non-ignored positions.
Examples
>>> logits = lucid.randn(2, 10, 32000)
>>> labels = lucid.randint(0, 32000, (2, 10))
>>> loss = MaskedLMMixin.compute_lm_loss(logits, labels)
>>> loss.shape
()