Grab a Coffee: Future-Aware Guidance for Discrete Diffusion with Compiled Objectives

Hua; Xu; Dongxin Li; Gwen Yidou-Weng; Guy Van den Broeck; Wei Wang; Anji Liu

cs.AISep 28, 2026

Grab a Coffee: Future-Aware Guidance for Discrete Diffusion with Compiled Objectives

Authors: Hua, Xu, Dongxin Li, Gwen Yidou-Weng, Guy Van den Broeck, Wei Wang, Anji Liu

Organizations: Edward · Data Science and Analytics Thrust, The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Department of Computer Science, University of California, Los Angeles, USA · School of Computing, National University of Singapore, Singapore

Abstract

Discrete diffusion models generate sequences by iteratively resolving multiple tokens in parallel, offering a flexible alternative to left-to-right generation. However, guiding this process with a sequence-level objective is difficult because the value of one unresolved token depends on the other tokens with which it can form a high-reward sequence. Enumerating all such completions makes the whole guidance computation grow exponentially with the number of unresolved positions. We introduce COFFEE, a plug-and-play framework that avoids this enumeration by separating sequence dependence from the objective. At each diffusion step, a target-free carrier absorbs the marginal token distributions predicted by the denoiser to construct a joint model over the unresolved tokens, while a compiled finite-state model records how their combinations affect the sequence-level preference. Pairing their states allows COFFEE to transfer global preferences to unresolved positions and sample a clean reconstruction without retraining the diffusion model. The same framework supports explicit hard constraints and learned soft objectives. We evaluate COFFEE across multiple symbolic, language, and biological benchmarks, where it achieves strong control results with task-dependent quality and diversity trade-offs. By making objectives available to inference rather than only evaluation, COFFEE brings joint conditioning, completion-weighted guidance, and optimization-based constraints into pretrained neural generation, showing the potential of neural-symbolic methods in diffusion guidance.

Figures & tables

Figure 1: Sequence-level guidance combines preferences with a model of possible completions. (a) Each root-to-leaf path specifies a clean reconstruction with neutral completion weights given by

q_{\theta}

. With two token choices at each of three unresolved positions, explicit enumeration produces

2^{3}

leaves. Calculating quantities in Eq. 1 with such way is therefore hard with exponential complexity. (b) Coffee uses state machines to represent token dependence and record pattern progress. Their product permits histories reaching the same state at the same position to share a continuation calculation, replacing path-by-path enumeration within the compiled model.

Figure 2: Joint states support tractable computation for joint sampling. The upper panel shows how Coffee transforms diffusion marginals into guided joint distributions. With compiled states, Coffee runs DP over a continuation table to obtain probability masses, based on which we can get normalized probabilities. The lower panel uses a toy example to dive into the DP process in the table to provide an intuitive understanding of Eq. 15 and 20 . In the table, columns index sequence boundaries and rows index joint states $(h,a)$ . Each $M_{i}(h,a)$ sums the weights of the remaining token-state paths, and generation paths sharing the same boundary state can therefore reuse one entry. The expanded example sums six consecutive states $(u,h^{\prime})$ to obtain $M_{4}(2,a_{\mathrm{init}})$ . Normalizing US ’s probability mass with the previous quantity gives the joint probabilities of the country.

	Dyck repair		Sudoku		CommonGen		RTP		Dyck repair		Sudoku		CommonGen		RTP		DeepSTARR		APARENT		K562		Protein
Method	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$	Control $\uparrow$	Quality $\downarrow$
Base	3.0	13.5	60.4	0.1	38.6	10.1	53.5	5.2	7.5	14.0	20.4	0.3	40.9	14.4	43.3	5.3	-0.4	728.3	2.3	728.3	44.3	728.3	11.0	1.7
CDD	46.9	14.8	99.3	0.0	46.1	30.4	55.7	17.2	62.5	14.7	90.7	0.0	40.9	18.8	81.7	13.1	3.5	1.2	7.9	682.1	83.7	70.1	10.5	1.7
CDM	0.2	11.7	100.0	0.0	57.2	33.8	65.6	5.4	7.7	13.8	100.0	0.0	41.5	14.5	54.3	4.9	0.6	429.1	10.1	691.9	72.4	250.9	10.2	1.7
D-CBG	14.9	14.3	99.8	0.0	39.4	10.4	62.2	5.0	23.1	14.5	98.7	0.0	30.7	32.7	72.4	7.1	-0.5	1799.7	2.2	1492.4	100.0	1.2	11.0	1.7
DG-TAG	100.0	14.9	100.0	0.0	29.7	9.8	87.9	7.2	100.0	14.9	100.0	0.0	32.9	14.7	66.9	5.4	3.6	2.6	24.3	6.3	97.2	3.4	14.5	1.5

