Source-linked AI summary

Internal Guidance as a Sampling Method for Tilted Diffusion Models

Curtis McDonald, Peter Bartlett

arXiv:2610.29086v1stat.MLcs.LG

TL;DR

Sampling reward-tilted densities from base-distribution samples can require external guidance terms that are intractable or approximated. This paper instead shifts the evaluation point inside a base denoiser and trains it with diagonal rescalings; it demonstrates the approach on small discrete sampling tasks.

  • Problem

    Sampling an exponential tilt from base-distribution samples is challenging when external guidance requires intractable terms or approximations that can accumulate errors.

    Source: §1 · pp. 1–2
  • Method

    The method uses hypercube denoiser symmetries and random diagonal rescalings so one base denoiser can be queried at shifted points for tilted sampling.

    Source: §1 · pp. 2–3
  • Results

    The approach is demonstrated on small Ising reward tilts, Sudoku generation and solving, and conditional MNIST class generation.

    Source: §1 · p. 3
  • Takeaways & Limitations

    A single denoiser can sample a family of tilted densities at inference time without fine-tuning or external guidance.

    Source: §1 · p. 2

    The demonstrations use small discrete targets, with scaling and more complex distributions left for future work.

    Source: §1 · p. 3

Abstract

from arXiv · show

Reward guidance methods sample an exponential tilt of a base density using a diffusion model trained with only samples from the base distribution. Existing methods use what we call external guidance: additive terms in the drift or multiplicative terms on the transition probabilities. We propose internal guidance which leaves the output of the base denoiser unmodified and instead changes the point at which it is evaluated. For discrete distributions embedded in the hypercube $\{-1,1\}^{d}$, the denoiser of any linearly tilted density is exactly the base denoiser evaluated at a shifted and rescaled input. In practice the base denoiser is a neural network trained on noised samples of the base distribution. The required evaluation points for sampling the tilted density would be out-of-distribution for basic noising of the base distribution. We modify the training algorithm with random diagonal rescalings of the data, so that the denoiser network learns both in- and out-of-distribution evaluations of the denoiser relative to the base distribution. A single denoiser network samples binary reward tilts, class labels, and infilling with no fine-tuning and no external guidance term.

1 Introduction

The paper addresses sampling tilted densities from base-distribution samples, where external guidance can require intractable terms or approximations. It proposes internal guidance through shifted denoiser evaluations and hypercube symmetries, with Λ focusing to improve coverage of the needed evaluation points.

1 Introduction

  • Existing external guidance modifies sampling paths or transition probabilities, while exact guidance can require an intractable term or approximations that introduce compounding errors.The paper frames this as a challenge when only base-distribution samples are available.Source: §1 · pp. 1–2
  • Internal guidance samples tilted densities by shifting the evaluation point inside the base denoiser rather than modifying its output.The symmetry maps are explicit shifts and rescalings.Source: §1 · p. 2
  • The network’s accuracy is uncontrolled at symmetry-mapped points that fall outside the base-noising training manifold.The model is trained for accuracy along marginals produced by noising the base distribution.Source: §1 · p. 2
  • For hypercube-embedded discrete distributions, constant norm removes quadratic terms in conditional log densities, enabling additional denoiser symmetries.These symmetries expand the equivalence class of evaluation points and can bring sampling evaluations near the training manifold.Source: §1 · pp. 2–3
  • Λ focusing uses random diagonal rescalings during training so the denoiser learns evaluations away from the typical noising path.The method aims to turn an out-of-distribution learning problem into an in-distribution one.Source: §1 · pp. 2–3
  • The paper tests small Ising reward tilts, Sudoku generation and solving, and conditional MNIST class generation, reserving scaling and more complex targets for future work.Conditional tasks and binary reward tilts can be represented as linear tilts in an extended state space by adding a noised and denoised reward coordinate.Source: §1 · pp. 2–3

2 Setup

The setup embeds discrete states in a hypercube and expresses reward tilting as a linear tilt by adding a reward coordinate. Sampling then uses stochastic localization, whose drift is the denoiser.

2 Setup

  • Finite discrete state spaces can be embedded in {−1, 1}^d, including token, constraint-satisfaction, image, and Ising distributions.The setup considers Gaussian noising of the embedded points.Source: §2 · pp. 3–4
  • Linear tilts and binary reward tilts are unified by promoting the reward value to an additional random coordinate.The reward coordinate is deterministically set by the reward function, and tilting it corresponds to a linear tilt in the extended state space.Source: §2 · p. 4
  • The sampler jointly noises and denoises the token string and reward coordinate, carrying a noisy reward in the added dimension.The extended state has d + 1 coordinates.Source: §2 · p. 4
  • On the hypercube, constant norm removes quadratic terms from the conditional densities, so denoiser evaluations reduce to linear tilts of the base density.The general mean-tilt operator defines the denoiser for the base density and its linear tilts.Source: §2 · pp. 4–5

