PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Optimal Large Diffusion Language Models

Jinjie Ni, Qian Liu, Chao Du, Longxu Dou, Hang Yan, Zili Wang, Tianyu Pang, Michael Qizhe Shieh

arXiv:2510.03280v2cs.LGcs.AIcs.CL

TL;DR

Large DLM training lacks systematic guidance on how compute, data repetition, model size, and optimization choices should be balanced. Quokka fits scaling laws across compute- and data-constrained regimes and studies major design choices, finding that DLMs need roughly 2–5× more data than AR models at equal FLOPs while some AR optimization rules transfer.

  • Problem

    DLM training knowledge remains near blank, with prior work largely heuristic or extrapolated from AR models across compute- and data-constrained regimes.

  • Method

    Quokka empirically models DLM loss across parameters, unique data, epochs, and compute using large-scale training runs, while ablating modeling and optimization designs.

  • Results

    DLMs require roughly 2–5× more data than AR models at the same FLOPs, and AR scaling conclusions for learning rate and batch size transfer directly.

  • Takeaways & Limitations

    DLM training should favor larger corpora relative to model size under compute constraints while retaining established AR choices for learning rate and batch size.

  • Takeaways & Limitations

    The validation-loss formulation lacks strict theoretical justification, may not hold universally, and fitted contours can overshoot epochs and underestimate model size.

Abstract

from arXiv · show

We introduce Quokka, the first systematic scaling law for diffusion language models (DLMs), encompassing both compute-constrained and data-constrained regimes, and studying the key modeling and optimization designs. Quokka is a good friend of Chinchilla and provides wider scopes. We hope the results would bring short-term practical guidance in DLMs training and long-term inspirations for the whole AI community.

1. Introduction

Quokka addresses the lack of systematic training guidance for large DLMs by studying compute- and data-constrained scaling alongside modeling and optimization choices. It finds that DLMs require substantially more data than AR models at equal compute, while several training-design conclusions transfer from AR models.

  • DLM motivation: DLMs’ bidirectional, any-order modeling supports tasks with noncausal dependencies, including coding, mathematics, and report generation.The cited motivation connects these modeling properties to on-the-fly context modification during generation.
  • Motivation: The work targets a gap in DLM training knowledge, where prior studies were largely heuristic or extrapolated conclusions from AR models.The two central regimes are fixed-compute allocation and fixed-data allocation, with additional modeling and optimization choices affecting end-of-training performance.
  • Scope and contribution: Quokka introduces a large-scale DLM scaling law covering compute-constrained and data-constrained regimes plus key modeling and optimization designs.The study empirically examines language-modeling loss and downstream evaluations across these factors.
  • Compute-constrained scaling laws for DLMs: DLMs require roughly 2–5× more data than AR models under the same FLOPs budget, while model size and data scale proportionally with compute.The allocation functions are estimated from models ranging from under 7M to over 11B parameters and datasets from 1B to over 260B tokens.
  • Data-constrained scaling laws for DLMs: The data-constrained analysis studies how unique data, model parameters, and repeated training interact, including when overfitting begins and how to allocate parameters and epochs.It treats fixed unique data with effectively unlimited compute and asks both how long a fixed model can train and how to jointly choose model size and repetitions.
  • Modeling and optimization designs: The study ablates transition kernels, diffusion schedules, curricula, loss formulation, learning rate, batch size, and weight decay; AR learning-rate and batch-size scaling conclusions transfer directly.The transfer result coexists with markedly different DLM scaling coefficients from AR models.

2. Preliminaries

This section reviews Chinchilla-style compute allocation and data-constrained effective-size formulations before introducing masked diffusion and its corruption–denoising processes. It emphasizes that effective-size discounting captures diminishing returns but fails to model overfitting.

  • Chinchilla Scaling Law: Chinchilla estimates compute-optimal AR training by varying parameters and tokens at fixed values, fixed FLOPs, or through a fitted loss surface.Its compute budget is approximately C≈6ND, and near-lockstep scaling follows when the fitted exponents are similar.
  • A data-constrained generalization: The data-constrained AR formulation discounts repeated tokens and excess parameters through effective data and model sizes, capturing diminishing marginal value.Unique tokens and compute-optimal parameters define the effective counterparts, with exponential half-life forms for repetitions and excess capacity.
  • Limitations of the AR formulation: The effective-size formulation assumes validation loss is non-increasing, but repeated training in practice increases loss through overfitting.The mismatch is described as a consequence of the bias–variance tradeoff.
  • Masked Diffusion Language Models: Masked diffusion uses an absorbing transition kernel that preserves discreteness, supports any-order modeling, and enables exact position-wise factorization and flexible likelihood estimation.The formulation also natively supports multi-token prediction and is presented as a strong competitor to AR modeling.
  • Forward corruption process: The forward process independently masks tokens according to α_t, the probability that a token remains clean at noise level t.The schedule is monotone, with α_0=1 and α_1=0, and the expected unmasked fraction equals α_t.
  • Reverse denoising process: The reverse process starts fully masked and progressively reveals tokens, while already revealed tokens remain fixed.Masked tokens either stay masked or are sampled from a data-prediction distribution according to the schedule.
  • Masked diffusion objective: Masked diffusion’s clean-token conditional depends only on the visible context, allowing denoising without an explicit time embedding.Its variational objective uses a schedule-dependent importance weight, which becomes 1/t under the linear schedule α_t=1−t.