Table 1: This table shows the generation results under full diffusion budgets. The warm Dream and blue LLaDA headers denote the two backbones and the black headers denote biological tasks. Coffee surpasses almost all the baselines on both control and quality metrics over different benchmarks.

Figure 3: Control across denoising budgets

b

(higher is better). Warm Dream and blue LLaDA titles distinguish the language backbones; black titles denote biological tasks. Hollow markers indicate archival or qualified settings documented in Sec. F . The results show that Coffee achieves the highest control effects over almost all the benchmarks, even with only a few diffusion steps.

Figure 4: Proposal trajectories and token-level guidance. (A) Evaluator-category transitions of complete clean proposals from 64 same-ID HepG2 rollouts per arm, with the objective disabled (

W_{C}\equiv 1

) or enabled and the identity carrier fixed. Ribbon widths count trajectories, not probability mass. (B) For a fixed K562 query and prefix at position 25, curves show token probabilities as future objective factors are enabled cumulatively; the state graph and all other factors remain fixed. Local omits future weights; Full includes all. Weights at two already observed downstream tokens reduce the probability of AAAAAA . Bars show incremental log odds of AAAAAA against the other candidates. The horizontal axis is objective position, not diffusion time. (C) CDF of Full-Local total variation over 224 masked positions in nine queries, with conditional TV averaged over prefixes under neutral occupancy. The diamond marks the maximum-TV position selected for (B). Protocols are in Sec. H.3 .

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Instance	Product state	Reward factor
Diffusion-only	identity	$1$
Carrier-only	$h_{i}$	$1$
Learned guidance	$(h_{i},a_{i})$	$\exp\{\lambda R_{\phi}(x)\}$
Hard constraint	$(h_{i},a_{i})$	$\mathbf{1}[a_{L}\in\mathcal{A}_{\rm accept}]$
Strict repair	$(h_{i},\text{bounded stack}_{i})$	$\mathbf{1}[x\in\mathcal{L},\ d_{H}(x,y)=d^{\star}(y)]$

Appendix

Table 2: Special cases of the same graph construction. The identity carrier is an ablation of the full method; the identity reward controller has one state.

Task / backbone	Base	CDD	CDM	D-CBG	DG-TAG	GILC-DB	MDM-VGB	Coffee
Dyck / Dream	6.249	96.966	51.076	82.366	121.766	443.709	406.640	13.940
Dyck / LLaDA	6.237	107.607	76.587	107.587	136.590	469.933	389.544	12.086
Sudoku / Dream	50.592	105.776	67.333	95.105	83.739	111.114	53.027	56.295
Sudoku / LLaDA	69.252	110.038	81.719	110.513	101.191	112.844	71.603–72.364	75.136
CommonGen / Dream	20.121	38.690	111.330	20.272 ‡	80.754	30.454–30.649	624.670	109.432
CommonGen / LLaDA	15.428	68.535	152.613	79.558	76.392	62.749	590.824	65.442

Appendix

Table 3: Complete online generation time in seconds for 64 final outputs at each task’s frozen Full budget. Base is a compute reference; bold identifies the fastest measured non-Base baseline within a GPU environment. BoN-4 is excluded. Static model loading, warmup, disk output, and independent evaluation are excluded. Per-request controller construction, inference, sampling, transfers, and actual extra model calls are included; early stopping is retained. † Measurements from a separate cloud GPU environment are shown for context and must not be divided by primary RTX 4090 measurements in the same cell. ‡ This D-CBG setting failed its predeclared DEV PPL cap; its measured time remains visible. The RTP/LLaDA DG-TAG and GILC ranges reverse order across repeats and use different native external batches. The LLaDA/Sudoku MDM range includes a different-card repeat; the same-GPU comparison with Coffee uses 72.364. w Reused static Graph, with complete generation timed; f first request including Graph capture. Times with different native batches, cohorts, or software stacks are descriptive unless the measurement conditions are matched.

Component	Formal statement	Empirical evidence
Backward DP	finite exact recursion	graph/support sizes
Ancestral replay	exact at $\tau_{a}=1$	sampling rule and seed
Scorer fit	convex fixed-design objective	solver status and validation metrics
Carrier fit	no global-optimum claim	epoch/loss/convergence flag
Diffusion host	frozen	checkpoint and schedule