2.1 Sampling Via Stochastic Localization

  • Stochastic localization uses the denoiser itself as the drift and produces a reverse sampling path for the target density.The paper notes that this makes denoiser-to-denoiser relationships simpler to construct than score-to-score relationships.Source: §2.1 · p. 5

3 Denoiser Symmetries

The paper develops symmetries that make denoiser evaluations for tilted sampling equivalent to evaluations at normalized points. Λ focusing and soft masking shape training evaluations to overlap with the sampling region.

3 Denoiser Symmetries

  • Under VP noising, typical evaluation vectors have norm near √(d + 1), and their coordinate magnitudes have a distribution independent of the base density.The squared norm has expectation d + 1 and variance 2(1 − a^4)(d + 1).Source: §3 · p. 5
  • Stochastic-localization evaluations for a linearly tilted target involve a shift and can have a dominant coordinate, unlike typical VP-noised base evaluations.For reward tilts, the large additive shift is in the reward coordinate.Source: §3 · pp. 5–6
  • Renormalization and Λ focusing provide additional equivalent evaluations, with Λ a diagonal rescaling that preserves the pure-linear-tilt property for hypercube-supported data.The Λ-focusing lemma evaluates the denoiser at the same signal level using a mapped observation.Source: §3 · p. 6
  • Mode 0 shifts mass toward the reward coordinate, while Mode 1 soft masking trains conditionals using random subsets of active coordinates.Soft masking sets excluded coordinates to zero so they do not influence the conditional expectation, providing partial-information training.Source: §3 · p. 7

3.1 Combined Results

  • Combined shifting, normalization, and Λ focusing recast training and tilted-sampling evaluations as base-denoiser evaluations at a norm-√(d + 1) vector and equivalent signal level.Reward-achieving focused training points and tilted-sampling points lie near the same one-hot reward pole, creating overlap between their regions.Source: §3.1 · pp. 7–8

4 Training and Sampling Algorithms

Training randomizes the signal level and diagonal focusing transform while retaining the mean-tilt operator as the regression minimizer. At sampling time, the algorithm shifts and renormalizes the state before querying the denoiser.

4 Training and Sampling Algorithms

  • For any joint training distribution over signal levels and diagonal rescalings, the mean-tilt operator remains the population minimizer of the regression objective.Standard VP noising is the special case with Λ = I and no renormalization.Source: §4 · p. 8
  • Training draws fresh signal levels and Λ-focusing matrices, mixing ordinary VP noising with rescaled or subset-focused examples.With probability 3/16, the described draws use plain VP noising.Source: §4 · p. 9
  • The sampler also uses self-conditioning, a common practice discussed separately to keep the main algorithm description simpler.The paper identifies the self-conditioning choice as an implementation detail.Source: §4 · p. 9
  • At each sampling step, the state is shifted by the tilt vector and renormalized before the network is queried.The terminal continuous state is decoded with one final denoiser call.Source: §4 · p. 9

5 Experiments

The experiments evaluate internal guidance on an Ising model, Sudoku generation and infilling, and conditional MNIST generation, using known log-odds or exact denoisers where available. Results include calibrated sampling on smaller targets and differing behavior across Sudoku scales and self-conditioning modes.

5.1 12-Dimensional Ising Model

  • The Ising experiment learns the ground-truth denoiser with low L2 error across tilt strengths and produces reward rates matching the log-odds curve.The 12-dimensional target is small enough to enumerate, providing a calibration test before larger state spaces.Source: §5.1 · p. 10

5.2 4x4 Sudoku

  • Both 4x4 Sudoku samplers reach 98% reward-achieving solutions at τ ≈3.5, but only the ground-truth-denoiser sampler tracks log-odds up to τ = 5.The network sampler starts to lag at higher tilt strengths.Source: §5.2 · p. 10
  • The network visits 287 of 288 reward-achieving 4x4 Sudoku solutions but fails the uniformity test (χ2 287 = 391.9, p < 10^-4).The exact-denoiser sampler visits all 288 solutions and passes the test (χ2 287 = 280.1, p = 0.60).Source: §5.2 · p. 10

5.3 6x6 Sudoku

  • In 6x6 Sudoku, 10 of 1,000 reward-achieving generations match training puzzles, consistent with no memorization given 3–14 expected matches under uniform sampling.The training set contains 0.88% of the reward-achieving puzzles.Source: §5.3 · p. 11

