data
BERTForPreTrainingOutput
extends
ModelOutputBERTForPreTrainingOutput(prediction_logits: Tensor, seq_relationship_logits: Tensor, loss: Tensor | None = None, mlm_loss: Tensor | None = None, nsp_loss: Tensor | None = None)Combined output for BERTForPreTraining.
Aggregates the masked-LM logits, next-sentence-prediction logits, and optional per-objective and combined losses produced by the full BERT pre-training pipeline of Devlin et al., 2018.
Parameters
prediction_logitsTensorMLM head logits of shape
(B, T, vocab_size) — one distribution
over the WordPiece vocabulary per input position.seq_relationship_logitsTensorNSP head logits of shape
(B, 2) — binary IsNext / NotNext scores
derived from the pooled [CLS] embedding.Sum of
mlm_loss and nsp_loss when both are available;
otherwise the single available loss, or None if neither label
set was supplied.Cross-entropy on masked positions when
labels was supplied.Binary cross-entropy on the NSP head when
next_sentence_label
was supplied.Notes
Reference: Devlin, Chang, Lee, and Toutanova, "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding", NAACL 2019 (arXiv:1810.04805) §3.1.
Examples
>>> import lucid
>>> from lucid.models.text.bert import BERTConfig, BERTForPreTraining
>>> cfg = BERTConfig(num_hidden_layers=2, hidden_size=128,
... num_attention_heads=2, intermediate_size=512)
>>> model = BERTForPreTraining(cfg).eval()
>>> input_ids = lucid.tensor([[101, 7592, 102, 2088, 102]])
>>> out = model(input_ids)
>>> out.prediction_logits.shape, out.seq_relationship_logits.shape
((1, 5, 30522), (1, 2))