Appendix

Table 4: Claim boundaries for “convergence.”

Task	Scored scope	Automatic controller construction	Terminal rule
RTP	generated continuation only	Mine and select train-only count features for $\mathbf{1}[\text{continuation toxicity}<0.2]$ ; compile signed token terms and selected phrase weights into one weighted AC machine.	zero for every state
Common Gen	generated sentence	For each record, enumerate standalone base, regular-plural, and possessive forms; tokenize them and compile the reachable lexical-prefix state plus a satisfied-concept bit mask.	$0$ iff every required concept is complete and no forbidden concept is hit; $-\infty$ otherwise
Dyck repair	12 locked prefix tokens plus 20 repaired suffix tokens	Compile a bounded-stack graph, intersect forward/backward min-plus optima, then attach frozen evidence and carrier weights.	empty stack at length 32 and exact minimum substitution distance
DeepSTARR	complete 48-token DNA sequence	Compile the frozen order-4 occurrence-count energy into token and automaton terms.	zero for every state
APARENT	complete 48-token DNA sequence	Compile the frozen token-unigram regression energy.	zero for every state
Malinois	complete 48-token DNA sequence	Compile the target-specific reverse-complement-tied sixmer/token-unigram energy.	zero for every state

Appendix

Table 5: How each task becomes a controller. Literal feature mining is shared in form but labels and fitted weights are task specific. CommonGen, Dyck, and Sudoku have no learned count-logit reward. Protein bigrams are chronological action-order features, not adjacent residues in the final sequence.

Task	Generation problem	Target / constraint	Primary control metric	Quality and audit metrics
RTP	prompt continuation	non-toxic continuation	ToxicBERT continuation toxicity and satisfaction	prompt-conditioned GPT-2-large corpus PPL
Common Gen	concept realization	terminal lexical DFA acceptance	fail-closed all-concept satisfaction	exact/normalized coverage; forbidden hits; unreachable support; conditional PPL
Dyck	fixed-length repair	valid minimum-edit completion	exact validity	all-output Hamming edit and legal-edge rate
Sudoku	clue-preserving completion	valid $9\times 9$ grid	exact validity	clue preservation and output format
DeepSTARR	unconditional DNA generation	developmental activity	official DeepSTARR H5 predictor	native masked LOO and failure count
APARENT	unconditional DNA generation	isoform usage	official APARENT H5 predictor	native masked LOO and diversity diagnostics

Appendix

Table 6: Task and evaluator contracts. The scorer used for guidance is not the sole evaluator.

Method	Implementation	Fitted object and supervision	Online rule	Reported compute
Base	host implementation	none	native reverse diffusion	$b$ reverse updates
CDD	same-backbone port	train-only projection expert	top- $k$ projection with augmented-Lagrangian inner updates	$b$ plus inner-loop count $I$
CDM	sequential Monte Carlo (SMC) port	train-only twist	$P$ particles with fixed proposal, effective sample size (ESS) threshold, and resampling law	task-dependent particle work; see below
D-CBG	same-backbone port	train-only corrupted-state classifier	classifier reweights token alternatives; DNA uses the declared full-vocabulary Taylor form	$b$ updates and candidate width
DG-TAG	method-inspired port	frozen task predictor	rescore top- $k$ successor states and update logits	$b$ updates and predictor calls
GILC-DB	adapted port	train-only task reward model	Monte-Carlo logit correction at every reverse update	$b$ reverse updates; MC samples reported separately

Appendix

Table 7: Executable baseline recipes. “Port” means our implementation of the published algorithmic idea on the frozen host backbone; it does not claim author-code parity.

$b$	Base	Coffee
$8$	6.976	6.809
$32$	27.903	27.238

Appendix

Table 8: Estimated backbone TFLOPs per output in the LLaDA/CommonGen efficiency runs. We use

F\approx 2PS

per forward, nominal

P=8\times 10^{9}

b

observed forwards, and mean padded lengths

S=54.50

(Base) and

53.20

( Coffee ). The length difference reflects batching. They exclude Coffee ’s dynamic programming and are neither profiled nor total FLOPs; other tasks require their own sequence lengths and model-call counts.

Field	COFFEE language default
Diffusion checkpoint	LLaDA-8B-Base
Generation / block length	32 / 32
Denoising budget	32
Evidence temperature $\tau_{e}$	1.0
Carrier-emission temperature $\tau_{m}$	1.0
Ancestral temperature $\tau_{a}$	1.0

