"""MARIA model: a frozen AR backbone + frozen MLM backbone + trained fusion head. Paper: "Enabling Autoregressive Models to Fill In Masked Tokens" (Israel et al. 2025) https://arxiv.org/abs/2502.06901 """ from __future__ import annotations import os from dataclasses import dataclass from typing import Optional, Tuple, Union import torch import torch.nn.functional as F from torch import nn from torch.nn import CrossEntropyLoss from transformers import ( AutoModelForCausalLM, AutoModelForMaskedLM, PreTrainedModel, ) from transformers.modeling_outputs import ModelOutput from .configuration_maria import AR_TOKEN_PROCESSORS, MariaConfig @dataclass class MariaLMOutput(ModelOutput): loss: Optional[torch.FloatTensor] = None logits: Optional[torch.FloatTensor] = None ar_hidden_states: Optional[torch.FloatTensor] = None mlm_hidden_states: Optional[torch.FloatTensor] = None class LinearHead(nn.Module): """Trained fusion head: [ar_hidden ; mlm_hidden] -> vocab logits.""" def __init__(self, ar_hidden_size: int, mlm_hidden_size: int, vocab_size: int): super().__init__() self.linear = nn.Linear(ar_hidden_size + mlm_hidden_size, vocab_size) def forward(self, ar_h: torch.Tensor, mlm_h: torch.Tensor) -> torch.Tensor: return self.linear(torch.cat([ar_h, mlm_h], dim=-1)) def _apply_ar_token_processor(input_ids: torch.Tensor, name: str) -> torch.Tensor: """Clamp tokens that fall outside the AR backbone's vocabulary.""" spec = AR_TOKEN_PROCESSORS[name] if spec["clamp_above"] < 0: return input_ids return torch.where(input_ids > spec["clamp_above"], spec["replacement"], input_ids) def _project_modernbert_hidden(mlm_model, hidden: torch.Tensor) -> torch.Tensor: """Apply ModernBERT's final-norm + head projection to encoder hidden states. These are normally applied inside ``ModernBertForMaskedLM.forward`` before the decoder; we replicate them here so the fusion head sees the same representation the original MLM decoder would have seen. """ return mlm_model.head(mlm_model.model.final_norm(hidden)) class MariaForMaskedLM(PreTrainedModel): """MARIA: an autoregressive model with masked-token infilling. The forward pass is autoregressive: the prediction for position ``i`` is conditioned on AR features for positions ``[0..i-1]`` (causal context) and MLM features for the entire sequence (bidirectional context). """ config_class = MariaConfig base_model_prefix = "maria" main_input_name = "input_ids" supports_gradient_checkpointing = False # Class-level hint so HF.from_pretrained knows these keys are tied and not "missing". _tied_weights_keys = ["ar_model.lm_head.weight", "mlm_model.decoder.weight"] _keys_to_ignore_on_load_missing = ["ar_model.lm_head.weight", "mlm_model.decoder.weight"] def __init__(self, config: MariaConfig): super().__init__(config) ar_dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype ar_cfg = config.get_ar_pretrained_config() mlm_cfg = config.get_mlm_pretrained_config() # Patch attn implementation onto the sub-configs so from_config respects it. ar_cfg._attn_implementation = config.ar_attn_implementation mlm_cfg._attn_implementation = config.mlm_attn_implementation self.ar_model = AutoModelForCausalLM.from_config(ar_cfg, torch_dtype=ar_dtype) self.mlm_model = AutoModelForMaskedLM.from_config(mlm_cfg, torch_dtype=ar_dtype) self.lm_head = LinearHead( ar_hidden_size=config.ar_hidden_size, mlm_hidden_size=config.mlm_hidden_size, vocab_size=config.vocab_size, ) # The backbones are frozen by construction; only ``lm_head`` ever trains. for p in self.ar_model.parameters(): p.requires_grad = False for p in self.mlm_model.parameters(): p.requires_grad = False # Re-tie embeddings inside each backbone so that bundled checkpoints can # ship only one copy of any tied weight (saves ~200MB and silences HF's # "weights not initialized" warning). self.ar_model.tie_weights() self.mlm_model.tie_weights() def tie_weights(self): # Called by PreTrainedModel.from_pretrained AFTER state_dict load. Forward # the call into each backbone so tied projections pick up the loaded embeds. self.ar_model.tie_weights() self.mlm_model.tie_weights() # ---------------------------------------------------------------------------- # Forward / training loss # ---------------------------------------------------------------------------- def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, denoised_ids: Optional[torch.LongTensor] = None, ) -> MariaLMOutput: """Compute fused logits and (optionally) cross-entropy loss. ``input_ids`` may contain mask tokens. The AR backbone is fed the denoised sequence (``labels`` filled in where present) so that during training it has access to the ground-truth left context. """ if labels is not None and (labels == -100).all(): return MariaLMOutput(loss=torch.tensor(0.0, requires_grad=True, device=self.device)) if denoised_ids is None and labels is not None: denoised_ids = torch.where(labels == -100, input_ids, labels) elif denoised_ids is None: denoised_ids = input_ids mask_id = getattr(self, "_mask_token_id", None) if mask_id is not None and torch.any(denoised_ids == mask_id): # The AR side must never see mask tokens; either supply ``labels`` that # cover every masked position or pre-fill ``denoised_ids`` yourself. return MariaLMOutput(loss=torch.tensor(0.0, requires_grad=True, device=self.device)) ar_ids = _apply_ar_token_processor(denoised_ids, self.config.ar_token_processor) if attention_mask is not None: attention_mask = attention_mask.bool() with torch.no_grad(): ar_outputs = self.ar_model(ar_ids, attention_mask=attention_mask, output_hidden_states=True) mlm_outputs = self.mlm_model(input_ids, attention_mask=attention_mask, output_hidden_states=True) # AR: features at position i predict token i+1 (causal shift). ar_h = ar_outputs.hidden_states[-1][:, :-1, :] # MLM: drop the first position (BOS); align with AR-predicted positions 1..N. mlm_h_raw = mlm_outputs.hidden_states[-1] if mlm_h_raw.dim() == 2: mlm_h_raw = mlm_h_raw.unsqueeze(0) mlm_h = _project_modernbert_hidden(self.mlm_model, mlm_h_raw)[:, 1:, :] logits = self.lm_head(ar_h, mlm_h).contiguous() loss = None if labels is not None: shifted_labels = labels[:, 1:].contiguous() loss = CrossEntropyLoss()( logits.view(-1, logits.size(-1)), shifted_labels.view(-1), ) return MariaLMOutput(loss=loss, logits=logits, ar_hidden_states=ar_h, mlm_hidden_states=mlm_h) # ---------------------------------------------------------------------------- # Likelihood # ---------------------------------------------------------------------------- @torch.no_grad() def compute_nll( self, input_ids: torch.LongTensor, labels: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, reduction: str = "mean", ) -> torch.Tensor: """Negative log-likelihood over positions where ``labels != -100``. ``reduction='mean'`` returns the per-token mean NLL (suitable for perplexity via ``exp``); ``reduction='sum'`` returns the total NLL. """ out = self.forward(input_ids=input_ids, attention_mask=attention_mask, labels=labels) shifted_labels = labels[:, 1:].contiguous().view(-1) nll = F.cross_entropy( out.logits.view(-1, out.logits.size(-1)), shifted_labels, reduction="sum", ignore_index=-100, ) if reduction == "sum": return nll n = (shifted_labels != -100).sum().clamp_min(1) return nll / n # ---------------------------------------------------------------------------- # Infilling # ---------------------------------------------------------------------------- @torch.no_grad() def infill( self, input_ids: torch.LongTensor, mask_token_id: int, attention_mask: Optional[torch.Tensor] = None, greedy: bool = True, temperature: float = 1.0, steps: int = 1, resample_rate: float = 0.3, eps: float = 1e-3, return_history: bool = False, ) -> Union[torch.LongTensor, list]: """Fill in every position where ``input_ids == mask_token_id``. With ``steps=1`` (the default) MARIA does a single left-to-right pass: one MLM forward over the whole sequence, then a KV-cached AR forward that emits one token per masked position. With ``steps>1`` MARIA re-masks a random ``resample_rate`` fraction of the previously-filled positions and re-runs the single-shot pass, using a linearly annealed temperature schedule. This iterative refinement is a small improvement on top of the single-shot pass, not the main contribution of the paper. """ assert input_ids.size(0) == 1, "infill() supports batch size 1" if steps == 1: return self._infill_single_pass( input_ids, mask_token_id, attention_mask, greedy, temperature ) # Iterative refinement. x = self._infill_single_pass( input_ids, mask_token_id, attention_mask, greedy=False, temperature=1.0, to_cpu=False ) history = [x.detach().cpu()] maskable = torch.ones_like(input_ids).bool() maskable[:, 0] = False maskable[:, -1] = False maskable[input_ids != mask_token_id] = False for t in torch.linspace(1.0, eps, steps, device=x.device): remask = (torch.rand_like(x, dtype=torch.float) < resample_rate) & maskable if not remask.any(): break x[remask] = mask_token_id x = self._infill_single_pass( x, mask_token_id, attention_mask, greedy=True, temperature=float(t), to_cpu=False ) if return_history: history.append(x.detach().cpu()) return history if return_history else x.detach().cpu() def _infill_single_pass( self, input_ids: torch.LongTensor, mask_token_id: int, attention_mask: Optional[torch.Tensor], greedy: bool, temperature: float, to_cpu: bool = True, ) -> torch.LongTensor: device = self.device x = input_ids.clone().to(device) if attention_mask is None: attention_mask = torch.ones_like(x) attention_mask = attention_mask.to(device) masked_idx = torch.unique((x == mask_token_id).nonzero()[:, 1]) if masked_idx.numel() == 0: return x.cpu() if to_cpu else x # Single bidirectional MLM pass over the masked sequence. mlm_out = self.mlm_model(x, attention_mask=attention_mask, output_hidden_states=True) mlm_h = mlm_out.hidden_states[-1] if mlm_h.dim() == 2: mlm_h = mlm_h.unsqueeze(0) mlm_h = _project_modernbert_hidden(self.mlm_model, mlm_h) # Left-to-right KV-cached AR pass that fills one mask at a time. ar_x = _apply_ar_token_processor(x, self.config.ar_token_processor) past_key_values = None prev = 0 for curr in masked_idx.tolist(): chunk = ar_x[:, prev:curr] cache_position = torch.arange(prev, curr, device=device) out = self.ar_model( chunk, attention_mask=attention_mask[:, prev:curr], past_key_values=past_key_values, position_ids=cache_position.unsqueeze(0), cache_position=cache_position, output_hidden_states=True, use_cache=True, ) past_key_values = out.past_key_values ar_h_last = out.hidden_states[-1][:, -1, :] logits = self.lm_head(ar_h_last, mlm_h[:, curr, :]) probs = torch.softmax(logits[0] / temperature, dim=-1) tok = torch.argmax(probs, dim=-1) if greedy else torch.multinomial(probs, 1).squeeze(-1) x[0, curr] = tok ar_x[0, curr] = _apply_ar_token_processor(tok.view(1, 1), self.config.ar_token_processor)[0, 0] prev = curr return x.cpu() if to_cpu else x # ---------------------------------------------------------------------------- # Embeddings # ---------------------------------------------------------------------------- @torch.no_grad() def get_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor: """Return concatenated [AR ; MLM] hidden states over positions 1..N-1.""" mlm_out = self.mlm_model(input_ids, output_hidden_states=True) mlm_h = mlm_out.hidden_states[-1] if mlm_h.dim() == 2: mlm_h = mlm_h.unsqueeze(0) mlm_h = _project_modernbert_hidden(self.mlm_model, mlm_h)[:, 1:, :] ar_ids = _apply_ar_token_processor(input_ids, self.config.ar_token_processor) ar_out = self.ar_model(ar_ids, output_hidden_states=True) ar_h = ar_out.hidden_states[-1][:, :-1, :] return torch.cat([ar_h, mlm_h], dim=-1)