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 steering with discrete diffusion combines guided proposals and search, but current variants use noisy single-sample gradient estimates and fixed resampling temperatures. GRAS reduces proposal variance and adapts resampling, achieving the best training-free rewards on regulatory DNA and protein design while matching or surpassing reward-fine-tuned models.

  • Problem

    Combined inference-time proposal-and-search methods lack clear attribution of reward and denoiser cost, while future reward must be estimated at noisy states.

    Source: §1 · p. 2, §2 · p. 3

  • Method

    GRAS combines variance-reduced guided proposals with an adaptive-temperature selector that standardizes per-step values into group-relative advantages.

    Source: §1 · p. 2, §3.2 · p. 5

  • 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.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The guided proposal captures most reward cheaply, while adaptive resampling delivers the remaining reward.

    Source: §1 · p. 2

    Selection concentrates the particle set, trading sample diversity for reward.

    Source: §6 · p. 8

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

This introduction motivates training-free reward steering of frozen discrete diffusion models and presents GRAS as a combination of variance-reduced guided proposals and adaptive resampling. It also frames the work as accounting for the reward and denoiser cost contributed by each component.

Introduction

  • Training-free steering keeps the pretrained generator frozen and biases sampling toward sequences that are both likely under it and high-reward, avoiding reward-specific retraining.The motivating rewards include enhancer expression, protein stability, and preference-model scores.Source: §1 · p. 1
  • Search methods select among multiple trajectories and gradient guidance biases sampling by reward gradients; prior work already combined them in TDS and TreeG’s gradient variant [12].GILC adds a straight-through reward gradient to clean-token logits [10] [29].Source: §1 · pp. 1–2
  • The paper separates proposal and selector contributions and reports denoiser-call budgets alongside reward; SVDD’s Rao–Blackwell value estimate raises its call count roughly fourfold.Table 4 positions GRAS as the only listed method combining training-free steering, gradients, search, and adaptive selection.Source: §1 · p. 2
  • The guided proposal captures most reward cheaply before saturating, while adaptive resampling supplies the remaining reward at a mode-collapse cost not detected by standard metrics.Figure 1 contrasts prior fixed-temperature resampling with GRAS’s adaptive-temperature selection.Source: §1 · p. 2, Figure 1 · p. 2
  • GRAS reduces guided-proposal variance with Gumbel–Rao for differentiable rewards and a leave-one-out baseline for non-differentiable rewards, then combines them with AMC, an adaptive-temperature selector whose scale is the only active ingredient.The proposal corrections add no denoiser calls; AMC standardizes per-step values into a group-relative advantage.Source: §1 · p. 2

2 Background

The background introduces masked discrete diffusion and reward-tilted sampling, then frames search and gradient guidance as complementary ways to approximate reward-aware reverse denoising. It also describes GILC’s straight-through gradient-guidance proposal for differentiable rewards.

Masked Discrete Diffusion

  • Masked diffusion generates sequences by progressively revealing tokens from an all-MASK state over T reverse steps.The forward process replaces tokens with MASK; the learned reverse process iteratively unmasks them.Source: §2 · p. 2
  • Each reverse step reveals masked positions with schedule-set probability, after which tokens remain frozen and positions are revealed independently from their clean predictions.The per-step reveal probability is at = (αs − αt)/(1 − αt).Source: §2 · p. 3

Reward-Tilted Sampling and the Soft Value Function

  • Reward-tilted sampling trades off generator likelihood against reward, and exact reverse-process sampling uses a Doob h-transform based on the soft value of future reward.The soft value is difficult to evaluate at noisy states because it depends on rewards of clean sequences.Source: §2 · p. 3
  • Search-based selection reweights or selects among base-kernel trajectories, while gradient guidance edits the reverse kernel; combining them allocates search over guided proposals.The paper adopts and analyzes the combined regime used by TDS, SMC-DDM, and TreeG.Source: §2 · p. 3

Differentiable Rewards and the GILC Proposal

  • For differentiable rewards, GILC uses a Gumbel–Softmax straight-through estimator and approximates the model Jacobian as identity to estimate discrete-sequence reward gradients [15] [2] [10].It adds the resulting guidance to masked-position logits; the guidance scale γ plays the role of 1/β.Source: §2 · p. 3

3 Method

GRAS combines variance-reduced guided proposals with adaptive-temperature selection as drop-in changes to an existing guided sampler. Its proposal reduction is matched to reward differentiability, while its selector adapts resampling to the per-step spread of values.

