class

BERTForPreTraining

extendsPretrainedModelMaskedLMMixin

BERTForPreTraining(config: BERTConfig)

source edit

BERT with the original joint MLM + NSP pre-training objective.

Combines the masked-language-modeling head (decoder weight tied to input embeddings when config.tie_word_embeddings is True) with the next-sentence-prediction head on top of the pooled [CLS] embedding. This is the exact head configuration used in Devlin et al., 2018 to train BERT-Base and BERT-Large from scratch.

Supply labels (MLM targets) and/or next_sentence_label (binary NSP target) to compute the corresponding losses; their sum is exposed as output.loss. Use this class only when reproducing the original pre-training recipe — newer encoder-only LMs typically drop NSP and use BERTForMaskedLM directly.

Parameters

configBERTConfig

BERT hyperparameters. config.tie_word_embeddings (default True) controls whether the MLM decoder weight is tied to the input embedding matrix.

Attributes

bertBERTModel

Underlying bidirectional encoder trunk.

clsnn.Module

Combined head holding both the MLM prediction projection (cls.predictions) and the NSP binary linear (cls.seq_relationship).

Notes

Reference: Devlin, Chang, Lee, and Toutanova, "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding", NAACL 2019 (arXiv:1810.04805) §3.1.

The combined loss when both objectives are supplied is

\mathcal{L}_{\mathrm{pretrain}} = \mathcal{L}_{\mathrm{MLM}} + \mathcal{L}_{\mathrm{NSP}}.

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, 103, 102, 2088, 102]])
>>> out = model(input_ids)
>>> out.prediction_logits.shape    # MLM logits  (B=1, T=6, V=30522)
(1, 6, 30522)
>>> out.seq_relationship_logits.shape   # NSP logits  (B=1, 2)
(1, 2)

Used by 1

lucid.models.text.bert

Constructors

dunder

init

→None

__init__(config: BERTConfig)

source edit

Instance methods

forward

→BERTForPreTrainingOutput

forward(input_ids: Tensor, attention_mask: Tensor | None = None, token_type_ids: Tensor | None = None, labels: Tensor | None = None, next_sentence_label: Tensor | None = None)

source edit

>>> 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, 103, 102, 2088, 102]]) >>> out = model(input_ids) >>> out.prediction_logits.shape # MLM logits (B=1, T=6, V=30522) (1, 6, 30522) >>> out.seq_relationship_logits.shape # NSP logits (B=1, 2) (1, 2)