5.4 9x9 Sudoku

  • For 9x9 Sudoku, carry self-conditioning reaches near 90% graded validity with an 8–12% predicted-reward gap, while two-pass validity plateaus near 60%.Two-pass tracks log-odds in its own reward prediction, but not in graded validity.Source: §5.4 · p. 11
  • Given 11/81 cells, 9x9 infilling maintains the initial values 92% of the time, achieves a reward 87% of the time, and meets both objectives 80% of the time.The cell-fixing tilt is combined with the reward tilt; given cells are not hard constraints.Source: §5.4 · p. 11

5.5 MNIST Digit Images

  • On MNIST, carry at τ = 0.51 produces the requested class 57.8% of the time against an exact 45.9%, while two-pass tracks the classifier-corrected rate through τ ≈1.Both modes produce 95% class-correct images at large tilts; higher tilt strengths exceed the judge CNN’s sensitivity.Source: §5.5 · pp. 12–13

AI Use Statement

The authors report using generative AI across data generation, modeling, coding, experiments, literature work, and writing. They state that proofs and final content were human verified and that AI did not provide critical ingredients for mathematical claims.

  • Generative AI supported synthetic data generation, modeling, proof assistance, methodology, experiments, coding, data work, analysis, and writing.The authors state that the first draft was written through human analysis and that proofs were human verified.Source: AI Use Statement · p. 13
  • The authors say AI did not provide critical ingredients for proofs or propose or refine hypotheses; logical proof steps were initially human designed.They note that “critical ingredients” is subjective.Source: AI Use Statement · p. 13
  • The authors reviewed AI-assisted work and take responsibility for the final content, including AI-assisted text, claims, and artifacts.Source: AI Use Statement · p. 13

A PROOFS

The proofs use the hypercube’s fixed coordinate magnitudes to reduce Gaussian-noised conditional densities to linear tilts. They establish that relevant conditional means can be represented through evaluations of the base denoiser at transformed inputs.

A PROOFS

  • The tilt strength β(a) can be inverted for any nonnegative strength, and rescaling the input vector gives an equivalent denoiser evaluation at a new signal level.The proof specifies the new level as β^-1(β(a)λ).Source: A PROOFS · pp. 18–19
  • On the hypercube, the x-dependent quadratic term in the reverse conditional density is constant and drops out, leaving a pure linear tilt of the base density.The proof uses x_i^2 = 1 for every coordinate.Source: A PROOFS · p. 19
  • With random noising parameters, the squared-loss minimizer depends on the observation only through the renormalized input H(ΛU, a).The proof derives the conditional mean for fixed parameters, then applies the tower property and renormalization.Source: A PROOFS · pp. 20–21

B.1 Visualization of Denoiser Symmetry

This appendix relates score-based diffusion and stochastic localization as reverse processes with different scalings. It explains the OU construction and shows that stochastic localization follows the OU reverse path under a time change and spatial rescaling.

B.2 Stochastic Localization

  • Score-based diffusion is variance preserving, starts from a standard normal draw, and uses the score as its drift; stochastic localization starts at the origin and uses the denoiser as its drift.Both processes arise from reverse-diffusion theory.Source: §B.2 · p. 22

B.2.1 Score-Based Diffusion from the Ornstein–Uhlenbeck Process √

  • The OU forward process yields a variance-preserving combination with signal level squared e^-2t, and its reverse drift is implemented by the score.Source: §B.2.1 · p. 24

B.2.2 The Stochastic Localization Time Scaling √

  • The stochastic-localization reverse path is initialized at zero, and its reverse drift is exactly the denoiser.The associated forward process is driven to zero at its terminal time.Source: §B.2.2 · p. 24
  • Matching signal levels defines the time change between the OU reverse process and stochastic localization.The appendix gives the relation e^-2(T-r) = s/(1+s).Source: §B.2.2 · p. 25
  • Stochastic localization is the OU reverse diffusion under a time change and spatial rescaling, with equivalent marginal densities along the two paths.Figure 6 illustrates the equivalence for an asymmetric one-dimensional mixture of normals.Source: §B.2.2 · p. 25, Figure 6 · p. 25

B.3 Log-Odds Calculation

The appendix derives exact reward and class log-odds under linear tilts and describes how reward rates are estimated. It also distinguishes two self-conditioning modes, including a limitation of carry for current error-control theory.