3.1 A Better Proposal: Zero-Cost Variance Reduction

  • For differentiable rewards, GR replaces the straight-through Jacobian with its conditional expectation given the hard sample, reducing reveal noise without an extra pass while preserving the mean and not increasing variance.The conditional reveal Gumbels have a closed-form truncated-Gumbel sampler.Source: §3.1 · p. 4
  • For non-differentiable rewards, PG-RO uses a score-function estimator with a leave-one-out reward baseline, reusing the same n reward evaluations without an extra forward pass.Each sample’s reward is centered by the mean reward of the other n −1 samples.Source: §3.1 · p. 4

3.2 A Better Search: Adaptive-Temperature Selection

  • AMC replaces fixed-temperature SMC resampling with αeff = ασ, adapting to per-step reward spread and avoiding overly cold or hot resampling without an extra denoiser call.The spread σ is floored at a small ϵ in the degenerate all-equal case; cached values make the change one line of vanilla SMC.Source: §3.2 · p. 5
  • Standardization’s centering does not change the resampling law for a constant baseline and affects it only up to O(1/K) for leave-one-out centering, leaving adaptive temperature as its operative effect.The method predicts that the selector’s benefit comes from adapting to σt, not from adding a leave-one-out baseline.Source: §3.2 · p. 5
  • With bounded rewards and an unbiased resampling scheme, fixed-temperature filtering consistently targets πγ, with γ controlling the guided base measure and α controlling the tilt.The guided filter’s reward-tilted proposal and reward-based weights act on orthogonal factors.Source: §3.2 · p. 5
  • AMC has a well-defined annealed target with O(1/K) bias from population estimates; when σt is constant it recovers πγ, but varying σt breaks the fixed-temperature filter’s telescoping exactness.Thus AMC intentionally targets a reweighted distribution when the per-step spread varies.Source: §3.2 · p. 6
  • GRAS combines a variance-reduced proposal with AMC selection: GRAS-GR for differentiable rewards and GRAS-RO for non-differentiable rewards, with no added denoiser cost.The method is designed as three drop-in changes to an existing guided sampler.Source: §3 · p. 3, §3.2 · p. 6

4 Experiments

Across regulatory DNA and protein design, GRAS achieves the strongest training-free reward reported, with protein results also showing strong joint structure-and-stability success. The experiments evaluate reward alongside fidelity to held-out or structural criteria.

4 Experiments

  • The comparison includes fine-tuning, gradient guidance, search, and guidance-plus-search methods, reproduced from official repositories under a common setup.The combination baselines closest to GRAS are SMC-DDM and TreeG.Source: §4 · p. 6

4.1 DNA Enhancer Sequence Design

  • DNA experiments use a masked discrete diffusion model with 128 reverse steps, guide on predicted HepG2 activity, and reserve held-out ATAC activity for evaluation.Reward is median predicted HepG2 activity (Pred-med); fidelity includes ATAC pass rate, 3-mer and JASPAR-motif correlations, and approximate log-likelihood (App-LL).Source: §4.1 · p. 6
  • On DNA, the guided proposal beats pure-search baselines, variance reduction improves it at no extra denoiser cost, and AMC achieves the best training-free reward, above the fine-tuned model and SMC-DDM.The held-out accessibility score rises with reward, while likelihood and motif match show the expected reward–naturalness tension.Source: §4.1 · pp. 6–7, Table 1 · p. 7

4.2 Protein Sequence Design

  • Protein experiments condition discrete diffusion on a target backbone, guide on oracle-predicted ddG stability, and re-fold sequences to assess structural fidelity.Metrics include median predicted stability (Pred-ddG), stabilizing-design fraction, self-consistency RMSD (scRMSD), well-folded fraction, and joint structure-and-stability success.Source: §4.2 · p. 7
  • On protein, GRAS achieves the best training-free stability in both reward settings; GRAS-GR leads joint structure-and-stability success at 82.1% versus DRAKES’s 78.6%, while the well-folded rate holds.The guided proposal beats pure-search baselines, and variance reduction improves it at no extra denoiser cost.Source: Table 2 · p. 7, §4.2 · p. 7

5 Ablation Study and Analysis

The ablations show that variance reduction improves guided proposals, while adaptive selection supplies most of the remaining reward gain. AMC’s adaptive temperature addresses changing value spread, but selection can sharply reduce diversity even when a common diversity metric appears unchanged.

Ablation study

  • The variance reducer improves proposals at no extra denoiser cost, plain SMC adds a small gain, and AMC supplies most remaining reward gain on both paths.AMC recovers about half the proposal-quality gap on both domains, including with the weaker non-differentiable proposal.Source: §5 · p. 8, Table 3 · p. 8

