Enabling Autoregressive Models to Fill In Masked Tokens

TL;DR

Autoregressive (AR) models are highly scalable and benefit from KV caching but struggle with infilling between contexts, whereas masked language models (MLMs) handle infilling but are less scalable. MARIA (Masked and Autoregressive Infilling Architecture) unifies a frozen pretrained MLM and a frozen AR by training a lightweight linear decoder on their concatenated hidden states to enable masked infilling while preserving AR-style inference. The method initializes the fusion layer with W3 = [W1/2; W2/2], equivalent to a product-of-experts, and optimizes an autoregressive distribution conditioned on unmasked context via the MARIA objective L_MARIA. Empirically, MARIA achieves state-of-the-art masked infilling compared with discrete diffusion baselines, delivering strong perplexity, sample quality (as judged by LLM-based assessments), and high throughput due to KV caching; it also enables flexible test-time strategies like simulated annealing for unconditional generation.

Abstract

Historically, LLMs have been trained using either autoregressive (AR) or masked language modeling (MLM) objectives, with AR models gaining dominance in recent years. However, AR models are inherently incapable of masked infilling, which is the ability to predict masked tokens between past and future context. In contrast, MLM models suffer from intrinsic computational inefficiencies during both training and inference that hinder their scalability. This work introduces MARIA (Masked and Autoregressive Infilling Architecture), a novel approach that leverages the strengths of both paradigms to achieve state-of-the-art masked infilling performance. MARIA combines a pre-trained MLM and AR model by training a linear decoder that takes their concatenated hidden states as input. This minimal modification enables the AR model to perform infilling while retaining its inherent advantages in terms of faster inference with KV caching. Our results demonstrate that MARIA significantly outperforms existing methods, namely discrete diffusion models, on masked infilling tasks.

Figures (5)

• Figure 1: MARIA architecture and training pipeline. MARIA takes two frozen pretained models: one MLM and one AR. As input, the MLM recieves the masked inputs and the AR model recieves the denoised inputs. We compute the hidden states under each model and perform truncating and shifting operations to ensure both hidden states model the same tokens. MARIA trains a linear layer to predict the logits of each masked input on the concatenated hidden states. This training scheme models an autoregressive distribution conditioned on unmasked tokens.
• Figure 2: Comparing evaluation loss curves for two different weight initializations. Product initialization (ours) is a far better weight initialization than random weight initialization, leading to faster training and better convergence.
• Figure 3: Throughput over input length. We show the throughput in tokens per second for sequences of given lengths at 0.5 masking rate. MARIA 1B exhibits the best performance, and MARIA-7B is comparable to DiffuLlama 7B. Decoding ModernBERT autoregressively is extremely inefficient at scale, and therefore is impractical in many circumstances.
• Figure 4: ELO scores for masked infilling. We perform infilling on downstream data with words masked 50 percent. Using GPT4o-mini as a judge we compute the ELO scores for each model respectively. MARIA 7B and 1B have the highest rating ELO rating under the Bradley-Terry model.
• Figure 5: Scaling test time compute for unconditional generation. We compare our simulated annealing inference approach for MARIA to our baseline discrete diffusion methods. MARIA 1B using simulated annealing effectively trades-off quality (as measured by generative perplexity) and with compute (measured in wall clock time).