Sophia Tang, Yuchen Zhu, Molei Tao, and Pranam Chatterjee

This is the repository for the paper A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding.

Masked discrete diffusion models (MDMs) offer a simple, stable likelihood-based framework for sequence generation, recently extended to any-length settings via token insertion. A2D2 is a unified framework for reward-guided fine-tuning of any-length MDMs that jointly optimizes the insertion and unmasking policies together with a quality-based inference schedule, converging to the intractable reward-tilted distribution without requiring target samples.

🃏 We derive the Radon–Nikodym derivative for the joint insertion–unmasking path measures, enabling theoretically guaranteed convergence to the reward-tilted sequence distribution.

🃏 We establish unmasking and insertion quality as tractable approaches for minimizing decoding error (compounding parallelization error), and train lightweight quality predictors alongside the policy.

🃏 We introduce the Adaptive Joint Decoding (AJD) loss, which provably yields the optimal path measure that generates the reward-tilted distribution while remasking low-quality tokens and dropping low-quality insertions at inference.

🃏 Empirically, A2D2 improves reward optimization while enhancing generation flexibility and accuracy over prior fixed-length fine-tuning and inference-time guidance methods.

We pre-train an any-length MDM on the SAFE dataset (Noutahi et al. 2024, ~950M molecules from ZINC and Unichem in SAFE notation) and fine-tune it with A2D2 to optimize QED (drug-likeness) and synthetic accessibility (SA). A2D2 jointly raises QED and lowers SA over the pre-trained baseline while increasing the fraction of valid, unique, drug-like, and synthesizable molecules. Code and instructions are in /a2d2_mol.

We pre-train an any-length peptide SMILES MDM on ~11M peptides (CycPeptMPDB, SmProt, CycloPs) and fine-tune with A2D2 on five therapeutic properties simultaneously: target-protein binding affinity, solubility, non-hemolysis, non-fouling, and permeability. A2D2 outperforms inference-time multi-objective guidance and fixed-length off-policy RL fine-tuning on almost all objectives, while improving the fraction of valid peptides. Code and instructions are in /a2d2_pep.

We additionally apply A2D2 to reward fine-tuning of any-length language MDMs (LLaDA / FlexMDM), optimizing math-reasoning correctness and format rewards (GSM8K / MATH), including infilling variants. Code is in /a2d2_language.

Each experiment directory contains its own README.md with environment setup, pretrained weight placement, fine-tuning commands, and evaluation instructions.

If you find this repository helpful for your publications, please consider citing our paper:

To use this repository, you agree to abide by the MIT License.