Effect of AMC

  • Swapping only SMC-DDM’s selector for AMC raises reward and every fidelity metric at both M counts, isolating a selector effect that transfers to other guided samplers.M denotes SMC-DDM’s reward-expansion sample count.Source: §5 · p. 8, Table 3 · p. 8

Adaptive-temperature mechanism

  • As σt falls 4.9× during denoising, a fixed temperature over-selects early; AMC divides out this changing spread, and reshaped schedules rank up > constant > down.By Proposition 1, standardization acts only through the scale ασt.Source: §5 · p. 8, Figure 2 · p. 8

Mode collapse under selection

  • Selection collapses diversity: mean pairwise Hamming distance stays flat while nearest-neighbour distance falls 1000×; less frequent resampling raises uniqueness from 0.271 to 0.990 for a 0.73 reward drop.The collapse occurs with AMC and search baselines; resampling frequency trades reward for diversity.Source: §5 · p. 8

6 Conclusion · Appendix Contents

GRAS combines a variance-reduced guided proposal with adaptive-temperature selection as drop-in sampler changes, achieving strong training-free reward alignment while exposing a diversity trade-off. The appendix lists proofs, related work, implementation and algorithm details, estimators, ablations, additional analysis, and metric definitions.

6 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 conclusion notes these gains come at a mode-collapse cost.Source: §6 · p. 8
  • GRAS combines a variance-reduced guided proposal with adaptive-temperature selection as drop-in changes to an existing masked discrete diffusion sampler.The proposal and selector form the method’s two components.Source: §6 · p. 8
  • Selection concentrates particles, creating a reward–diversity trade-off with resampling frequency; guarantees are asymptotic in particle count K and target a reweighted rather than exact tilted posterior.The conclusion points to §5 for the resampling-frequency trade-off and Proposition 3 for the guarantee.Source: §6 · p. 8

Appendix Contents

  • Appendix A covers deferred proofs, while Appendix B presents related work and positioning.Source: Appendix Contents · p. 11
  • Appendices C–E list implementation details, sampler and selector algorithms, and base proposal estimators.Source: Appendix Contents · p. 11
  • Appendices F–H list the full-result ablation study, additional analysis, and metric definitions.Source: Appendix Contents · p. 11

A Deferred Proofs · B Related Work and Positioning · B.1 Anatomy of Guidance and Selection in Prior Samplers

The paper positions GRAS as a guidance-and-search method whose main distinction is adaptive selection, and gives proofs for the resampling law, guided-filter target, and AMC convergence. Its related-work analysis distinguishes GRAS’s guidance and selection mechanisms from prior samplers.

A Deferred Proofs

  • The resampling law is exactly invariant to a constant baseline and invariant to a leave-one-out baseline up to O(1/K); standardization’s operative effect is replacing α with ασ.The proof also shows the constant-baseline law depends only on centered increments.Source: A Deferred Proofs · pp. 11–12
  • With bounded reward, unbiased resampling, and the terminal-value identity, the guided filter converges almost surely to πγ ∝ qγ e^(r/α); at γ = 0, its target is p⋆ with β = α.The proposition states convergence as particle count K → ∞.Source: A Deferred Proofs · p. 12
  • AMC’s self-normalized estimator converges almost surely to πAMC with O(1/K) adaptation bias, and πAMC approaches πγ as σt approaches a constant.The stated assumptions include bounded reward and ασt ≥ ϵ > 0.Source: A Deferred Proofs · p. 13

B Related Work and Positioning

  • Prior discrete-diffusion guidance includes posterior-mean classifier guidance / DPS [7], exact CTMC guidance [22], GILC’s clean-logit correction [10], and entropy-aware guidance [29].Gradient-free methods SGDD [6], CSMC [25], and IterRef [18] use reward in accept/reject decisions rather than a gradient.Source: B Related Work and Positioning · pp. 13–14
  • Among combined samplers, GRAS replaces fixed-temperature SMC selection and TreeG’s hard top-A beam [12] with adaptive per-step temperatures, which the paper says deliver reward after the proposal saturates.The positioning table contrasts methods by whether they guide, select, or combine both.Source: B Related Work and Positioning · p. 14, Table 4 · p. 14
  • Learning-based steering includes DDPP [26] and fine-tuning through sampled-trajectory rewards in DRAKES [31] or offline preferences in SDPO [13], trading inference-time flexibility for a trained model.The paper compares against DRAKES as the fine-tuned reference on its benchmarks.Source: B Related Work and Positioning · p. 14