Appendix

Table 9: Default COFFEE language configuration before benchmark-specific DEV selection. Temperatures and

\lambda

are independent controls.

Benchmark	Development data	Final panel	Frozen selection rule and interpretation
CommonGen	993 records, with 992 in the corrected common pool	1,497 disjoint records	Maximize exact lexical acceptance under the common PPL cap, followed by coverage, PPL, and strength tie breaks.
RTP	Method-independent development panel	Common 1,024-prompt stress panel selected by Base toxicity	Maximize non-toxicity under the same-budget quality rule. Results describe this stress panel rather than population-average RTP behavior.
Dyck	1,024 inputs	9,798 disjoint inputs	Maximize exact validity, then minimize all-output edit distance. Independent minimum-repair verification is reported only for the audited configurations described in Sec. H.2 .
Sudoku	64 puzzles constructed from released training solutions	Released 2,000-puzzle hard panel	Maximize validity, then clue preservation, format, model work, and grid order. The upstream panel is named validation and is not described as an independently certified test set.
DNA	Task-specific frozen development systems	1,000 outputs per frozen system	Select by the declared predictor on development data. Native LOO and diversity remain separate diagnostics. Comparisons are distributional rather than molecule paired.
Protein-1YCR	Common development geometries	Common 400-geometry panel	Maximize Success@1Å, then minimize mean motif RMSD and evaluator work. Particle methods retain their additional model and evaluator calls.

Appendix

Table 10: Selection contracts for the results discussed in the main paper. No final-panel metric is used to choose a configuration.

Method	Parameters searched on development data	Computation fixed before final evaluation
COFFEE	Language evidence, carrier, and ancestral temperatures in $\{0.3,0.6,0.9\}$ , followed by reward strengths in $\{0.5,1,2,4,8\}$ for declared representatives. Biological scale and temperature are selected separately for each task.	Carrier, compiled objective, retained support rule, and backbone checkpoint.
CDD	Projection strength $\eta\in\{0.2,0.4,0.6,0.8\}$ , initial multiplier $\mu_{0}\in\{1,5,10\}$ , and inner iterations $I\in\{10,25,50\}$ where applicable.	Projection expert, candidate width, and outer update schedule.
CDM	Twist strength $\beta\in\{0.1,0.2,2\}$ for language, with task-specific biological refinements.	Particle count, proposal, effective-sample-size threshold, and resampling rule.
D-CBG	Classifier strength $\gamma\in\{0.5,1,2\}$ , with separately selected full-vocabulary scales for Malinois.	Corrupted-state classifier and candidate rule.
DG-TAG	Rescoring strength $\gamma\in\{1,5,20\}$ , plus the declared $0.1$ CommonGen refinement.	Frozen task predictor and top- $k$ successor set.
GILC-DB	Language strength $\beta\in\{0.1,0.3,1,3,10\}$ , with task-specific biological scale refinements.	Reward model, Monte Carlo sample count, and candidate width.

Appendix

Table 11: Development-time search axes for the methods shown in the main comparison. Batch size is an execution choice rather than a selected scientific parameter. Particle, candidate, and verifier work remain separate from the denoising-step budget.

Main-text result	Additional evidence	Required interpretation
Dyck repair	The compiler constructs the minimum-substitution valid support before probabilistic replay. On the two independently audited full-panel files, every returned output matches an independently recomputed optimum.	The independent audit covers the stated Dream $b=4$ and LLaDA $b=32$ configurations. It is not a blanket empirical certificate for every budget.
Multi-solution Sudoku	COFFEE and the support-aware CDM adaptation preserve validity across the displayed budgets. Sec. H.1 separately compares coherent joint replay with independent exact marginals.	Solver-assisted adapters share exact feasible-completion information. Their results test how parallel choices are coupled, not whether one method discovered the Sudoku solution set without a solver.
CommonGen and RTP	CommonGen uses a record-conditioned hard automaton. RTP uses a learned soft objective on a common 1,024-prompt stress panel.	Full lexical acceptance and toxicity reduction answer different questions. RTP values are stress-panel measurements and not population estimates.
Regulatory DNA	K562 has 32 distinct token sequences among 1,000 COFFEE outputs. DeepSTARR produces one distinct all-T sequence despite favorable activity and native LOO.	Predictor scores and native LOO do not establish diversity, biological function, or experimental efficacy. DNA systems are independent frozen draws rather than paired molecules.
Protein-1YCR	Success@1Å and mean motif RMSD are computed on the common 400-geometry panel.	Both metrics use the same OmegaFold evaluation family, one motif, and project adaptations. They are not independent wet-lab validation.

