class

BERTForMaskedLM

extendsPretrainedModelMaskedLMMixin

BERTForMaskedLM(config: BERTConfig)

source edit

BERT with a tied masked-language-modeling head.

Implements the masked-LM half of the Devlin et al. (2018) pre-training objective. A two-layer projection (dense + GELU + LayerNorm) maps each hidden state to vocabulary logits via a decoder whose weight matrix is tied to the input word_embeddings table when config.tie_word_embeddings is True. Use for pre-training from scratch, continued pre-training on domain corpora, or fill-in-the-blank inference.

Parameters

configBERTConfig

BERT hyperparameters. config.tie_word_embeddings (default True) controls whether the decoder weight is bound to the input embedding matrix to halve the parameter count.

Attributes

bertBERTModel

Underlying bidirectional encoder trunk.

clsnn.Module

Masked-LM prediction head with its own dense + LayerNorm transform and an output decoder of shape (hidden_size, vocab_size).

Notes

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

When labels is supplied the head computes

\mathcal{L}_{\mathrm{MLM}} = -\frac{1}{|M|} \sum_{i \in M} \log p_{\theta}(x_i \mid x_{\setminus M})

over the set $M$ of masked positions, with positions where the label equals -100 excluded from the sum.

Examples

>>> import lucid
>>> from lucid.models.text.bert import BERTConfig, BERTForMaskedLM
>>> cfg = BERTConfig(num_hidden_layers=2, hidden_size=128,
...                  num_attention_heads=2, intermediate_size=512)
>>> model = BERTForMaskedLM(cfg).eval()
>>> input_ids = lucid.tensor([[101, 7592, 103, 102]])   # [CLS] hello [MASK] [SEP]
>>> out = model(input_ids)
>>> out.logits.shape   # (B=1, T=4, V=30522)
(1, 4, 30522)

Used by 2

Constructors

dunder

init

→None

__init__(config: BERTConfig)

source edit

Initialise the module and validate the supplied config.

Parameters

configModelConfig

Must be an instance of the subclass's declared config_class.

Raises

TypeError

If config_class has not been set on the concrete subclass, or if config is not an instance of config_class.

Instance methods

forward

→MaskedLMOutput

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

source edit

>>> import lucid >>> from lucid.models.text.bert import BERTConfig, BERTForMaskedLM >>> cfg = BERTConfig(num_hidden_layers=2, hidden_size=128, ... num_attention_heads=2, intermediate_size=512) >>> model = BERTForMaskedLM(cfg).eval() >>> input_ids = lucid.tensor([[101, 7592, 103, 102]]) # [CLS] hello [MASK] [SEP] >>> out = model(input_ids) >>> out.logits.shape # (B=1, T=4, V=30522) (1, 4, 30522)