PaperlayerView exact revision on arXiv

Source-linked AI summary

GRAS: Guided Reduced-Variance Proposals and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion

Kwanyoung Kim

arXiv:2608.26585v1cs.LGcs.CEq-bio.QMstat.ML

TL;DR

Training-free reward alignment for discrete diffusion must combine guidance and search without retraining, but existing pipelines use noisy proposal gradients and fixed-temperature resampling. GRAS adds zero-cost variance reduction and adaptive-temperature selection, attaining the best training-free reward across regulatory DNA and protein design while exposing a diversity cost.

  • Problem

    Training-free steering combines gradient guidance and search, but guided proposals can have high-variance gradient estimates and selectors can ignore per-step reward spread.

  • Method

    GRAS combines Rao–Blackwellized reveals for differentiable rewards, leave-one-out baselines for non-differentiable rewards, and AMC adaptive-temperature selection.

  • Results

    GRAS attains the best training-free reward on regulatory DNA and protein design, matches or surpasses a reward-fine-tuned model, and remains effective for non-differentiable rewards.

  • Takeaways & Limitations

    The guided proposal captures most reward cheaply, while adaptive selection supplies additional reward and transfers across guided samplers.

  • Takeaways & Limitations

    Reward gains reduce sample diversity through selection, and AMC guarantees target a reweighted rather than exact tilted posterior asymptotically in particle count K.

Abstract

from arXiv · show

Discrete diffusion models have become a strong, widely adopted class of generators for sequence data, and steering them toward a downstream reward at inference time, without any retraining, is increasingly important. Such training-free steering is done by gradient guidance, by search, or by combining the two. We study the combined regime and identify two weaknesses in how it is usually run: the guided proposal estimates its gradient from a single noisy sample, and the search then resamples particles at a fixed temperature that ignores how rewards spread across each denoising step. We address both with a small set of changes that add no denoiser cost. For the proposal, we lower the estimator variance with a Rao-Blackwellized reveal for differentiable rewards and a leave-one-out baseline for non-differentiable ones; for the search, we standardize the per-step values into a group-relative advantage and prove it collapses to a single active ingredient, an adaptive resampling temperature. We call the resulting method Guided Reduced-variance proposals and Adaptive Selection (GRAS). GRAS is simple yet effective: across regulatory DNA and protein design it attains the best training-free reward, outperforming prior training-free methods and matching or surpassing a reward-fine-tuned model, and it remains effective even for non-differentiable rewards.

1 Introduction

Discrete diffusion generators are naturalness-oriented, so training-free inference-time steering combines reward-gradient guidance with search. GRAS targets weaknesses in this combined regime through variance-reduced proposals and adaptive selection without extra denoiser cost.

  • Motivation: Discrete diffusion generates sequences by iteratively unmasking tokens after training to reverse MASK corruption.These models scale well and perform strongly on language, code, regulatory DNA, and proteins.
  • Motivation: Inference-time steering keeps the pretrained generator frozen and biases sampling toward sequences that are both natural and high-reward.It avoids retraining for each new reward, unlike fine-tuning or reinforcement learning.
  • Prior approaches: Search selects high-scoring trajectories, while gradient guidance biases sampling using reward gradients; these families are routinely combined.Prior combined methods include TDS and TreeG’s gradient variant.
  • GRAS: Prior combined methods use fixed-temperature resampling, which can let one particle take all the mass, whereas GRAS changes the proposal and selector.GRAS uses a Jacobian-free proposal, variance-reduced correction, and adaptive-temperature selection.
  • GRAS: GRAS attributes reward gains to proposal and selector components while making denoiser-call budgets explicit.The accounting treats cheap proposals as a strength and exposes the cost of heavier value estimates.
  • GRAS: The guided proposal captures most reward cheaply before saturating, while adaptive resampling supplies remaining reward at a mode-collapse cost.The stated contributions include zero-cost variance reductions and the AMC adaptive-temperature rule.

