Research map / Reinforcement learning / RL from feedback

Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design

ICLR 2025 · Reinforcement learning › RL from feedback · approach: RL fine-tuning of generative models · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes DRAKES, which fine-tunes pretrained masked discrete diffusion models to maximize a differentiable reward while staying close to the pretrained model via KL, by backpropagating through sampled trajectories using the Gumbel-Softmax trick, and shows gains on DNA enhancer and protein stability design.

Problem

Fine-tuning pretrained discrete diffusion models on biological sequences to optimize downstream reward objectives (enhancer activity, protein stability) while preserving naturalness, where sampled CTMC trajectories are non-differentiable.

Approach

Formulate reward maximization with KL regularization to a pretrained masked discrete diffusion model as a CTMC control problem, and make the sampling trajectory differentiable via straight-through Gumbel-Softmax so rewards can be backpropagated directly through the trajectory; theory shows the optimum is proportional to exp(r/α) p_pre.

Key results

• On regulatory DNA design, DRAKES reaches 92.5% ATAC-Acc and 0.887 3-mer correlation versus 45.3% and 0.848 for TDS, while keeping log-likelihood near pretrained levels.
• On protein inverse folding, DRAKES achieves a 78.6% success rate versus 63.6% for SMC and 34.4% for the pretrained model.
• Without the KL term (DRAKES w/o KL), protein scRMSD rises to 7.307, showing over-optimization that the KL penalty mitigates.

Measured effects reported in the paper

Surprising findings

• Classifier-free guidance performs very poorly on protein design (1.3% success rate) because the labeled Megascale dataset covers only a few hundred backbones.
• DRAKES without KL achieves higher predicted activity on DNA but lower independent ATAC-Acc, indicating reward over-optimization.

Limitations

• Requires differentiable reward models, which may not hold for all reward types.
• Evaluation is in silico only; no wet-lab validation.
• Requires a balance between truncation length and memory cost in backpropagation through the trajectory.

Experimental setup

Datasets: MPRA enhancer dataset (~700k sequences) PDB training set Megascale protein stability dataset JASPAR

Benchmarks: enhancer activity oracle (HepG2) ATAC-seq chromatin accessibility classifier protein stability (Pred-ddG) oracle with ESMFold scRMSD

Metrics: Pred-Activity ATAC-Acc 3-mer Corr JASPAR Corr App-Log-Lik Pred-ddG scRMSD success rate

Baselines: CG SMC TDS CFG DRAKES w/o KL

Builds on: masked discrete diffusion classifier guidance for discrete diffusion Gumbel-Softmax DPS

Closest papers by meaning

• Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-based Decoding (NeurIPS 2025 · similarity 0.91)
• Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction (ICLR 2025 · similarity 0.91)
• Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models (NeurIPS 2024 · similarity 0.91)
• Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design (ICML 2025 · similarity 0.90)
• Steering Generative Models with Experimental Data for Protein Fitness Optimization (NeurIPS 2025 · similarity 0.89)
• Feedback Efficient Online Fine-Tuning of Diffusion Models (ICML 2024 · similarity 0.89)