Work done as final project for CS 242 (Computing at Scale) at Harvard, taught by H.T. Kung.

Large language model (LLM) specialization typically involves routing, mixture of experts, or post-training approaches, which can learn redundant parameters across domains. Diffusion language models (DLMs) provide specialization at sampling time via iterative denoising and classifier-free guidance (CFG). In this project, we study whether a sampler-conditioned DLM, equipped with domain embeddings and guidance, can compose multiple capabilities at sampling time and compete with autoregressive (AR) baselines under a fixed parameter budget. At the ~120M scale, we find that AR models outperform DLMs. However, within the diffusion regime, with two-domain specialization, the sampler-conditioned DLM consistently improves over an unguided DLM. Sampling-time domain conditioning is empirically useful inside the diffusion family but does not yet replace AR specialization, at least at 120M scale. However, DLMs are known to be more parameter efficient than ARs at larger-scale.

Monolithic LLMs can underperform on specialized domains. Cloud providers have addressed the issue by training math, code, and chat variants. However, these strategies (joint multi-domain models, MoE, and domain-specific fine-tuning) all duplicate core circuits such as syntax, world knowledge, and basic reasoning, which fragments capacity. DLMs instead move specialization from parameters to the sampler.

We pre-train all models on FineWeb-Edu-10B and inject domain specialization using GSM8K (math) and CBT (reading comprehension) datasets. To ablate architecture and conditioning, we implement six 120M parameter models with matched compute per sample and identical data pipelines: Joint ARM (one AR model trained on both domains, without explicit domain token), MoE ARM (same architecture as joint ARM except MLPs are replaced by MoE blocks and router), separate ARMs (two separate ~60M parameter ARMs, each fine-tuned on one domain), joint DLM (one DLM trained jointly on both domains, without CFG), separate DLMs (two separate ~60M parameter DLMs, each fine-tuned on one domain, without CFG), and sampler-conditioned DLM (one DLM trained on both domains, using domain embeddings for sampling-time conditioning). For ARMs, conditioning strength increases from implicit multi-domain, to MoE, to separate domain-specific fine-tuned model. For DLMs, we compare joint multi-domain training, per-domain specialists, and sampler conditioned learners.

CFG was introduced to steer the trajectory of generative image diffusion models without training an auxiliary classifier. Inspired by that mechanism, we inject domain embeddings at sampling time to bias every denoising step toward joint-domain trajectory while still sharing all weights. We design a sampler-conditioned DLM that uses domain embeddings for sampling-time conditioning, without duplicating learned circuits. Our sampler-conditioned model follows the discrete Block Diffusion process where tokens are grouped into blocks of size $B=16$. For each block $b$, we sample a corruption rate $t_b \in [t_{\ell}, t_{u}]$ and then draw i.i.d per-token Bernoulli masks with probability $t_b$ where every position in the block shares a noise rate but is masked independently. The corrupted prefix is: $$ \mathbf{x}_t = \mathbf{m}_t \cdot \mathbf{x}_{\text{mask}} + (1-\mathbf{m}_t) \cdot \mathbf{x}_0 $$ where $\mathbf{m}_t \in [0,1]^T$ stores these Bernoulli draws. We concatenate $[x_t; x0]$ and apply the Block Diffusion attention mask (block diagonal, offset block causal, block causal), thus a noisy position can read its own block, already-denoised blocks in the same document, and the clean suffix. Call the transformer outputs logits on the masked tokens $\ell_{\theta}(\mathbf{x}_t, \mathbf{p}, \mathbf{e}_c)$. We optimize the block-weighted cross entropy objective: $$ L_{\text{block}} = \mathbb{E} \left[ \frac{\mathbf{m}_t}{t_b + 10^{-3}} \cdot \text{CE}\left(\ell_\theta (\mathbf{x}_t, t, \mathbf{p}, \mathbf{e}_c), \mathbf{x}_0 \right) \right] $$ At sampling, we choose the desired domain token and run the denoiser once per block prefix. Guided and unguided generations differ only in whether $\mathbf{e}_c$ is set to the learned embedding or to zero. Therefore, specialization is kept within the guidance embedding while maintaining a shared backbone.

DLMs are most useful for controlled text generation, where they can synthesize coherent, stylistically consistent outputs across diverse domains. They also support sequence editing and infilling by identifying the tokens and spans that require revision, which can iteratively improve overall output quality. Additionally, by removing the serial dependency in output token generation, DLMs in general can offer opportunities for faster inference compared to traditional AR approaches. Compared to sampler-conditioned DLMs, by offloading specialization from parameters to the sampler, it may be possible to achieve effective multi-domain learning without duplication in model parameters, which could lead to more compact model architectures that are competitive with the current large autoregressors.

Large Language Diffusion Models [Paper] [Code]