2 Background

Masked discrete diffusion reverses token masking through iterative unmasking, while reward tilting favors high-reward sequences without abandoning generator likelihood. Guidance and search provide complementary approximations to the resulting soft-value reweighting.

  • Masked discrete diffusion: Masked diffusion replaces sequence tokens with MASK during corruption and generates by progressively revealing them from an all-MASK sequence.Once revealed, tokens are frozen, and masked positions are revealed independently from clean-token predictions.
  • Masked discrete diffusion: At each denoising step, a network outputs clean-token categorical distributions from the partially masked state.The schedule determines the per-step reveal probability.
  • Reward tilting: Reward-tilted sampling trades generator likelihood against reward, with β→∞ recovering the generator and β→0 concentrating on reward maximizers.The exact reverse-process construction uses a Doob h-transform weighted by a soft value function.
  • Reward tilting: The soft value is the log-sum-exp expected future reward from a noisy state, but it depends on clean-sequence rewards and must be estimated at noisy states.Search and gradient guidance are complementary approximations to this reweighting and can be composed.
  • Inference-time methods: Search-based methods reweight or select particles using value estimates, whereas gradient guidance edits the reverse kernel through reward gradients.For masked diffusion, gradient guidance tilts clean-token logits; its shared challenge is estimating the reward gradient.

3 Method

GRAS combines zero-cost variance reduction for guided proposals with AMC, an adaptive-temperature selector for search. Its proofs isolate adaptive scaling as the selector’s active mechanism, while experiments and analysis expose both reward gains and diversity costs.

  • Method overview: GRAS uses a variance-reduced guided proposal and adaptive-temperature selector as drop-in changes to an existing guided sampler.The proposal and selector are designed without additional denoiser calls.
  • 3.1 A Better Proposal: Zero-Cost Variance Reduction: For differentiable rewards, Gumbel–Rao Rao–Blackwellizes reveal noise without increasing estimator variance.The conditional truncated-Gumbel sampler preserves the mean and does not increase variance.
  • 3.1 A Better Proposal: Zero-Cost Variance Reduction: For non-differentiable rewards, PG-RO subtracts a leave-one-out reward baseline from the score-function estimator using the same reward evaluations.It requires no relaxation, Gumbel temperature, model Jacobian, or additional forward pass.
  • 3.2 A Better Search: Adaptive-Temperature Selection: AMC standardizes per-step value increments, replacing fixed temperature α with α_eff = ασ and flooring σ in degenerate cases.Cached values make this a one-line SMC change with no denoiser call; the scaling makes α reward-scale invariant.
  • 3.2 A Better Search: Adaptive-Temperature Selection: Only the adaptive scale affects selector behavior: centering or leave-one-out baselines are invariant up to O(1/K), while the effective temperature changes.The predicted adaptive-temperature effect tracks σ_t, whereas selector leave-one-out centering has a null effect.
  • 3.2 A Better Search: Adaptive-Temperature Selection: The guided filter has a well-defined target under bounded rewards and unbiased resampling, with γ controlling the base measure and α controlling tilt temperature.AMC is consistent for an annealed reweighted target, with population estimates contributing O(1/K) bias; constant σ recovers the fixed-temperature target.

4 Experiments