3. Compute-Constrained Scaling Law for Diffusion Language Models

Quokka estimates compute-optimal DLM scaling through IsoFLOPs experiments and a parametric loss fit. Both approaches find that parameters and training data scale nearly proportionally with compute, while DLMs require substantially more data than AR models at fixed FLOPs.

  • 3. Compute-Constrained Scaling Law for Diffusion Language Models: 2–5× more data than AR models is optimal for DLMs at the same FLOPs budget, favoring smaller models and larger corpora.Both empirical approaches support this compute-constrained allocation pattern.
  • 3.1. Approach 1: IsoFLOPs Profiles: The IsoFLOPs approach varies model size across nine fixed budgets from 3 × 10^18 to 1 × 10^21 FLOPs, fits each loss minimum, and extrapolates optimal allocations.The loss curves are fit with parabolas to estimate the parameter count at each minimum.
  • 3.1. Approach 1: IsoFLOPs Profiles: Doubling model size requires doubling training data under compute-optimal scaling, with fitted exponents a=0.51 and b=0.49.The fitted formulas are N≈0.0216C^0.514 and D≈7.7C^0.486.
  • 3.3. Optimal Model Scaling: At LLaDA’s 1.1 × 10^23 FLOPs, the compute-optimal allocation is a 15B-parameter model with 1.2T tokens rather than its 8B model with 2.3T tokens.Quokka’s comparison illustrates the practical effect of its larger data multiplier and smaller optimal model allocation.
  • 3.2. Approach 2: Fitting a Parametric Loss Function: The parametric loss fit decomposes loss into irreducible error, finite model-capacity error, and finite-data error, then derives the compute-efficient frontier.The fitted loss contours and isoFLOPs cross-sections are compared with real data points.
  • 3.2. Approach 2: Fitting a Parametric Loss Function: The parametric approach independently yields a≈0.50 and b≈0.50, reinforcing that parameters and data scale at the same pace under fixed compute.The fit uses Huber loss with L-BFGS optimization and grid-searched initializations.

4. Data-Constrained Scaling Law for Diffusion Language Models

The data-constrained scaling law models validation loss across model size, unique data, and training epochs, capturing overfitting and guiding allocation under limited data. Experiments show U-shaped loss, data-dependent overfitting, and practical optima, while fitting limitations constrain universal interpretation.

  • Empirical validation-loss landscape: Validation loss eventually rises with repeated training, smaller models overfit more slowly, more unique data delays overfitting, and minimum loss improves before worsening with excess capacity.These patterns were observed across varying parameter scales, unique data sizes, and epochs.
  • Validation-loss formulation: The proposed formulation extends Chinchilla with an effective dataset size D′ depending on epochs, model size, and unique data, reproducing the full learning–overfitting cycle.At small e, effective data grows and loss falls; at large e, the overfitting penalty shrinks effective data and loss rises.
  • Fitted scaling behavior: Overfitting onset scales roughly as e_opt ∝ U_D^0.39/N^0.55, so larger unique-data budgets support more training while larger models overfit sooner.The fitted contours also show local optima and generally favor larger models and more epochs as unique-token budgets increase.
  • Caveats: The formulation lacks strict theoretical justification, assumes a single descent despite some double-descent runs, and can overestimate epochs or underestimate model size.Predicted low validation losses may also be unattainable because of fitting error, while absolute losses depend on the validation set and tokenizer.

5. Key Modeling and Optimization Choices