B.1 Anatomy of Guidance and Selection in Prior Samplers

  • GRAS differentiates reward only on clean logits η and uses the identity approximation ∂η/∂zt ≈ I to avoid a backward pass through the denoiser.TDS differentiates a continuous score in xt, while SMC-DDM and TreeG-G linearize reward in discrete state zt; all three differentiate through the denoiser.Source: §B.1 · p. 15, B Related Work and Positioning · p. 15
  • TreeG keeps top-A candidates without resampling or importance weights, whereas TDS and SMC-DDM use proposal correction log(pθ/q), which GRAS omits by design.Among particle filters, GRAS alone uses an adaptive rather than fixed resampling temperature.Source: §B.1 · p. 15, B Related Work and Positioning · p. 15

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

The section makes the gradient and proposal differences explicit: GRAS truncates a known gradient expression and inserts its correction in clean-logit space, preserving the reveal schedule unlike post-hoc kernel tilts. It distinguishes this placement from GRAS’s contribution of retaining the per-position product and acting through selection.

(i) Which derivative

  • The retained gradient is a logit-space correction that can be added directly to the logits defining the reveal distribution, unlike a state-space correction applied externally.The approximation ∂η/∂zt ≈ I is exact only for an identity denoiser; it changes the target question from how to move the noisy state to how to move the clean logits.Source: §B.2 · p. 16

(ii) Where the tilt is inserted, and why the schedule notices

  • GRAS tilts the clean-prediction softmax while leaving MASK untouched; post-hoc TDS [32] and SMC-DDM [23, Eq. 8] normalize the tilt across revealed tokens and MASK.The revealed-token numerators match; the distinction is whether the normalizer applies only to the revealed branch or across both branches.Source: §B.2 · p. 16
  • GRAS’s reveal probability is invariant to adding a constant to the guidance, whereas the post-hoc form’s reveal probability can approach one or zero as that constant grows or falls.At the natural gauge, nonconstant guidance systematically accelerates unmasking under the post-hoc form; this can commit tokens earlier than the schedule intends, and its correction depends on particle reweighting.Source: §B.2 · pp. 16–17

(iii) Is putting the correction in the logits ours alone?

  • Logit-space placement comes from GILC [10], not GRAS; GRAS’s contribution is keeping the per-position product intact and acting through selection.Among the training-free gradient-plus-search samplers discussed, the section says none places the correction in logit space.Source: §B.2 · p. 17

C Implementation Details

Experiments use a shared reproduction protocol with fixed denoising and search budgets, while guidance and selection settings vary by task. The implementation also applies Gumbel–Rao proposals and matches search-method denoiser-call budgets where applicable.

C Implementation Details

  • DNA uses T=128, γ=11000, n=10, M=8, and α=0.5; protein uses T=50, γ=1000, n=20, M=8, and α=0.5.Search methods use N=20 particles/candidates; results average three seeds, with 640 DNA sequences generated per seed. Guidance and held-out evaluation oracles follow Wang et al. [31].Source: C Implementation Details · p. 17
  • Gumbel–Rao lifts the protein proposal from 0.64 to 0.70 for free, while its final effect is within noise, so M=8 is retained for consistency with DNA.The final effect is reported in Table 12.Source: C Implementation Details · p. 17
  • SMC-DDM [23] and TreeG-G [12] use official repositories on the shared backbone and match the other search methods’ denoiser-call budget; TreeG-G is omitted from protein comparisons.The adaptive-temperature selector uses the up schedule, and Table 6 lists each method’s NFE underlying Tables 1–2.Source: C Implementation Details · p. 17

D Sampler and Selector Algorithms · D.1 Closed Form of the Importance Correction · E Base Proposal Estimators

The appendix specifies the guided sampler and AMC, derives the importance correction, and details the differentiable and non-differentiable base proposal estimators underlying the reduced-variance methods.

D Sampler and Selector Algorithms

  • The guided sampler replaces base-model unmasking with a guided reveal and uses either reduced-variance estimator; AMC flags adaptive-temperature weighting as its departure from vanilla SMC.The SVDD reference is the published selector run unmodified on the same guided reveal.Source: D Sampler and Selector Algorithms · pp. 17–18

D.1 Closed Form of the Importance Correction

  • Adding log(pθ/qγ) to the log-weights restores p⋆ for every γ, with the correction supported only on positions revealed at that step.All quantities in the correction are already materialized by the guided reveal, so it requires no additional forward or backward pass; under AMC it is added after standardization.Source: §D.1 · pp. 18–19

