"""Bidirectional Transformer for masked-diffusion language modeling. The model is intentionally simple and self-contained: a token + learned positional embedding, a stack of pre-norm bidirectional Transformer blocks, and a tied LM head. There is **no causal mask** — the masked-diffusion objective lets every position attend to every other position, which is the source of the "implicit data augmentation over token orderings" that makes diffusion shine in data-constrained regimes (Prabhudesai et al., 2025, arXiv:2507.15857). It subclasses ``PreTrainedModel`` so ``save_pretrained`` / ``from_pretrained`` and Hub uploads work out of the box. The forward pass returns logits over the real vocabulary **plus** the absorbing ``[MASK]`` column (index ``vocab_size``); the mask column is ignored when computing loss and when scoring. """ from __future__ import annotations import math import torch import torch.nn as nn from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput from .config import MaskedDiffusionConfig class TransformerBlock(nn.Module): """Pre-norm bidirectional Transformer block (MHSA + GEGLU-free MLP).""" def __init__(self, config: MaskedDiffusionConfig) -> None: super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_eps) self.attn = nn.MultiheadAttention( embed_dim=config.n_embd, num_heads=config.n_head, dropout=config.dropout, batch_first=True, ) self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_eps) hidden = config.ffn_mult * config.n_embd self.mlp = nn.Sequential( nn.Linear(config.n_embd, hidden), nn.GELU(), nn.Linear(hidden, config.n_embd), nn.Dropout(config.dropout), ) def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor | None) -> torch.Tensor: h = self.ln_1(x) # need_weights=False keeps the attention fast; mask is True where padded. attn_out, _ = self.attn( h, h, h, key_padding_mask=key_padding_mask, need_weights=False ) x = x + attn_out x = x + self.mlp(self.ln_2(x)) return x class MaskedDiffusionLM(PreTrainedModel): """A bidirectional Transformer trained with an absorbing-state diffusion loss.""" config_class = MaskedDiffusionConfig base_model_prefix = "mdlm" supports_gradient_checkpointing = False # Tells HF that lm_head.weight is tied to the input embedding (so # save_pretrained does not treat it as an illegal shared tensor). Recent # transformers (>=4.53) require the {target: source} dict form; older ones # (e.g. the eval pipeline's 4.51.3) iterate it as keys, so the dict is # backward-compatible. A bare list crashes get_expanded_tied_weights_keys(). _tied_weights_keys = {"lm_head.weight": "tok_emb.weight"} def __init__(self, config: MaskedDiffusionConfig) -> None: super().__init__(config) self.config = config # +1 embedding row for the absorbing [MASK] state at index vocab_size. self.tok_emb = nn.Embedding(config.num_embeddings, config.n_embd) self.pos_emb = nn.Embedding(config.n_positions, config.n_embd) self.drop = nn.Dropout(config.dropout) self.blocks = nn.ModuleList( [TransformerBlock(config) for _ in range(config.n_layer)] ) self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_eps) self.lm_head = nn.Linear(config.n_embd, config.num_embeddings, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.tok_emb.weight self.post_init() # ── HF plumbing ──────────────────────────────────────────────────────── def get_input_embeddings(self) -> nn.Module: return self.tok_emb def set_input_embeddings(self, value: nn.Module) -> None: self.tok_emb = value def get_output_embeddings(self) -> nn.Module: return self.lm_head def set_output_embeddings(self, value: nn.Module) -> None: self.lm_head = value def _init_weights(self, module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) # ── Forward ────────────────────────────────────────────────────────────── def forward( self, input_ids: torch.LongTensor, attention_mask: torch.Tensor | None = None, labels: torch.LongTensor | None = None, layer_duplication_factor: int | None = None, token_type_ids: torch.Tensor | None = None, **kwargs, ) -> MaskedLMOutput: """Run the bidirectional encoder. Args: input_ids: (B, T) token ids, may contain ``mask_token_id``. attention_mask: (B, T) 1 for real tokens, 0 for padding. labels: (B, T) original tokens at masked positions, ``-100`` elsewhere. If given, a plain (unweighted) cross-entropy is returned in ``loss``. The weighted MDLM loss lives in ``masking.diffusion_loss`` and is used by the training loop instead. layer_duplication_factor: optional inference-time "reasoning depth" override (repeats the middle blocks). Defaults to the config value. token_type_ids / **kwargs: accepted and ignored. This model is a single-segment bidirectional encoder, but HF tokenizers emit ``token_type_ids`` by default and some eval harnesses (e.g. the official ``reading`` task) call ``model(**tokenizer(...))``, so we must tolerate these extra arguments instead of crashing. Returns: ``MaskedLMOutput`` with ``logits`` of shape (B, T, vocab_size + 1). """ x = self._encode(input_ids, attention_mask, layer_duplication_factor) logits = self.lm_head(x) loss = None if labels is not None: loss = nn.functional.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100 ) return MaskedLMOutput(loss=loss, logits=logits) def _encode( self, input_ids: torch.LongTensor, attention_mask: torch.Tensor | None, layer_duplication_factor: int | None, ) -> torch.Tensor: """Run the encoder stack and return final hidden states (B, T, n_embd).""" B, T = input_ids.shape device = input_ids.device pos = torch.arange(T, device=device).unsqueeze(0).expand(B, T) x = self.drop(self.tok_emb(input_ids) + self.pos_emb(pos)) # MultiheadAttention expects True where a position should be *ignored*. key_padding_mask = None if attention_mask is not None: key_padding_mask = attention_mask == 0 dup = layer_duplication_factor or self.config.layer_duplication_factor for block in self._expanded_blocks(dup): x = block(x, key_padding_mask) return self.ln_f(x) def _expanded_blocks(self, dup: int): """Return the block sequence, optionally repeating the middle blocks. With ``dup == 1`` this is just ``self.blocks``. With ``dup > 1`` the interior blocks (all but the first and last) are applied ``dup`` times, giving extra "reasoning depth" at no parameter cost — the duplicated-layer idea from the proposal's COMPS / entity-tracking hypotheses. """ if dup <= 1 or self.config.n_layer <= 2: return self.blocks first, *middle, last = list(self.blocks) return [first, *(middle * dup), last] class MaskedDiffusionModel(MaskedDiffusionLM): """Headless variant: same weights, returns hidden states instead of logits. Registered under ``AutoModel`` in ``auto_map``. The official BabyLM GLUE fine-tuning harness loads encoders with ``AutoModel.from_pretrained`` and feeds ``last_hidden_state`` (B, T, n_embd) into its own classification head, so this class must NOT return logits — ``MaskedLMOutput.logits`` would be (B, T, vocab+1) and the harness would mistake it for the encodings. The parameter names are identical to :class:`MaskedDiffusionLM`, so any checkpoint loads into either class unchanged. """ def forward( # type: ignore[override] self, input_ids: torch.LongTensor, attention_mask: torch.Tensor | None = None, layer_duplication_factor: int | None = None, token_type_ids: torch.Tensor | None = None, **kwargs, ) -> BaseModelOutput: x = self._encode(input_ids, attention_mask, layer_duplication_factor) return BaseModelOutput(last_hidden_state=x)