B.3 Log-Odds Calculation

  • For a binary reward tilt, the conditional distributions within the rewarded and unrewarded classes remain unchanged, while their class odds shift with tilt strength.The experiments use the training-set fraction of reward-achieving examples as the baseline probability for Ising and Sudoku.Source: §B.3 · p. 26
  • For a one-hot class tilt, the target class gains a factor e^4τ in its odds, while odds among non-target classes remain proportional to their original values.Source: §B.3 · p. 26
  • Reward rates are estimated from sampled outcomes and compared with exact log-odds using 95% Clopper–Pearson intervals.At N = 2000, the largest resolvable log-odds is approximately 7.6; high-tilt discrepancies can correspond to less than one percentage point in reward rate.Source: §B.3 · p. 27
  • Comparing the sampler’s internal reward prediction with exact and externally graded rewards distinguishes an incorrect tilted reward marginal from token–reward inconsistency.Source: §B.3 · p. 27

B.4 Self-Conditioning

  • Two-pass uses a self-estimate in a second network evaluation, whereas carry reuses the previous iteration’s estimate and breaks the memory-one Markov structure assumed by current error-control theory.Source: §B.4 · p. 28

Implementation

Implementation results vary across the Sudoku and MNIST tasks and self-conditioning modes. The appendix also specifies a deterministic, tilt-dependent time grid used for sampling.

Implementation

  • On 6x6 Sudoku, two-pass tracked performance across log-odds levels when confidence intervals were considered.Source: Implementation · p. 29
  • On 9x9 Sudoku, carry tracked log-odds better than two-pass and saturated at 90% reward-achieving solutions, with less mismatch between predicted and actual reward.Two-pass saturated at 60% and showed significant prediction-versus-reward differences.Source: Implementation · p. 29
  • On MNIST, two-pass followed the log-odds curve better, while carry oversampled reward-achieving solutions at medium tilt strengths.The proposed explanation that carry may reinforce early predictions is stated as a conjecture.Source: Implementation · p. 29

B.5 Discrete Step Size Scheduling

  • The sampling time grid is deterministic, computed once, and depends on the tilt through its norm; the experiments use T = 20 and K = 1000.The steps are finest near t = 0 and grow toward T.Source: §B.5 · pp. 30–31

C.1 12-Dimensional Ising Model

The 12-dimensional Ising model serves as a small, exactly enumerable test case for denoiser learning and sampling under reward tilts. The network matches the untilted model’s moments, maintains bounded denoiser error across tested tilts, and tracks reward log-odds.

C.1 12-Dimensional Ising Model

  • The nonlinear reward tilt gives reward when exactly 10 of the 12 bits are positive, and reward information is appended as a 13th coordinate.This converts the reward tilt into a linear tilt on the extended state.Source: §C.1 · p. 31

Extended Results

  • The untilted model matches the target Ising model’s mean and covariance, effectively recreating its first and second moments.The 12-dimensional target is small enough to enumerate the ground-truth denoiser for comparison.Source: §C.1 · p. 32
  • 0.12 is the uniform bound on learned-denoiser L2 error across tested tilt strengths and signal levels, for τ from 0 to 3.6.Error increases with tilt strength; the bound can be converted to a terminal KL-error bound at the final sample point.Source: §C.1 · p. 32
  • At 2,000 samples per τ, the model tracks the reward log-odds perfectly across the tested tilt range.The range extends to the maximum Λ focusing used in training, d + 1.Source: §C.1 · p. 32

C.2 4x4 Sudoku

The 4×4 Sudoku experiment tests tilted sampling where the ground-truth denoiser can be enumerated. The network learns the denoiser well and samples high-validity puzzles, but lags at the strongest tilts and shows a mild preference among solutions.

C.2 4x4 Sudoku

  • The 4×4 state has 32 bits and only 288 valid puzzles, while the independent-row weak generator produces a proper puzzle by chance only 0.081% of the time.The training set therefore mixes in valid solutions; its stated expected injection is 7,500 examples among 150,000 samples.Source: §C.2 · pp. 34–35

Extended Results

  • The learned denoiser’s L2 error increases with tilt strength but remains uniformly bounded by 0.5 across tilt and signal levels.The ground-truth denoiser is enumerable because the reward set contains 288 valid grids and the weak generator factorizes across rows.Source: §C.2 · p. 35
  • The network sampler tracks log-odds up to 98% reward-achieving puzzles, then under-samples at higher tilts even while still producing over 98% valid puzzles.The exact-denoiser sampler tracks the log-odds across all tested levels.Source: §C.2 · p. 36

Early Decoding

  • Confidence-based early decoding records likely terminal bit values and allows partially specified Sudoku grids to be read along the sampling trajectory.A coordinate is fixed when the magnitude of its denoiser prediction crosses a threshold; a fixed bit can still change sign later.Source: §C.2 · p. 37