The ablations examine transition kernels, schedules, curricula, losses, and optimization hyperparameters in diffusion language models. Masked diffusion and mostly linear scheduling perform strongly, while several established autoregressive optimization practices transfer to DLMs.

  • Transition kernels: Masked diffusion consistently outperforms uniform diffusion across pretraining losses and downstream metrics by a wide margin.Uniform models achieve low loss on noisy-versus-clean identification, suggesting transformation into correct embeddings is the harder challenge.
  • Schedules and curricula: Linear schedules are strongest and most stable across most settings, while cosine performs worst and poly2 performs better on HellaSwag.A Gaussian curriculum reduces loss faster early and slightly improves end-of-training performance on both benchmarks.
  • Loss formulations: MaskGIT loss converges faster initially, whereas the principled diffusion loss is compared as the likelihood-bound objective for masked generative models.The supplied results passage introduces this comparison but does not report its final performance outcome.
  • Optimization hyperparameters: DLM learning-rate behavior transfers from autoregressive models, with minimal end-of-training differences and a slight advantage for 1e-4.The results support reusing established autoregressive learning-rate practices.
  • Optimization hyperparameters: Weight decay is unnecessary for single-epoch convergence but improves multi-epoch DLM results, while DLMs remain relatively robust to repeated data without it.Removing weight decay severely harms autoregressive models in the multi-epoch setting and consistently increases parameter norms across ablations.

6. Related Work

Prior scaling-law work established compute–data trade-offs and later addressed data scarcity, but systematic scaling laws for diffusion language models remained limited. Quokka extends this line by studying DLM scaling and related modeling choices more comprehensively.

  • Scaling-law foundations: Chinchilla showed that under fixed compute, optimal performance requires scaling model size and training data together rather than simply enlarging models.This challenged earlier power-law interpretations that emphasized ever-larger models.
  • Data-constrained scaling: Data-constrained scaling laws address limited unique data by modeling repeated exposure through effective model and dataset sizes, but assume non-increasing validation loss.That assumption conflicts with observed overfitting during repeated training.
  • Broader research direction: Scaling-law research has broadened beyond pretraining loss toward downstream performance, inference dynamics, theoretical grounding, architecture, data quality, and task-specific requirements.This broader agenda motivates more refined laws that integrate multiple sources of variation.
  • Diffusion language models: Before Quokka, DLM scaling evidence was limited to low-FLOPs comparisons and provided few scaling-law coefficients or broader insights.Other work demonstrated competitive large-scale DLMs, commercial systems, and hybrid approaches, but did not supply the same systematic scaling-law scope.

7. Discussions

Scaling laws are useful high-level guidance, but practical allocation also reflects deployability, adoption, hardware, and training-time constraints. Their coefficients can shift across architectures and datasets, so precise choices require empirical tuning.

  • Commercial allocation can depart from scaling-optimal choices because deployability, customer adoption, hardware compatibility, and flexible training time also matter.
  • Scaling laws should be viewed as high-level guidance because their coefficients can shift across architectures and datasets.Precise choices require empirical tuning under specific constraints.

A. Implementation Details

The experiments use a modified Megatron-LM setup with corpus-specific training data and c4-val evaluation. Compute-constrained runs apply heavy smoothing and specified learning-rate, schedule, sequence-length, and batch-size settings.

  • Experiments used a heavily modified Megatron-LM codebase, with Nemotron-CC for compute-constrained runs, c4-en for data-constrained runs, and c4-val for validation.Token budgets were randomly sampled without additional filtering.
  • Compute-constrained runs used Gaussian smoothing with a 301-step window, reducing variance by approximately 13× at the cost of a mild 40-step lag.Learning rates were 2e−4 below 8B parameters and 1.25e−4 above 8B, with cosine decay and Warmup-Stable reuse.
  • All models used sequence length 2048, while batch size increased from 256 to 1024 across model-size ranges for stability.
  • Alternative fitting figures compare Equation (15)'s predicted contours and validation losses with real data points, while noting epoch overestimation and model-size underestimation.

B. Alternative Data-Constrained Formulas and Fitting Results

Alternative data-constrained formulas decompose learning and overfitting effects, but their fits remain substantially worse than Equation (15). The appendix compares predicted contours, allocations, and validation losses against observed data.

  • B. Alternative Data-Constrained Formulas and Fitting Results: Alternative formulas achieve losses of 31.52 and 23.8 over 23,145 data points, versus 9.78 for Equation (15).
  • B.1. Additive Overfitting Term v1: Equation (18) decomposes the data-constrained law into learning loss and an additive overfitting penalty.
  • B.2. Additive Overfitting Term v2: Equation (22) uses a more complicated additive overfitting penalty after additional trials, with contours and observed-data comparisons shown in Figures 21–23.
  • B.2. Additive Overfitting Term v2: Table 5 reports FLOPs and token allocations from approaches 2 and 3, noting that the fitting approach underestimates N_opt for very large models.
  • B.2. Additive Overfitting Term v2: Figures 18–20 show Equation (18)'s predicted contours, optimal allocations, and comparisons with actual data points and validation losses.
  • B.2. Additive Overfitting Term v2: Table 6 lists the architecture details for all models trained in the work.
Loading 2510.03280v2…