GRAS is evaluated on regulatory DNA and protein design against training-free, guided, search-based, combined, and fine-tuned baselines. Across both domains, variance-reduced proposals and AMC improve reward, while fidelity remains supported by held-out or structural metrics.

  • Experimental setup: GRAS is evaluated on regulatory DNA and protein inverse-folding benchmarks using shared discrete-diffusion setups and reward-specific fidelity metrics.DNA uses predicted HepG2 activity with held-out ATAC and naturalness measures; protein uses predicted stability with refolding-based structural measures.
  • Baselines: The comparison includes fine-tuning, gradient guidance, search, and guidance-plus-search baselines reproduced under a common setup.The closest combined baselines are SMC-DDM and TreeG.
  • DNA results: On DNA, the guided proposal beats pure-search baselines, variance reduction improves it without extra denoiser calls, and AMC yields the best training-free reward.GRAS also exceeds the reward-fine-tuned model and SMC-DDM, while using a fraction of SVDD’s denoiser calls; the non-differentiable setting follows the same pattern at lower reward.
  • DNA results: DNA fidelity improves with reward for held-out accessibility, while likelihood and motif match exhibit the expected reward–naturalness tension.The accessibility result is interpreted as generalization rather than reward hacking because the ATAC oracle is held out from guidance.
  • Protein results: On protein, GRAS attains the best training-free stability above the reward-fine-tuned model in differentiable and non-differentiable settings.GRAS-GR leads joint structure-and-stability success at 82.1% versus DRAKES at 78.6%, while the well-folded rate holds.

5 Ablation Study and Analysis

Ablations attribute gains to both proposal variance reduction and AMC, with AMC providing most of the remaining reward after the proposal saturates. Its adaptive temperature compensates for changing value spread, but selection causes substantial diversity loss.

  • Adaptive-temperature mechanism: AMC scales resampling temperature with per-step value spread σt, softening selection early and sharpening it late, unlike plain SMC’s fixed temperature.The reshaped schedule ranks up > constant > down, and σt falls 4.9× during denoising.
  • Ablation study: Every GRAS component helps: variance reduction improves the proposal at no extra denoiser cost, plain SMC adds a small increment, and AMC provides most remaining gain.AMC recovers about half the proposal-quality gap on both domains, including for the weaker non-differentiable proposal.
  • Ablation study: Swapping SMC-DDM’s selector for AMC raises reward and every fidelity metric at both Monte-Carlo sample counts.This isolates the improvement to the selector rather than the sample budget and shows transfer to other guided samplers.
  • Adaptive-temperature mechanism: Standardization benefits the selector through the scale ασt, while a leave-one-out baseline has a null selector effect at O(1/K).The observed gain tracks σt, supporting adaptive temperature as the active ingredient.
  • Mode collapse under selection: AMC collapses diversity: nearest-neighbour distance falls 1000× while mean pairwise Hamming distance stays flat.Reducing resampling frequency raises uniqueness from 0.271 to 0.990 for a 0.73 reward drop.

6 Conclusion

GRAS combines variance-reduced guided proposals with adaptive-temperature selection as a training-free method for masked discrete diffusion reward alignment. It achieves strong reward across DNA and protein design, including non-differentiable rewards, while exposing a mode-collapse cost and theoretical scope limits.

  • Conclusion: GRAS combines a variance-reduced guided proposal with an adaptive-temperature selector as drop-in changes to an existing masked discrete-diffusion sampler.The method targets training-free reward alignment without adding denoiser cost.
  • Conclusion: GRAS attains the best training-free reward on regulatory DNA and protein design, matches or surpasses a reward-fine-tuned model, and remains effective for non-differentiable rewards.The reported benefit comes with a mode-collapse cost that the paper surfaces.
  • Limitations: Selection concentrates particles and trades reward for diversity, while the guarantees are asymptotic in K and target a reweighted rather than exact tilted posterior.These limitations define the method’s practical and theoretical scope.

A Deferred Proofs

The proofs formalize two invariances and convergence properties: baseline subtraction does not materially alter resampling, while the adaptive particle systems converge under boundedness assumptions. They also position the method against prior guidance-and-search combinations.

  • Propositions: Constant baselines cancel from resampling softmax weights, leaving the law dependent only on centered increments.This invariance is exact for any constant baseline.
  • Propositions: Leave-one-out baselines change the resampling law only up to O(1/K), while standardization acts through replacing α with ασ.The adaptive scale, rather than mean subtraction, is the operative selection change.
  • Propositions: Under bounded rewards, the guided filter’s self-normalized estimator converges almost surely to the reward-tilted target πγ as particle count K grows.The result assumes unbiased resampling and the stated terminal identity.
  • Propositions: AMC converges almost surely to a well-defined annealed target with O(1/K) adaptation bias, and approaches πγ when ασt becomes constant.The analysis treats AMC as an adaptive Feynman–Kac particle system.
  • Positioning: Unlike prior combined methods, GRAS uses adaptive selection rather than fixed-temperature resampling or hard top-A beam selection.The paper explicitly separates proposal reward from selector reward and reports denoiser-call budgets.