Extended Results

  • In 3,000 reward-achieving generations, the network visits 287 of 288 solutions but fails the uniformity test, indicating a non-strong preference for some solutions.The exact-denoiser sampler visits all solutions and passes the test.Source: §C.2 · p. 37

C.3 6x6 Sudoku

The 6×6 Sudoku experiment scales the setting to a much larger puzzle space, where the tilted ground-truth denoiser cannot be enumerated. The reported memorization check is consistent with no memorization, while high-tilt log-odds estimates have limited precision.

C.3 6x6 Sudoku

  • For 6×6 Sudoku, each cell uses 3 bits, leaving two illegal codes and yielding a 2^108-bit-string state space; the generator enforces no column or box constraints, so proper-puzzle probability is only 1.8 × 10−10.Source: §C.3 · p. 40

Extended Results

  • With 2000 samples, confidence intervals on high-tilt log-odds become uninformative, although those tilts already produce over 98% correct puzzles.The reward-tilted denoiser is not enumerated for this setting.Source: §C.3 · p. 41
  • 10 of 917 valid generations exactly matched training rows, consistent with no memorization under the stated uniform-sampling null.The training set covers 0.88% of valid grids; the null expectation is 8.1 exact matches, with a 95% interval of 3–14.Source: §C.3 · p. 43

C.4 9x9 Sudoku

The 9×9 Sudoku experiment compares two self-conditioning modes and evaluates infilling and novelty. Carry tracks reward log-odds better than two-pass, while generated valid grids are reported as no closer to training data than held-out solutions.

C.4 9x9 Sudoku

  • The 9×9 state uses 324 core bits, and its 4-bit cell encoding has seven illegal code words per cell.The target space contains approximately 6.67 × 10^21 proper puzzles.Source: §C.4 · p. 43

Extended Results

  • Two-pass self-conditioning plateaus near 60% valid puzzles, while carry tracks log-odds better and plateaus near 90% reward-achieving puzzles.Two-pass predicts over 90% reward-achieving puzzles while only about 60% are actually valid; carry’s prediction-validity discrepancy is about 5–10%, versus 30% for two-pass.Source: §C.4 · pp. 44–45
  • Two-pass and carry produce comparable proportions of puzzles within 10 bits of a valid puzzle, although only 60% of two-pass puzzles are valid.For two-pass, 90% are within 10 bits of a valid puzzle.Source: §C.4 · p. 46, Figure 21 · p. 46

Infilling

  • 80% to 100% of infilling samples both retain the given cells and produce a valid puzzle, depending on how many cells are initially given.With fewer initial cells, validity falls and the sampler does not always preserve the given cells.Source: §C.4 · p. 46, Table 1 · p. 46

Memorization Check

  • The generated and held-out valid-grid distance distributions to training data are almost identical, and no generation exactly matches a reference grid.Among 1,000 generations, 869 are valid; the closest generation differs from every reference grid in 39 cells.Source: §C.4 · pp. 46–47, Figure 23 · p. 48

C.5 MNIST Conditional Image Generation

The MNIST experiments compare two self-conditioning modes for class-conditional generation, measuring class log-odds and generated-image quality. Two-pass tracks the expected tilt more closely, while carry overshoots across the rising tilt range.

C.5 MNIST Conditional Image Generation

  • MNIST images use 6,272 core bits with Gray-code pixels, and each class is tilted through its full 10-bit one-hot reward pattern.One class receives a positive tilt and the other nine negative tilts; this gives class log-odds a slope of 4.Source: §C.5 · p. 48
  • The model was trained for 30,000 Adam steps, and its held-out loss decreased from 720.6 at step 1,000 to 632.1 at step 30,000 with no sign of overfitting.Training used batches of 2,048 and a 3,555,256-parameter network.Source: §C.5 · p. 49
  • Classifier-scored success saturates at the classifier’s class-dependent sensitivity, making tilts at τ ≥1.69 uninformative about further improvements.The authors report that the corrected exact rate is within two standard errors of this limit for every class at those tilts.Source: §C.5 · p. 49
  • Two-pass tracks the corrected exact class log-odds through τ ≈1, while carry overshoots; at τ = 0.51, carry yields 57.8% versus an exact 45.9%.The untilted rate is correct under both modes, isolating the carry discrepancy to its response to tilt.Source: Table 2 · p. 50, Table 3 · p. 50, §C.5 · p. 50
  • Generated digits show similar quality and variety to held-out images, with no observed mode collapse or degradation in digit quality under tilting.The visual comparison includes held-out digits and samples from both self-conditioning modes.Source: §C.5 · p. 54, Figure 28 · p. 54, Figure 29 · p. 55
Loading 2610.29086v1…