Appendix

Table 12: Additional evidence and claim boundaries for the results interpreted in the main text.

Method	$b=4$	$b=8$	$b=16$	$b=32$	Full ( $b=64$ )
Joint replay	120	120	120	120	120
Independent exact marginals	68	100	114	118	120
Uniform joint	120	120	120	120	120
D-CBG	13	40	71	113	120
Joint reference	120	120	120	120	120
Maximum cells committed per step	14	7	4	2	1

Appendix

Table 13: Multi-solution Sudoku: joint consistency under parallel commitments. Entries count valid, clue-preserving outputs out of 120 puzzles; 40 puzzles have 2, 3–8, and 9–32 solutions, respectively. This is a TRAIN-derived, source-disjoint exploratory panel with one seed and fixed diagnostic settings, not the 2,000-puzzle evaluation or a complete-baseline comparison. Full (

b=64

) permits at most one new cell per step, with 10–53 actual calls. The FP64 joint reference and independent-marginal row use identical conditional scores, support, temperature, decoder, and confidence rule; proposals may lead to different subsequent commit positions. Uniform joint still uses model-based commit ordering.

Carrier and sampling rule	Lexical coverage (%)	PPL $\downarrow$
H–Neutral	22.01	16.81
H–Feasible	100.00	20.28
H–Full	100.00	18.93
H–Full-marginal	83.98	18.34
I–Neutral	30.99	19.56
I–Feasible	100.00	23.71

Appendix

Table 14: CommonGen component diagnostic on 256 reused validation inputs and three seeds per arm. PPL is the prompt-conditioned GPT-2-large corpus PPL over generated continuations. No output is excluded from scoring.

Explore similar work

Jun 11, 2026cs.LG

A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding

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.

Sophia Tang, Yuchen Zhu, Molei Tao +1

Department of Computer and Information Science, University of Pennsylvania · School of Mathematics, Georgia Institute of Technology · Department of Bioengineering, University of Pennsylvania

Diffusion Models

Ask PDF

May 10, 2026cs.AI

Primal-Dual Guided Decoding for Constrained Discrete Diffusion

Discrete diffusion models generate structured sequences by progressively unmasking tokens, but enforcing global property constraints during generation remains an open challenge. We propose primal-dual guided decoding, an inference-time method that formulates constrained generation as a KL-regularised optimisation problem and solves it online via adaptive Lagrangian multipliers. At each denoising step, the method modifies token logits through an additive, constraint-dependent bias, with multipliers updated by mirror descent based on constraint violation. The bias arises as the optimal KL-regularised projection of the constraint, so the constrained distribution remains as close as possible to the model's unconstrained distribution while still satisfying the constraint. The method requires no retraining and no additional model evaluations beyond standard sampling, supports multiple simultaneous constraints, and provides formal bounds on constraint violation. We evaluate our approach on topical text generation, molecular design, and music playlist generation, showing that a single algorithm instantiated via domain-specific scoring functions improves constraint satisfaction while preserving relevant domain-specific quality metrics.

Federico Tomasi, Dmitrii Moor, Alice Wang +1

Spotify

Diffusion Models Diffusion Dynamics Decoding Biases

Ask PDF

Jun 4, 2026cs.LG

Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction

Controllable generation with discrete diffusion models is often hindered by high computational overhead or the need for retraining. In this paper, we present \underline{\textbf{G}}radient-\underline{\textbf{I}}nformed \underline{\textbf{L}}ogit \underline{\textbf{C}}orrection (\textbf{GILC}), a plug-and-play framework that efficiently estimates guidance signals by repurposing the pretrained denoising network as a variational proxy. To circumvent the gradient instability inherent in high-dimensional discrete spaces, we introduce a Jacobian-free mechanism that directly corrects the clean prediction logits, facilitating stable and effective guidance. Our method accommodates both differentiable and non-differentiable reward functions. Extensive experiments across DNA, protein sequence, and molecular generation tasks demonstrate that GILC achieves state-of-the-art performance without additional training, frequently outperforming fine-tuning approaches.

Hongkun Dou, Zike Chen, Fengji Li +2

Beihang University, Beijing, China · Zhongguancun Academy, Beijing, China.

Diffusion Models Guidance Logit Lens Jacobian

Ask PDF