B.2 Guidance Made Explicit: SMC-DDM, TreeG-G, and Ours Side by Side

The comparison makes guidance differences explicit: GRAS truncates the denoiser Jacobian and inserts the correction in clean-logit space, while also preserving the intended reveal schedule. This placement avoids schedule leakage caused by post-hoc tilting.

  • Which derivative: GRAS keeps the reward and closed-form softmax factors but replaces the denoiser Jacobian with the identity, eliminating backward differentiation through the denoiser.The surviving correction is a logit-space vector that can be added directly to clean logits.
  • Where the tilt is inserted: GRAS inserts the exponential tilt inside the clean-logit softmax and leaves the MASK branch untouched, unlike post-hoc full-support renormalization.The numerators match across forms; the difference is where the normalizer is applied.
  • Schedule leakage: Post-hoc tilting changes the reveal probability in a state- and step-dependent way, with guidance strictly accelerating unmasking under the stated conditions.The leakage grows with guidance scale and depends on the clean prediction distribution.
  • Schedule leakage: The GRAS proposal is invariant to additive shifts in the guidance signal, preventing arbitrary reward baselines from changing the unmasking schedule.Under post-hoc tilting, the same shift can drive the schedule toward revealing everything or nothing.
  • Why placement matters: Preserving the schedule matters because earlier token commitment occurs when guidance is strongest and clean predictions are least reliable.The paper states that equation (43) removes this discrepancy at the proposal level.
  • Novelty boundary: The logit-space placement is inherited from GILC; GRAS’s distinct contribution is retaining the per-position product while acting through adaptive selection.The paper narrows its novelty claim to the combined sampler comparison and selection mechanism.

C Implementation Details

Experiments use a shared evaluation setup across regulatory DNA and protein design, with fixed denoising schedules, particle counts, guidance settings, and Monte Carlo sample sizes. Baselines are reproduced from official repositories where available.

  • Common setup: Experiments use T=128 denoising steps on DNA and T=50 on protein, with N=20 particles or candidates for search methods.These settings define the shared benchmark protocol.
  • Guidance settings: DNA uses γ=11000, n=10, M=8, and α=0.5, while protein uses γ=1000, n=20, M=8, and α=0.5.M denotes Gumbel–Rao samples in the reported setup.
  • Baseline evaluation: SMC-DDM and TreeG-G run from official repositories under matched denoiser-call budgets, but TreeG-G is omitted from protein comparisons because no official protein pipeline exists.Table 6 reports the denoiser-call budget for every evaluated method.

D Sampler and Selector Algorithms

The sampler and selector algorithms combine a guided reveal with reduced-variance proposal estimators and AMC’s adaptive-temperature selection. The correction is computed from quantities already available during the guided reveal, avoiding extra denoiser passes.

  • Algorithms: The guided sampler replaces base-model unmasking with guided reveals using either reduced-variance estimator, while AMC changes vanilla SMC through one adaptive-temperature weighting line.The algorithms use denoiser, reward, guidance-scale, Monte Carlo, particle, temperature, and step inputs.
  • Guided correction: The guided filter’s log-weight correction restores the reward-tilted target for every guidance scale γ while preserving the masked/unmasked split.The proposal ratio is supported on positions revealed at that step.
  • Compute: The correction adds no denoiser call because all quantities required by it are already materialized by the guided reveal.Under AMC, it is added after standardizing the per-particle increments.
  • Proposal estimators: Gumbel–Rao reduces variance for differentiable rewards by replacing a single-draw straight-through Jacobian with an M-sample conditional average.M=0 recovers the base CORRECTION-DB estimator.
  • Proposal estimators: For non-differentiable rewards, CORRECTION-PG uses an unbiased score-function estimator, and RLOO subtracts each sample’s leave-one-out reward mean.The estimator requires neither a relaxation nor a model Jacobian.
  • Variance reduction: RLOO improves the proposal over the full-group baseline, whereas selector standardization is inert up to O(1/K) except for its effective temperature change.This motivates placing RLOO on the proposal and standardization on the selector.

