A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding
Sophia Tang, Yuchen Zhu, Molei Tao, Pranam Chatterjee
Abstract
Discrete diffusion models offer a simple and stable likelihood-based framework for sequence generation, recently extended to any-length settings via token insertion. Principled reward-guided fine-tuning for any-length discrete diffusion, however, remains largely unexplored. We introduce Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding (A2D2), a unified framework for reward-guided fine-tuning of any-length discrete diffusion models via joint optimization of the insertion and unmasking policies together with a quality-based inference schedule. We derive the Radon-Nikodym derivative for the joint insertion-unmasking path measures, enabling theoretically guaranteed convergence to the intractable reward-tilted sequence distribution without requiring target samples. Building on this, we establish unmasking and insertion quality as tractable approaches for minimizing decoding error and introduce the Adaptive Joint Decoding (AJD) loss, which provably yields the optimal path measure that generates the reward-tilted distribution. Empirically, A2D2 improves reward optimization while enhancing generation flexibility and accuracy over prior fixed-length fine-tuning and inference-time guidance methods.
AI Impact Assessments
(1 models)Scientific Impact Assessment: A2D2 — Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding
1. Core Contribution
A2D2 addresses a genuine gap in the discrete diffusion literature: principled reward-guided fine-tuning for any-length masked discrete diffusion models (MDMs). While prior work has explored RL fine-tuning for fixed-length MDMs, any-length models introduce a combinatorially larger action space spanning both token unmasking and variable-length insertions. The paper's key theoretical contributions are:
The problem formulation is well-motivated: any-length generation is essential for molecules (variable SMILES lengths), peptides, and code infilling, where the output length is unknown a priori.
2. Methodological Rigor
The theoretical development is thorough. The paper derives the RND for joint CTMCs (Proposition 4.1), proves unique minimizers for both quality losses (Propositions 3.1, 3.3), establishes connections between quality maximization and compounding parallelization error (Proposition 3.2), and shows the AJD loss converges to the optimal path measure (Proposition C.4). The proofs, provided extensively in the appendix, follow a coherent chain from observations about non-overlapping rates through to the final loss derivation.
However, several concerns arise:
3. Potential Impact
Drug discovery and peptide design: The molecule and peptide experiments demonstrate meaningful improvements in multi-objective optimization. For peptides targeting TfR, A2D2 improves validity from ~10% to ~49% while substantially improving binding affinity, solubility, and permeability. These are practically relevant improvements for therapeutic design pipelines.
Language reasoning: The GSM8K results (+25 points Pass@1 at 128 steps) and HumanEval infilling improvements are notable, particularly because they demonstrate that reward-aligned any-length diffusion can compete with fixed-length approaches on standard NLP benchmarks. The finding that A2D2 concentrates accuracy in low-step regimes has practical implications for inference efficiency.
Broader methodological impact: The RND derivation for joint CTMCs and the quality-adaptive inference framework could generalize to other variable-length generation settings beyond MDMs, potentially influencing work on edit-based models, insertion transformers, and structured prediction with variable outputs.
4. Timeliness & Relevance
This paper arrives at a critical juncture. Discrete diffusion models are rapidly gaining traction for language (LLaDA, MDLM, DiffuCoder) and biological sequences. The extension to any-length generation (Kim et al., 2025a) is recent, and this is among the first papers to address RL fine-tuning in this setting. The code infilling and math reasoning experiments connect to the active area of inference-time scaling for diffusion LLMs.
The any-length setting is particularly timely for biological applications where sequence length is a design variable (e.g., peptide length optimization), making this more than an incremental extension of fixed-length fine-tuning.
5. Strengths & Limitations
Strengths:
Limitations:
6. Additional Observations
The connection between quality maximization and CPE minimization (Section 3, Proposition C.2) provides useful theoretical insight, but the paper does not empirically verify that CPE actually decreases during training. The ablation studies (Appendix G) are informative but limited to molecules; analogous ablations for language would strengthen the claims about the quality predictors' contribution.
Generated Jun 12, 2026
Comparison History (13)
Paper 2 directly tackles a fundamental problem in scientific discovery: extracting governing equations from noisy, high-dimensional data. Its approach bridges machine learning with physics, biology, and neuroscience, offering broad, cross-disciplinary applications. While Paper 1 presents strong algorithmic improvements for sequence generation, Paper 2's potential to uncover fundamental laws of nature from observational data provides a more profound potential impact across diverse scientific fields.
Paper 1 addresses a critical and highly timely challenge in AI safety and LLM alignment: the tension between helpfulness and harmlessness in RLHF. By applying mechanistic interpretability to uncover how these objectives interfere at the neuron level, it provides fundamental insights that can directly impact how future foundational models are aligned. While Paper 2 offers robust theoretical advancements in discrete diffusion models, Paper 1's focus on understanding and solving core bottlenecks in widely deployed LLM alignment gives it a broader potential impact across both AI research and real-world deployment.
Paper 2 (A2D2) likely has higher scientific impact because it contributes a broadly applicable, theoretically grounded framework for reward-guided fine-tuning and decoding in any-length discrete diffusion models, including a Radon–Nikodym derivation and convergence guarantees. This can influence multiple areas of sequence generation (NLP, code, biological sequences) and provides reusable principles/losses (AJD) beyond a single benchmark. Paper 1 is impressive and timely for automated theorem proving, but appears more system/engineering- and benchmark-driven with narrower cross-field methodological generality.
Paper 1 addresses reward-guided fine-tuning for sequence generation using discrete diffusion, a highly active area with broad applications in NLP and computational biology. Its theoretical contributions and empirical improvements offer immediate, widespread utility. While Paper 2 introduces an elegant, physics-informed neural layer, its impact is likely confined to the more specialized intersection of geometric deep learning and higher gauge theory, giving Paper 1 a broader and more immediate scientific footprint.
Paper 2 presents a fundamental algorithmic advance in generative AI with strong theoretical guarantees for any-length discrete diffusion. Its unified framework for reward-guided fine-tuning has broad applicability across multiple domains requiring sequence generation, such as NLP and computational biology. While Paper 1 provides a valuable dataset and benchmark for a specific subfield of chemistry, Paper 2's methodological innovation offers a wider breadth of impact and a foundational contribution to machine learning.
Paper 2 likely has higher scientific impact due to its direct, large-scale real-world application to climate modeling and carbon-cycle uncertainty, a highly timely and societally critical area. AI4Land’s outputs (global high-resolution reconstructions, open-source emulators, HPC-enabled pipeline, planned coupling to digital twins) can be broadly used by Earth system science, remote sensing, ecology, and policy communities, increasing breadth of impact. Paper 1 is methodologically innovative for discrete diffusion fine-tuning, but its impact is more specialized within ML sequence generation and may face faster turnover in the field.
FlexRank addresses a broadly impactful problem—adaptive deployment of large models across varying compute budgets—with a practical 'train-once, deploy-everywhere' paradigm applicable to LLMs and ViTs. Its breadth of applicability across model types and deployment scenarios gives it wider real-world impact. While A2D2 makes rigorous theoretical contributions to discrete diffusion fine-tuning (Radon-Nikodym derivatives, convergence guarantees), it targets a narrower subfield. FlexRank's timeliness regarding efficient LLM deployment and its potential to change deployment practices across industry and academia give it higher estimated impact.
Paper 2 (A2D2) has higher potential impact due to greater novelty and breadth: it proposes a unified, theoretically grounded framework (Radon–Nikodym path-measure derivation, provable convergence, optimality of AJD) for reward-guided fine-tuning in any-length discrete diffusion—an emerging paradigm relevant across sequence generation tasks. Its contributions are more general and could influence diffusion-based decoding, RL fine-tuning, and probabilistic modeling. Paper 1 is timely and practically valuable (quantized LRM decoding + kernels) but is more incremental and scoped to NVFP4 inference and reasoning LLM deployment.
Paper 2 addresses a critical bottleneck in proteomics by improving de novo peptide sequencing. Its training-free, plug-and-play approach yields massive performance gains (up to 39.1%), offering immediate, high-impact applications in biological research and drug discovery. While Paper 1 provides rigorous theoretical advancements in machine learning, Paper 2 demonstrates more immediate real-world scientific utility.
Paper 2 likely has higher scientific impact due to its strong real-world application potential and cross-field reach: zero-shot, compositional generation of whole-cortex fMRI dynamics conditioned on language and spatial priors could directly affect neuroscience, cognitive science, and experimental design (counterfactual/in-silico experiments). It is timely given rapid adoption of diffusion/flow methods and foundation-model-style conditioning. Paper 1 is methodologically rigorous and novel within discrete diffusion fine-tuning, but its impact is more concentrated within ML sequence generation, with less immediate domain-transformative application.