E Base Proposal Estimators

  • CORRECTION-DB and CORRECTION-PG share the mean ∇ηE[r] and can be used interchangeably in Algorithm 3.They are the base estimators reduced by GR and PG-RO, respectively.Source: E Base Proposal Estimators · p. 19
  • The differentiable base estimator averages straight-through gradients from n hard one-hot reward evaluations, but its single-draw Gumbel Jacobian has high variance.Gumbel–Rao replaces that Jacobian with an M-sample conditional average at no extra pass; M=0 recovers CORRECTION-DB.Source: E Base Proposal Estimators · p. 19
  • For non-differentiable rewards, CORRECTION-PG uses reward-weighted REINFORCE scores with a standardized group-relative advantage, requiring no relaxation, Gumbel temperature, or model Jacobian.It is higher variance than the differentiable path and uses RLOO as its proposal baseline.Source: E Base Proposal Estimators · p. 19
  • RLOO outperforms the full-group GRPO baseline on the proposal, while both baselines are inert to O(1/K) on the selector.Accordingly, the non-differentiable final method uses RLOO on the proposal and standardization only on the selector.Source: E Base Proposal Estimators · p. 19

F Full-Result Ablation Study · G Additional Analysis

The appendix gives a full, all-metric ablation of GRAS’s guided-proposal build-up stages and outlines additional checks of adaptive tempering, variance reduction, and reward–diversity trade-offs.

F Full-Result Ablation Study

G Additional Analysis

G.1 Ablation Study … Protein sequence design.

The merged sections test how variance-reduced proposals and adaptive selection affect reward, protein stability, and diversity. They also examine the mechanisms behind adaptive tempering and clarify how reward and fidelity are measured.

G.1 Ablation Study

  • Selection closes 51.4% of the proposal-to-final reward gap on DNA and 48.3% on protein, despite the non-differentiable proposal being up to 4× weaker.Adaptive tempering gains most where plain SMC gains least; increasing the policy-gradient proposal budget to n=80 reaches only 6.29 Pred-med on DNA, below AMC’s 6.61 at 25T.Source: §G.1 · p. 20, Table 8 · p. 21

G.2 Secondary Checks of Adaptive Tempering

  • The per-step spread σ_t drifts 4.9×, and increasing temperature schedules rank above constant, which ranks above decreasing schedules.A fixed temperature implicitly decreases selection pressure over the trajectory, over-selecting early when σ_t is largest; the increasing schedule dominates constant on every protein metric.Source: §G.2 · p. 20, Table 9 · p. 21
  • A 3× reward scaling moves DNA Pred-med by 0.078 under AMC versus 0.471 under plain SMC, supporting transfer of α without retuning.Standardization by σ_t and cancellation of the reward shift in the softmax explain this scale invariance.Source: §G.2 · pp. 20–21, Table 10 · p. 21

G.3 Isolating the Matched Variance Reducers

  • On DNA, proposal RLOO adds +0.346 (t=3.61), while selector placement is inert or harmful and using it on both cancels; Gumbel–Rao saturates near M=8.On protein, Gumbel–Rao helps the proposal but mildly hurts after AMC, an interaction the authors report without explaining.Source: §G.3 · p. 21, Table 11 · p. 22, Table 12 · p. 22

G.4 Protein Structure versus Stability

  • GRAS-GR reaches 1.29 Pred-ddG from 0.85 with selection, while joint success rises from 74.5% to 82.1% and well-folded rate holds at 91.6%.Its joint success rate also edges DRAKES (82.1% vs. 78.6%), and it attains the best training-free ddG.Source: §G.4 · p. 21, Table 8 · p. 21

G.5 Diversity Metric Blindness

  • Across 47 configurations, mean-Hamming stays flat while nearest-neighbour distance collapses 1000×, so mean-Hamming misses mode collapse.The reward–diversity frontier shows that fewer resampling steps recover diversity at a reward cost, providing a user control over that tradeoff.Source: §G.5 · p. 21, Table 13 · p. 22

H Metric Definitions

  • Reward metrics measure optimization toward the guidance objective, whereas fidelity metrics assess whether designs remain realistic rather than adversarial artifacts.Examples include DNA Pred-med for reward and held-out ATAC%, natural-sequence statistics, and generator likelihood for fidelity; protein measures include predicted ddG, scRMSD, and joint success.Source: H Metric Definitions · p. 22, Protein sequence design · p. 22

Protein sequence design.

  • Fewer AMC resampling steps recover diversity, measured by nearest-neighbour distance and uniqueness, at a reward cost.The resampling frequency provides a single control for trading off reward against mode collapse.Source: Table 14 · p. 23

Related papers

Loading 2608.26585v1…