F Full-Result Ablation Study

The full ablation traces how each guided proposal and search component contributes to reward across DNA and protein design. It shows that adaptive selection and variance reduction provide distinct gains, with selection closing a substantial proposal-to-final gap.

  • Ablation design: The full ablation evaluates every build-up stage of both guided proposals across all metrics, with final settings marked separately.The study uses mean±std results from three seeds with N=20 for DNA and protein.
  • DNA results: 6.956±0.188 Pred-med is reached by GRAS-RO after GILC-PG proposal, plain SMC, AMC, and RO stages.The DNA ladder increases from 4.859±0.074 for the proposal to 6.956±0.188 for GRAS-RO.
  • Selection contribution: Selection closes 51.4% of the proposal-to-final gap on DNA and 48.3% on protein, despite the non-differentiable proposal being up to 4× weaker.Recovery is defined as 1 − (final gap)/(proposal gap), comparing DB−PG differences before and after AMC.
  • Selection contribution: Raising the policy-gradient proposal budget to n=80 reaches only 6.29 Pred-med on DNA, below AMC’s 6.61 at 25T.The result indicates that additional proposal effort does not substitute for adaptive selection in this ablation.

G.2 Secondary Checks of Adaptive Tempering

Secondary checks support adaptive tempering as the active selection mechanism: reward spread changes substantially during denoising, increasing schedules outperform fixed or decreasing schedules, and reward scaling requires little retuning. The checks also expose proposal-specific variance-reduction behavior and a diversity limitation under selection.

  • Adaptive tempering: The per-step reward spread drifts 4.9× along the denoising trajectory, motivating temperature adaptation.AMC divides by the changing per-step spread rather than using one fixed calibration.
  • Adaptive tempering: Increasing > constant > decreasing is the effective-temperature schedule ranking, with the increasing schedule best on protein and DNA.A fixed temperature implicitly over-selects early when reward spread is largest; the main results therefore use the increasing schedule.
  • Reward-scale invariance: A 3× reward scaling moves Pred-med by 0.078 under AMC versus 0.471 under plain SMC, so α transfers without retuning.Standardization removes reward-scale effects while the softmax shift cancels.
  • Variance-reduction checks: The leave-one-out baseline improves the proposal by +0.346 but is inert or harmful on the selector, matching the proposition’s prediction.Applying it to both stages cancels the gain, yielding 6.66 versus 6.61 without either baseline.
  • Variance-reduction checks: Gumbel–Rao improves monotonically with sample count on the proposal and nearly saturates by M=8 on DNA, while its protein interaction with AMC is mixed.On protein, it helps the proposal but mildly hurts after AMC; the paper reports this interaction without explaining it.
  • Protein structure versus stability: On protein, the selector raises ddG from 0.85 to 1.29 while joint success increases from 74.5% to 82.1% and well-folded rate remains near 90%.GRAS-GR also surpasses DRAKES on ddG and edges it on joint structure-and-stability success, 82.1% versus 78.6%.
  • Diversity limitation: Nearest-neighbour distance collapses 1000× across 47 configurations while mean-Hamming stays flat, revealing metric blindness to mode collapse.Reducing resampling frequency raises uniqueness from 0.271 to 0.990 at a 0.73 reward drop, exposing a reward–diversity trade-off.
Loading 2608.26585v1…