Introduction
Diffusion models designed via Gaussian noising have become the dominant generative algorithm for continuous distributions [29 , 13 , 32 ] , and continue to show promise for discrete targets as well [7 , 16 , 35 ] . The original theory of diffusion models was established for unconditional generation [31 , 32 ] : given access to an initial pool of samples from a distribution π \pi π π \pi π π \pi π r r r τ \tau τ π ~ ( x ) ∝ π ( x ) e τ r ( x ) \tilde{\pi}(x) \propto\pi(x)e^{\tau r(x)} π ~ ( x ) ∝ π ( x ) e τ r ( x )
Existing methods utilize what the authors refer to as external guidance , implementing an additive vector field called a guidance term to modify SDE paths at sampling time [10 , 28 , 11 ] , or a density ratio to reweight Markov chain conditional probabilities [33 , 24 , 17 , 12 ] . Methods such as classifier-free guidance [14 ] implement ad hoc and closed-form guidance terms, but are known to structurally sample the wrong invariant measure and introduce sampling artifacts such as corrupted samples and mode collapse [8 , 4 ] . Exact guidance schemes derive the analytically correct guidance terms, but result in an intractable object that must either be learned in an expensive neural network training run, or approximated by a tractable analog such as an empirical average that introduces compounding errors at sample time.
A recent line of work [15 , 21 , 20 , 22 ] has opened up a new method for reward-tilted sampling we call internal guidance . Rather than modify the output of the denoiser of the base distribution with additive or multiplicative terms outside of the object, we realize we can instead shift our evaluation point inside the existing base denoiser. Recent work of Moitra et al. [21, Lemma 3.3] has shown for an internal mapping function H H H
G t i l t ( u , a ) = G b a s e ( u ~ = H 1 ( u , a ) , a ~ = H 2 ( u , a ) ) . G_{\mathrm{tilt}}(u,a)=G_{\mathrm{base}}(\tilde{u}=H_1(u,a),\tilde{a}=H_2(u,a)). G tilt ( u , a ) = G base ( u ~ = H 1 ( u , a ) , a ~ = H 2 ( u , a )) . These symmetry mappings are closed-form objects involving shifting and rescaling inputs, derived from the interactions of linear tilts with the quadratic log-density terms introduced via Gaussian convolution.
Internal guidance has clear benefits. There is no intractable external term to approximate, and the symmetry maps H H H
The catch with this method is one of coverage. If we had access to the oracle denoiser G G G g θ g_\theta g θ G G G L 2 L^2 L 2 ( u ~ , a ~ ) (\tilde{u},\tilde{a}) ( u ~ , a ~ )
The relationship of the symmetry map can be captured by an equivalence class: for a tilted density π ~ \tilde{\pi} π ~ ( u , a ) (u,a) ( u , a ) ( u ~ , a ~ ) (\tilde{u},\tilde{a}) ( u ~ , a ~ ) ( u ~ , a ~ ) (\tilde{u},\tilde{a}) ( u ~ , a ~ ) ( u , a ) (u,a) ( u , a ) [21 ] can be interpreted as showing this equivalence class exists and is non-empty, and established one symmetry map which creates the class. However, this symmetry alone is not enough to force overlap between the equivalence class of points produced by tilted sampling and the high probability region of the base distribution training manifold the denoiser network was trained on.
To have equivalent points that live near the training manifold, we need to introduce additional symmetry relationships that expand the equivalence class. We restrict our study to discrete distributions embedded in the hypercube { − 1 , 1 } d \{-1,1\}^{d} { − 1 , 1 } d ∥ x ∥ 2 = d \lVert x\rVert^2=d ∥ x ∥ 2 = d
Conditional generation tasks such as infilling and reward tilting by binary rewards can be cast as linear tilts in an extended state space by promoting the reward coordinate to an additional random variable noised and denoised when sampling.
Figure 1. Sketch of denoiser symmetry. The training manifold is where the denoiser network has learned the base denoiser. The sampling manifold is the typical vectors that will be encountered when sampling a tilted density. The equivalence relationships can find a point with equal evaluation of the base denoiser near the training manifold, not all equivalence relationships map back to the region of the training manifold.
We establish the theory of denoiser symmetry , the evaluation of the denoiser of a linear tilt to a base density can be recast as an evaluation of the base denoiser at a shifted evaluation point. We establish additional symmetries specific to hypercube embedded discrete distributions, by the removal of quadratic terms in the conditional log-density.
We introduce diagonal matrix rescalings to our hypercube embedding during training called Λ \Lambda Λ focusing , turning an out-of-distribution learning problem into an in-distribution learning problem, and allowing the base denoiser to be evaluated at locations off the typical noising path.
We demonstrate performance on small discrete sampling targets such as a non-linear reward tilt to a small 12 dimension Ising model with ground truth comparison denoiser; generating and solving 4x4, 6x6, and 9x9 Sudoku puzzles; and conditional class generation of MNIST digit images. These problems are small enough to act as a test bed where the quality of our sampling can be compared against exact ground truth values as a diagnostic of effective sampling, with future work to focus on scaling and more complex target distributions for internal guidance.
Setup
We consider discrete distributions embedded in the hypercube X → { − 1 , 1 } d \mathcal{X} \to\{-1,1\}^{d} X → { − 1 , 1 } d [7 ] , referred to as “analog bits”, although the structural results of linear tilts and internal guidance we derive here were not known at the time. Any discrete distribution can be embedded in the hypercube by assigning each element x x x X \mathcal{X} X 2 d 2^d 2 d e ( x ) ∈ { − 1 , 1 } d e(x) \in\{-1,1\}^d e ( x ) ∈ { − 1 , 1 } d
Our target density π \pi π { − 1 , 1 } d \{-1,1\}^d { − 1 , 1 } d π \pi π v ∈ R d v \in\mathbb{R}^d v ∈ R d π \pi π π v ( x ) ∝ π ( x ) e ⟨ v , x ⟩ \pi_v(x) \propto\pi(x)e^{\langle v,x\rangle} π v ( x ) ∝ π ( x ) e ⟨ v , x ⟩ r r r { − 1 , 1 } \{-1,1\} { − 1 , 1 } π \pi π τ > 0 \tau> 0 τ > 0 r r r q τ ( x ) ∝ π ( x ) e τ r ( x ) q_\tau(x) \propto\pi(x)e^{\tau r(x)} q τ ( x ) ∝ π ( x ) e τ r ( x )
To unify these two objectives, we promote the reward value r ( x ) r(x) r ( x ) x ∗ = r ( x ) x_* = r(x) x ∗ = r ( x ) d + 1 d+1 d + 1 x i x_i x i { − 1 , 1 } \{-1,1\} { − 1 , 1 } x ∗ x_* x ∗ { − 1 , 1 } \{-1,1\} { − 1 , 1 } π ( x ∗ ∣ x ) = 1 { x ∗ = r ( x ) } \pi(x_*|x) = 1\{x_* = r(x)\} π ( x ∗ ∣ x ) = 1 { x ∗ = r ( x )} extended base density with d + 1 d+1 d + 1 π \pi π e ∗ = ( 0 , ⋯ , 0 , 1 ) ∈ R d + 1 e_* = (0,\cdots,0,1) \in\mathbb{R}^{d+1} e ∗ = ( 0 , ⋯ , 0 , 1 ) ∈ R d + 1
π ( x , x ∗ ) = π ( x ) 1 { r ( x ) = x ∗ } \pi(x,x_*) = \pi(x)1\{r(x) = x_*\} π ( x , x ∗ ) = π ( x ) 1 { r ( x ) = x ∗ } π ( x ) e τ r ( x ) = π ( x , x ∗ ) e τ ⟨ e ∗ , ( x , x ∗ ) ⟩ . \pi(x)e^{\tau r(x)} = \pi(x,x_*)e^{\tau\langle e_*,(x,x_*)\rangle}. π ( x ) e τ r ( x ) = π ( x , x ∗ ) e τ ⟨ e ∗ , ( x , x ∗ )⟩ . As such, linear tilts by vectors v v v π \pi π d + 1 d+1 d + 1
For values a , σ > 0 a,\sigma> 0 a , σ > 0 π \pi π
U = a X + σ Z , X ∼ π , Z ∼ N ( 0 , I ) . U = aX + \sigma Z,\ X \sim\pi,\ Z \sim N(0,I). U = a X + σ Z , X ∼ π , Z ∼ N ( 0 , I ) . The denoiser of π \pi π X X X U U U
F b a s e ( u , a , σ ) = E π [ X ∣ u , a , σ ] = ∑ x π ( x ) e − 1 2 σ 2 ∥ u − a x ∥ 2 E π [ e − 1 2 σ 2 ∥ u − a X ∥ 2 ] . F^{\mathrm{base}}(u,a,\sigma) = E_\pi[X|u,a,\sigma] = \sum_x \frac{\pi(x)e^{-\frac{1}{2\sigma^2}\lVert u-ax\rVert^2}}{E_\pi[e^{-\frac{1}{2\sigma^2}\lVert u-aX\rVert^2}]}. F base ( u , a , σ ) = E π [ X ∣ u , a , σ ] = x ∑ E π [ e − 2 σ 2 1 ∥ u − a X ∥ 2 ] π ( x ) e − 2 σ 2 1 ∥ u − a x ∥ 2 . The denoiser defines the drift which implements the reverse SDE sampling algorithm. A special case is the variance preserving (VP) combination, when a 2 + σ 2 = 1 a^2+\sigma^2=1 a 2 + σ 2 = 1 σ = 1 − a 2 \sigma=\sqrt{1-a^2} σ = 1 − a 2 F b a s e ( u , a ) F^{\mathrm{base}}(u,a) F base ( u , a )
In the extended state space, we are really considering the joint noising of the token bit string and the reward coordinate at the same time, u = ( u t o k e n , u ∗ ) u=(u^{\mathrm{token}},u_*) u = ( u token , u ∗ )
( U t o k e n U ∗ ) = a ( x t o k e n x ∗ ) + 1 − a 2 ( Z Z ∗ ) \begin{pmatrix} U^{\mathrm{token}} \\ U_* \end{pmatrix}
= a\begin{pmatrix} x^{\mathrm{token}} \\ x_* \end{pmatrix}
+ \sqrt{1-a^2}\begin{pmatrix} Z \\ Z_* \end{pmatrix} ( U token U ∗ ) = a ( x token x ∗ ) + 1 − a 2 ( Z Z ∗ ) a ∈ [ 0 , 1 ] , X , X ∗ ∼ π , Z ∼ N ( 0 , I ) , Z ∗ ∼ N ( 0 , 1 ) a \in[0,1], X,X_* \sim\pi, Z \sim N(0,I), Z_* \sim N(0,1) a ∈ [ 0 , 1 ] , X , X ∗ ∼ π , Z ∼ N ( 0 , I ) , Z ∗ ∼ N ( 0 , 1 )
Thus our sampler carries a “noisy reward” value in its d + 1 d+1 d + 1
Denoisers are expectations of quadratic tilts to the base density. By restricting our base density to be supported on the hypercube, all quadratic terms can be dropped from the exponent of the conditional densities that define the denoiser expectation. As such, all denoiser evaluations of π \pi π π \pi π π \pi π π \pi π π \pi π
Definition 1. For any vector v ∈ R d + 1 v \in\mathbb{R}^{d+1} v ∈ R d + 1 a ∈ [ 0 , 1 ] a \in[0,1] a ∈ [ 0 , 1 ] β ( a ) = a 1 − a 2 \beta(a) = \frac{a}{1-a^2} β ( a ) = 1 − a 2 a
G ( v , a ) = ∑ x x π ( x ) e β ( a ) ⟨ v , x ⟩ ∑ x ˉ π ( x ˉ ) e β ( a ) ⟨ v , x ˉ ⟩ . G(v,a) = \sum_x x \frac{\pi(x)e^{\beta(a)\langle v,x\rangle}}{\sum_{\bar{x}}\pi(\bar{x})e^{\beta(a)\langle v,\bar{x}\rangle}}. G ( v , a ) = x ∑ x ∑ x ˉ π ( x ˉ ) e β ( a ) ⟨ v , x ˉ ⟩ π ( x ) e β ( a ) ⟨ v , x ⟩ . The base denoiser F b a s e F^{\mathrm{base}} F base G G G
Sampling Via Stochastic Localization
Our sampling algorithm will be the stochastic localization path, which is equivalent to a time shifting and rescaling of the marginal densities of score-based diffusion. See Montanari [23 ] and our discussion in Appendix § 9 .§ 2 for further connection. The benefit of stochastic localization is the drift term is exactly the denoiser, rather than being related to the score via Tweedie’s formula [25 ] . This makes it simpler to construct denoiser-to-denoiser relationships rather than score-to-score relationships.
Definition 2 (Stochastic Localization). For a target density π \pi π X 0 = 0 X_0 = 0 X 0 = 0
d X t = F π ( X t , t , t ) d t + d B t , dX_t = F^\pi(X_t,t,\sqrt{t})dt + dB_t, d X t = F π ( X t , t , t ) d t + d B t , then the marginals at time t t t
X t = t X + t Z , X ∼ π , Z ∼ N ( 0 , I ) . X_t = tX + \sqrt{t}Z, X \sim\pi, Z \sim N(0,I). X t = tX + t Z , X ∼ π , Z ∼ N ( 0 , I ) . For large T T T X T / T X_T/T X T / T 1 / T 1/\sqrt{T} 1/ T
Denoiser Symmetries
First, we consider properties of the training manifold distribution. Typical points produced under VP noising of the base distribution are norm d + 1 \sqrt{d+1} d + 1
Lemma 1 (Constant Norm Shell). Let π \pi π { − 1 , 1 } d + 1 \{-1,1\}^{d+1} { − 1 , 1 } d + 1 U = a X + 1 − a 2 Z U = aX+\sqrt{1-a^2}Z U = a X + 1 − a 2 Z X ∼ π X \sim\pi X ∼ π Z ∼ N ( 0 , I ) Z \sim N(0,I) Z ∼ N ( 0 , I ) a ∈ [ 0 , 1 ] a \in[0,1] a ∈ [ 0 , 1 ] a a a
E ∥ U ∥ 2 = d + 1 E\|U\|^2 = d+1 E ∥ U ∥ 2 = d + 1 Var ∥ U ∥ 2 = 2 ( 1 − a 4 ) ( d + 1 ) \operatorname{Var}\|U\|^2 = 2(1-a^4)(d+1) Var ∥ U ∥ 2 = 2 ( 1 − a 4 ) ( d + 1 ) U U U d + 1 \sqrt{d+1} d + 1 ϵ > 0 \epsilon> 0 ϵ > 0 P ( ∣ ∥ U ∥ 2 d + 1 − 1 ∣ > ϵ ) ≤ 2 ( 1 − a 4 ) ϵ 2 ( d + 1 ) P\left(\left|\frac{\|U\|^2}{d+1}-1\right|>\epsilon\right) \le\frac{2(1-a^4)}{\epsilon^2(d+1)} P ( d + 1 ∥ U ∥ 2 − 1 > ϵ ) ≤ ϵ 2 ( d + 1 ) 2 ( 1 − a 4 )
The magnitude vector ( ∣ U 1 ∣ , … , ∣ U d + 1 ∣ ) (|U_1|,\ldots,|U_{d+1}|) ( ∣ U 1 ∣ , … , ∣ U d + 1 ∣ ) ∣ N ( a , 1 − a 2 ) ∣ |N(a,1-a^2)| ∣ N ( a , 1 − a 2 ) ∣ π \pi π
We then consider the points produced under a stochastic localization path of a tilted target density. The non-VP scaling of localization paired with a shift term from the linear tilt results in evaluation points that can be of various norm and have one large dominant coordinate in magnitude.
Lemma 2 (Stochastic Localization Tilt). Let X t X_t X t t t t π \pi π τ v \tau v τ v φ = − 1 + 5 2 ≈ 0.618 \varphi= \frac{-1+\sqrt{5}}{2} \approx0.618 φ = 2 − 1 + 5 ≈ 0.618 φ = 1 − φ 2 ⟹ φ = β − 1 ( 1 ) \varphi= 1-\varphi^2 \Longrightarrow\varphi= \beta^{-1}(1) φ = 1 − φ 2 ⟹ φ = β − 1 ( 1 ) φ \varphi φ τ v \tau v τ v
E π v [ X ∣ X t = t X + t Z ] = G ( τ v + X t , φ ) \mathbb{E}_{\pi_v}[X\mid X_t=tX+\sqrt{t}Z]=G(\tau v+X_t,\varphi) E π v [ X ∣ X t = tX + t Z ] = G ( τ v + X t , φ ) = ∑ x x π ( x ) e ⟨ τ v + X t , x ⟩ E π [ e ⟨ τ v + X t , x ⟩ ] . =\sum_x x\frac{\pi(x)e^{\langle\tau v+X_t,x\rangle}}{\mathbb{E}_{\pi}[e^{\langle\tau v+X_t,x\rangle}]}. = x ∑ x E π [ e ⟨ τ v + X t , x ⟩ ] π ( x ) e ⟨ τ v + X t , x ⟩ . In particular, for a reward tilt v = e ⋆ v=e_\star v = e ⋆ d + 1 d+1 d + 1
Based on these two lemmas alone, we see the typical evaluation points for G G G
First, any G G G d + 1 \sqrt{d+1} d + 1 a ~ \tilde{a} a ~
Lemma 3 (Renormalization). For any vector v v v a a a d + 1 \sqrt{d+1} d + 1 v v v a ~ \tilde{a} a ~
G ( v , a ) = G ( d + 1 v ∥ v ∥ , a ~ ) G(v,a)=G\left(\sqrt{d+1}\frac{v}{\lVert v\rVert},\tilde{a}\right) G ( v , a ) = G ( d + 1 ∥ v ∥ v , a ~ ) C = β ( a ) ∥ v ∥ d + 1 , a ~ = 1 2 ( 1 C 2 + 4 − 1 C ) . C=\beta(a)\frac{\lVert v\rVert}{\sqrt{d+1}},\qquad\tilde{a}=\frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right). C = β ( a ) d + 1 ∥ v ∥ , a ~ = 2 1 ( C 2 1 + 4 − C 1 ) . Second, for the base density we can multiply our vectors by any diagonal matrix Λ \Lambda Λ Λ \Lambda Λ
Lemma 4 (Λ \Lambda Λ Let Λ \Lambda Λ U U U X ∼ π X\sim\pi X ∼ π
U = a Λ X + 1 − a 2 Z , X ∼ π , Z ∼ N ( 0 , I ) , a ∈ [ 0 , 1 ] . U=a\Lambda X+\sqrt{1-a^2}Z,\quad X\sim\pi,\ Z\sim N(0,I),\ a\in[0,1]. U = a Λ X + 1 − a 2 Z , X ∼ π , Z ∼ N ( 0 , I ) , a ∈ [ 0 , 1 ] . Then the denoiser is an evaluation of G G G Λ \Lambda Λ u u u
E [ X ∣ U = a Λ X + 1 − a 2 Z ] = G ( Λ u , a ) . \mathbb{E}[X\mid U=a\Lambda X+\sqrt{1-a^2}Z]=G(\Lambda u,a). E [ X ∣ U = a Λ X + 1 − a 2 Z ] = G ( Λ u , a ) . We construct our diagonal matrix Λ \Lambda Λ tr ( Λ 2 ) = d + 1 \operatorname{tr}(\Lambda^2)=d+1 tr ( Λ 2 ) = d + 1 d \sqrt{d} d q ∈ R d q\in\mathbb{R}^d q ∈ R d λ ∈ [ 1 , d + 1 ] \lambda\in[1,\sqrt{d+1}] λ ∈ [ 1 , d + 1 ]
Λ j , j = 1 + 1 d − λ 2 d q j , j ≠ d + 1 , Λ d + 1 , d + 1 = λ . \Lambda_{j,j}=\sqrt{1+\frac{1}{d}-\frac{\lambda^2}{d}q_j},\ j\ne d+1,\qquad\Lambda_{d+1,d+1}=\lambda. Λ j , j = 1 + d 1 − d λ 2 q j , j = d + 1 , Λ d + 1 , d + 1 = λ . This effectively moves our embedding points x x x { − 1 , 1 } d + 1 \{-1,1\}^{d+1} { − 1 , 1 } d + 1 e d + 1 e_{d+1} e d + 1 Figure 2 . We consider two choices of q q q
Figure 2. Reward based Λ \Lambda Λ
Mode 0: Reward Focusing Here we set q = 1 ⃗ d q=\vec{1}_d q = 1 d Λ \Lambda Λ
Mode 1: Infill Training Let A + A_+ A + A − A_- A − π \pi π
π ( x ∣ x A + = 1 , x A − = − 1 ) = lim τ → ∞ π ( x ) e τ ( ⟨ e A + , x ⟩ − ⟨ e A − , x ⟩ ) . \pi(x\mid x_{A_+}=1,x_{A_-}=-1)=\lim_{\tau\to\infty}\pi(x)e^{\tau(\langle e_{A_+},x\rangle-\langle e_{A_-},x\rangle)}. π ( x ∣ x A + = 1 , x A − = − 1 ) = τ → ∞ lim π ( x ) e τ (⟨ e A + , x ⟩ − ⟨ e A − , x ⟩) . As such, pick a set of indexes A A A ∣ A ∣ = k \lvert A\rvert=k ∣ A ∣ = k q q q d d d A A A q j = 1 { j ∈ A } d / k q_j=\mathbf{1}\{j\in A\}\sqrt{d/k} q j = 1 { j ∈ A } d / k A C A^C A C
E [ X ∣ U A , U = a X + 1 − a 2 Z ] = G ( Diag ( 1 A ) u , a ) . \mathbb{E}[X\mid U_A,U=aX+\sqrt{1-a^2}Z]=G(\operatorname{Diag}(\mathbf{1}_A)u,a). E [ X ∣ U A , U = a X + 1 − a 2 Z ] = G ( Diag ( 1 A ) u , a ) . We refer to this concept as “soft masking”, similar in spirit to masked diffusion models [2 , 26 , 27 ] where certain tokens are set to a ‘MASK’ token and do not influence the conditional expectation. Here, we can get a similar effect by shutting off certain indexes with a q j = 0 q_j=0 q j = 0 U U U U A U_A U A
Combined Results
Combining these symmetries of shifting, normalizing, and Λ \Lambda Λ Λ \Lambda Λ d + 1 \sqrt{d+1} d + 1
Corollary 1. Let U U U Λ \Lambda Λ
U = a Λ X + 1 − a 2 Z , X ∼ π , Z ∼ N ( 0 , I ) , a ∈ [ 0 , 1 ] , U = a\Lambda X + \sqrt{1-a^2}Z,\ X \sim\pi,\ Z \sim N(0,I),\ a \in[0,1], U = a Λ X + 1 − a 2 Z , X ∼ π , Z ∼ N ( 0 , I ) , a ∈ [ 0 , 1 ] , then the denoiser at this observation point is equal to an evaluation of the base denoiser at a norm d + 1 \sqrt{d+1} d + 1 a ~ \tilde{a} a ~
E π [ X ∣ U = a Λ X + 1 − a 2 Z ] = G ( d + 1 Λ u ∥ Λ u ∥ , a ~ ) E_\pi[X|U = a\Lambda X + \sqrt{1-a^2}Z] = G\left(\sqrt{d+1}\frac{\Lambda u}{\|\Lambda u\|},\tilde{a}\right) E π [ X ∣ U = a Λ X + 1 − a 2 Z ] = G ( d + 1 ∥Λ u ∥ Λ u , a ~ ) C = a 1 − a 2 ∥ Λ u ∥ d + 1 , a ~ = 1 2 ( 1 C 2 + 4 − 1 C ) . C = \frac{a}{1-a^2}\frac{\|\Lambda u\|}{\sqrt{d+1}},\quad\tilde{a} = \frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right). C = 1 − a 2 a d + 1 ∥Λ u ∥ , a ~ = 2 1 ( C 2 1 + 4 − C 1 ) . Corollary 2. Let X t X_t X t t t t π \pi π τ v \tau v τ v d + 1 \sqrt{d+1} d + 1 a ~ \tilde{a} a ~
E π v [ X ∣ X t = t X + t Z ] = G ( d + 1 τ v + X t ∥ τ v + X t ∥ , a ~ ) E_{\pi_v}[X|X_t = tX + \sqrt{t}Z] = G\left(\sqrt{d+1}\frac{\tau v + X_t}{\|\tau v + X_t\|},\tilde{a}\right) E π v [ X ∣ X t = tX + t Z ] = G ( d + 1 ∥ τ v + X t ∥ τ v + X t , a ~ ) C = ∥ τ v + X t ∥ d + 1 , a ~ = 1 2 ( 1 C 2 + 4 − 1 C ) . C = \frac{\|\tau v + X_t\|}{\sqrt{d+1}},\quad\tilde{a} = \frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right). C = d + 1 ∥ τ v + X t ∥ , a ~ = 2 1 ( C 2 1 + 4 − C 1 ) . At sampling time, we encounter evaluations of the G G G a ~ \tilde{a} a ~ d + 1 \sqrt{d+1} d + 1 d + 1 d+1 d + 1 Λ \Lambda Λ e ∗ , e_*, e ∗ ,
Training and Sampling Algorithms
In training we may noise the data under any joint distribution of signal level a a a Λ \Lambda Λ d + 1 \sqrt{d+1} d + 1 a ~ \tilde{a} a ~ G G G Λ = I \Lambda= I Λ = I
Lemma 5 (General regression minimizer). Let μ \mu μ ( a , Λ ) (a,\Lambda) ( a , Λ ) a ∈ [ 0 , 1 ] a \in[0,1] a ∈ [ 0 , 1 ] Λ \Lambda Λ ( a , Λ ) ∼ μ (a,\Lambda) \sim\mu ( a , Λ ) ∼ μ X ∼ π X \sim\pi X ∼ π Z ∼ N ( 0 , I ) Z \sim N(0,I) Z ∼ N ( 0 , I ) U = a Λ X + 1 − a 2 Z U = a\Lambda X + \sqrt{1-a^2}Z U = a Λ X + 1 − a 2 Z H : R d + 1 × [ 0 , 1 ] → S × [ 0 , 1 ] H : \mathbb{R}^{d+1} \times[0,1] \to\mathcal{S} \times[0,1] H : R d + 1 × [ 0 , 1 ] → S × [ 0 , 1 ]
H ( v , a ) = ( d + 1 v ∥ v ∥ , a ~ ) H(v,a) = \left(\sqrt{d+1}\frac{v}{\|v\|},\tilde{a}\right) H ( v , a ) = ( d + 1 ∥ v ∥ v , a ~ ) C = β ( a ) ∥ v ∥ d + 1 , a ~ = 1 2 ( 1 C 2 + 4 − 1 C ) . C = \beta(a)\frac{\|v\|}{\sqrt{d+1}},\quad\tilde{a} = \frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right). C = β ( a ) d + 1 ∥ v ∥ , a ~ = 2 1 ( C 2 1 + 4 − C 1 ) . Then over all measurable functions g g g G G G
G = argmin g E ∥ X − g ( H ( Λ U , a ) ) ∥ 2 . G = \underset{g}{\operatorname{argmin}}\ E\|X-g(H(\Lambda U,a))\|^2. G = g argmin E ∥ X − g ( H ( Λ U , a )) ∥ 2 . We minimize the empirical loss by SGD in Algorithm 1, drawing a fresh ( a , Λ ) (a,\Lambda) ( a , Λ ) β \beta β [ β min , β max ] = [ 0.02 , 50 ] [\beta_{\min},\beta_{\max}] = [0.02,50] [ β m i n , β m a x ] = [ 0.02 , 50 ] a = a ( β ) a=a(\beta) a = a ( β ) β ( a ) = a / ( 1 − a 2 ) \beta(a)=a/(1-a^2) β ( a ) = a / ( 1 − a 2 ) p 0 p_0 p 0 q = 1 d q=1_d q = 1 d k k k p λ p_\lambda p λ λ = 1 \lambda=1 λ = 1 λ \lambda λ [ 1 , d + 1 ] [1,\sqrt{d+1}] [ 1 , d + 1 ] q = 1 d q=1_d q = 1 d λ = 1 \lambda=1 λ = 1 p 0 = 3 / 4 p_0=3/4 p 0 = 3/4 p λ = 1 / 4 p_\lambda=1/4 p λ = 1/4 3 / 16 3/16 3/16
Algorithm 1 Batch SGD step for the tilt–mean network g θ g_\theta g θ
Require: Data { X i } ⊂ { − 1 , 1 } d + 1 \{X_i\}\subset\{-1,1\}^{d+1} { X i } ⊂ { − 1 , 1 } d + 1 d + 1 d+1 d + 1 N N N γ \gamma γ
L ← 0 \mathcal{L}\leftarrow0 L ← 0
for n = 1 , … , N n=1,\ldots,N n = 1 , … , N do
Draw x ∼ { X i } x\sim\{X_i\} x ∼ { X i } Z ∼ N ( 0 , I d + 1 ) Z\sim\mathcal{N}(0,I_{d+1}) Z ∼ N ( 0 , I d + 1 ) β ∼ LogUnif [ β min , β max ] \beta\sim\operatorname{LogUnif}[\beta_{\min},\beta_{\max}] β ∼ LogUnif [ β m i n , β m a x ] a ← a ( β ) a\leftarrow a(\beta) a ← a ( β )
W.p. p 0 p_0 p 0 q ← 1 d q\leftarrow1_d q ← 1 d k ∼ Unif { 2 , … , d } k\sim\operatorname{Unif}\{2,\ldots,d\} k ∼ Unif { 2 , … , d } q ← d / k 1 S q\leftarrow\sqrt{d/k}1_S q ← d / k 1 S k k k S S S
W.p. p λ p_\lambda p λ λ ← 1 \lambda\leftarrow1 λ ← 1 λ ∼ LogUnif [ 1 , d + 1 ] \lambda\sim\operatorname{LogUnif}[1,\sqrt{d+1}] λ ∼ LogUnif [ 1 , d + 1 ]
Λ j , j ← 1 + 1 d − λ 2 d q j 2 \Lambda_{j,j}\leftarrow\sqrt{1+\frac{1}{d}-\frac{\lambda^2}{d}q_j^2} Λ j , j ← 1 + d 1 − d λ 2 q j 2 j ≤ d j\le d j ≤ d Λ d + 1 , d + 1 ← λ \Lambda_{d+1,d+1}\leftarrow\lambda Λ d + 1 , d + 1 ← λ
η ← ( a Λ 2 x + 1 − a 2 Λ Z ) \eta\leftarrow(a\Lambda^2x+\sqrt{1-a^2}\Lambda Z) η ← ( a Λ 2 x + 1 − a 2 Λ Z )
if q = 1 d q=1_d q = 1 d λ = 1 \lambda=1 λ = 1 then
( v ∗ , s ∗ ) ← ( η , a ) (v^*,s^*)\leftarrow(\eta,a) ( v ∗ , s ∗ ) ← ( η , a )
else
C ← β ( a ) ∥ η ∥ d + 1 C\leftarrow\beta(a)\frac{\|\eta\|}{\sqrt{d+1}} C ← β ( a ) d + 1 ∥ η ∥ a ~ ← 1 2 ( 1 C 2 + 4 − 1 C ) \tilde a\leftarrow\frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right) a ~ ← 2 1 ( C 2 1 + 4 − C 1 )
v ∗ ← d + 1 η / ∥ η ∥ v^*\leftarrow\sqrt{d+1}\eta/\|\eta\| v ∗ ← d + 1 η /∥ η ∥
L ← L + ∥ x − g θ ( v ∗ , s ∗ ) ∥ 2 \mathcal{L}\leftarrow\mathcal{L}+\|x-g_\theta(v^*,s^*)\|^2 L ← L + ∥ x − g θ ( v ∗ , s ∗ ) ∥ 2
θ ← θ − γ ∇ θ L \theta\leftarrow\theta-\gamma\nabla_\theta\mathcal{L} θ ← θ − γ ∇ θ L
We note our sampler also makes use of self-conditioning [7 ] which is a common practice in the field. We discuss the form of self-conditioning in Appendix § 9 .§ 4 to avoid over complication of the discussion in the main body.
Sampling a linear tilt follows Corollary 2: at each step the state is shifted by the tilt vector and renormalized before querying the network, shown in Algorithm 2. The terminal X T X_T X T
Algorithm 2 Stochastic-localization sampling from the tilted density μ ∝ π e ⟨ ν , ⋅ ⟩ \mu\propto\pi e^{\langle\nu,\cdot\rangle} μ ∝ π e ⟨ ν , ⋅ ⟩
Require: Trained network g θ g_\theta g θ ν = τ i n f i l l ( e A + − e A − ) + τ r e w a r d e ∗ \nu=\tau_{\mathrm{infill}}(e_{A_+}-e_{A_-})+\tau_{\mathrm{reward}}e_* ν = τ infill ( e A + − e A − ) + τ reward e ∗ 0 = t 0 < t 1 < ⋯ < t K = T 0=t_0<t_1<\cdots<t_K=T 0 = t 0 < t 1 < ⋯ < t K = T
X ˙ t 0 ← 0 \dot X_{t_0}\leftarrow0 X ˙ t 0 ← 0
for k = 0 , … , K − 1 k=0,\ldots,K-1 k = 0 , … , K − 1 do
δ ← t k + 1 − t k \delta\leftarrow t_{k+1}-t_k δ ← t k + 1 − t k
η ← ν + X t k \eta\leftarrow\nu+X_{t_k} η ← ν + X t k
C ← ∥ η ∥ d + 1 C\leftarrow\frac{\|\eta\|}{\sqrt{d+1}} C ← d + 1 ∥ η ∥ s ∗ ← 1 2 ( 1 C 2 + 4 − 1 C ) s^*\leftarrow\frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right) s ∗ ← 2 1 ( C 2 1 + 4 − C 1 )
v ∗ ← d + 1 η / ∥ η ∥ v^*\leftarrow\sqrt{d+1}\eta/\|\eta\| v ∗ ← d + 1 η /∥ η ∥
Draw ξ ∼ N ( 0 , I d + 1 ) \xi\sim\mathcal{N}(0,I_{d+1}) ξ ∼ N ( 0 , I d + 1 )
X t k + 1 ← X t k + δ g θ ( v ∗ , s ∗ ) + δ ξ X_{t_{k+1}}\leftarrow X_{t_k}+\delta g_\theta(v^*,s^*)+\sqrt{\delta}\xi X t k + 1 ← X t k + δ g θ ( v ∗ , s ∗ ) + δ ξ
η ← ν + X T \eta\leftarrow\nu+X_T η ← ν + X T ( v ∗ , s ∗ ) (v^*,s^*) ( v ∗ , s ∗ ) η \eta η
return x ^ ← sign ( g θ ( v ∗ , s ∗ ) ) \hat{x}\leftarrow\operatorname{sign}(g_\theta(v^*,s^*)) x ^ ← sign ( g θ ( v ∗ , s ∗ ))
Experiments
We demonstrate the performance of the method in three experimental settings: a 12-dimensional Ising model with a non-linear reward tilt, 4x4, 6x6, and 9x9 Sudoku unconditional grid generation and puzzle solving (infilling), and MNIST conditional digit generation. Full technical descriptions of the target density, training set, model architecture, and additional details for each experiment can be found in Appendix § 10 . All experiments are sampled up to a stochastic localization time T = 20 T = 20 T = 20 { t k } k = 0 K \{t_k\}_{k=0}^{K} { t k } k = 0 K K = 1000 K = 1000 K = 1000 τ \tau τ
12-Dimensional Ising Model
The 12-dimensional Ising model has a reward tilt which gives reward when the number of positive bits is exactly 10. The state space of the Ising model is small enough to enumerate giving exact ground truth for the denoiser to compare our network against. In Appendix § 10 Figure 8 we see the network denoiser has learned the ground truth denoiser well with low L 2 L^2 L 2 § 10 Figure 9 . The Ising model is a small scale experiment with ground truth denoiser to calibrate the method before proceeding to larger state spaces.
Figure 8. Comparison of L 2 L^2 L 2
Figure 9. Ising model log-odds across tilt strengths.
4x4 Sudoku
The 4x4 Sudoku problem has a small enough state space to enumerate the ground truth denoiser. We compare the network denoiser to the ground truth in L 2 L^2 L 2 § 10 Figure 11 , the network learns the denoiser well. For the log-odds curve in Appendix § 10 Figure 12 we can consider the stochastic localization algorithm implemented with either the ground truth denoiser or the network estimation. Both methods track the log-odds up to 98% reward-achieving solutions at τ ≈ 3.5 \tau\approx3.5 τ ≈ 3.5 τ = 5 \tau= 5 τ = 5
Figure 11. Comparison of L 2 L^2 L 2 4 × 4 4 \times4 4 × 4
Figure 12. 4 × 4 4 \times4 4 × 4
With 288 reward-achieving puzzles, we can track if the network is uniformly sampling these puzzles. In 3,000 generated reward-achieving puzzles, the exact-denoiser sampler visits all 288 solutions, each between 4 and 21 times, and passes a χ 2 \chi^2 χ 2 χ 287 2 = 280.1 \chi^2_{287} = 280.1 χ 287 2 = 280.1 p = 0.60 p = 0.60 p = 0.60 χ 287 2 = 391.9 \chi^2_{287} = 391.9 χ 287 2 = 391.9 p < 10 − 4 p < 10^{-4} p < 1 0 − 4 § 10 Figure 15 .
(a) 4x4 Sudoku chi-squared test network samples.
(b) 4x4 Sudoku chi-squared test ground truth samples.
Figure 15. Figure 15
An added benefit of embedding in the hypercube is we can do early bit-by-bit decoding of the continuous vector X t X_t X t t t t ∣ E [ X i ∣ X t ] ∣ ≥ 0.75 |E[X_i|X_t]| \ge0.75 ∣ E [ X i ∣ X t ] ∣ ≥ 0.75 § 10 Figure 13 .
Figure 13. Sudoku 4x4 denoiser trajectory and early decoding. Partially complete discrete token sequences fixed when denoiser coordinates pass magnitude threshold.
6x6 Sudoku
For 6x6 Sudoku, the state space is large and there are vastly more reward-achieving puzzles than the 4x4 problem. In the 4x4 problem, the state space was size 2 32 2^{32} 2 32 χ 2 \chi^2 χ 2
For 6x6 Sudoku the state space is 2 108 2^{108} 2 108
We cannot fully enumerate the denoiser for reward tilts in the 6x6 setting, but we can show the L 2 L^2 L 2 τ = 0 \tau= 0 τ = 0 § 10 Figure 17. The network also properly tracks the log-odds across tilt strengths in Appendix § 10 Figure 18.
We check memorization by the rate the network reproduces a reward-achieving puzzle provided in the training set. With the training set containing 0.88% of the reward-achieving puzzles, in 1,000 reward-achieving generations proper uniform sampling would produce exact matches with the training set by chance 3–14 with a 95% interval. We observe 10 matches in 1,000 samples, consistent with no memorization.
9x9 Sudoku
For 9x9 Sudoku the state space is too large to enumerate a ground truth denoiser. The number of reward-achieving puzzles is so large there is a near 0 probability of recreating a training set puzzle under proper sampling. The training set is 10 million Sudoku grids with 5% being reward-achieving puzzles.
The 9x9 Sudoku target density is the first example where the effect of self-conditioning is noticeable, self-conditioning is discussed in more detail in Appendix B.4. In summary, we consider two modes of self-conditioning: mode “two-pass” reuses the network’s denoiser evaluation at a given time and maintains a memory one Markov structure. Mode “carry” reuses a previous time point’s denoiser value and breaks this Markov structure but is more commonly used in modern methods.
The two methods give different results when tracking the log-odds in Appendix § 10 Figures 20a and 20b. The reward coordinate is denoised jointly with the tokens, so the sampler reports its own prediction of reward in the d + 1 d + 1 d + 1 d + 1 d + 1 d + 1 d + 1 d + 1 d + 1 § 10 Figure 21), consistent with a few bits mis-sampled late in the trajectory being penalized by the sharp validity reward.
The model performs both unconditional grid generation and infilling, which some may call “solving” a puzzle, by adding a cell-fixing tilt to the reward tilt (Figure 3 ). This type of sampling is natural for the diffusion model here which operates under linear tilts and decodes in any order, but would be difficult for a left-to-right autoregressive model which would need to backtrack against later given values. We use a reward tilt strength and infill tilt strength of τ = 3 \tau= 3 τ = 3 § 10 Table 1 and Figure 4 . Note the sampler may not maintain the initial given cell values as they are a tilt vector not a hard constraint. Given 71/81 cells the model maintains the initial cells and achieved reward in 100% of 2000 draws. Given 11/81 cells, the initial cell values are maintained 92% of the time, a reward achieved 87% of the time, and both objectives satisfied 80% of the time.
(a) Unconditional generation from a blank grid.
(b) Conditional generation from a partially given puzzle.
Figure 3. Sudoku 9x9 completions sampled from the tilt-operator network.
Figure 4. Probability of correct grid infilling across input cell numbers. Given no initial cells, should return a valid grid 95% of the time, we measure saturation at 87% without any infilling.
Finally, the model has no indication of memorizing the training set. Of 1000 unconditional draws at τ = 3 \tau= 3 τ = 3
MNIST Digit Images
The MNIST data set consists of 60,000 images of hand drawn digits 0-9 and 10,000 hold out examples. We encode reward as a 10-bit one-hot class block, and generating class tilts along the whole block with 1 class being positively tilted and the other 9 negatively tilted. Generated images are classified by an independent CNN classifier (98.74% held-out accuracy). Since the scorer has a ceiling to its classifier accuracy, the observable success rate saturates at the classifier’s per-class sensitivity; points beyond τ ≈ 1.69 \tau\approx1.69 τ ≈ 1.69
In Appendix § 10 Figure 29 we see generated pools of images show similar variety and quality as hold out examples not used for training indicating no mode collapse or degradation of digit quality under tilting.
In Appendix § 10 Figures 24-27 self-conditioning mode two-pass and carry are also compared for their tracking of the log-odds. In this experiment, carry mode seems to over-sample a class at medium tilts. Two-pass tracks the classifier-corrected exact rate within its confidence interval through τ ≈ 1 \tau\approx1 τ ≈ 1 τ = 0.51 \tau= 0.51 τ = 0.51 τ \tau τ
Acknowledgments
We gratefully acknowledge the support of the NSF for FODSI through grant DMS-2023505, of the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031883 and 814639, and of the ONR through MURI award N000142112431, and a research grant from Bridgewater to the Simons Institute.
AI Use Statement
In this work, we used generative AI tools to generate synthetic data sets, help develop theoretical models or conceptual frameworks, formulate mathematical claims, assist in the writing of proofs, design or provide feedback on research methodology or experiments, implement methods, clean and reformat datasets, support qualitative and thematic data analysis, interpret results. AI assisted coding was used in the project to generate synthetic training sets, write training code, and write evaluation code. AI was used to suggest hyper-parameter choices and present analysis of outputs. All code and outputs also human verified. All theorems and proofs are original constructs of the authors. A first draft written with purely human analysis. This draft is then presented to the AI for suggested edits, which may be incorporated by the authors into a refined proof. All proofs are human verified as correct and began with the authors own steps of logic as the foundation.
We have not used generative AI tools to provide critical ingredients for proving mathematical claims, propose or refine hypotheses, or assist with translation. “Critical ingredients” for proofs is subjective. The authors do not consider AI as presenting any of the critical ideas that make the proofs work, but rather suggesting useful notation and suggesting more succinct summaries of steps. All logical steps in the proofs were human designed initially.
Additionally, we used generative AI tools to suggest experimental parameters, create or edit software code, summarize or analyse existing literature, discover research topics or identify gaps, brainstorming, sourcing/searching for information, edit a research paper to improve readability, identify relevant literature.
We have reviewed all AI-assisted work. We take responsibility for the final content of this work, including text, claims or artifacts produced with the aid of generative AI.
PROOFS
Proof of Lemma 1
Proof. Throughout the proof X X X Z Z Z X i 2 = 1 X_i^2 = 1 X i 2 = 1 i i i
Condition (1):
Note the moments E [ Z i ] = E [ Z i 3 ] = 0 E[Z_i] = E[Z_i^3] = 0 E [ Z i ] = E [ Z i 3 ] = 0 E [ Z i 2 ] = 1 E[Z_i^2] = 1 E [ Z i 2 ] = 1 E [ Z i 4 ] = 3 E[Z_i^4] = 3 E [ Z i 4 ] = 3
The random variable U U U X X X Z Z Z Z Z Z
E [ ∥ U ∥ 2 ] = a 2 E [ ∥ X ∥ 2 ] + 2 a 1 − a 2 ⟨ E [ X ] , E [ Z ] ⟩ + ( 1 − a 2 ) E [ ∥ Z ∥ 2 ] = a 2 E [ ∥ X ∥ 2 ] + ( 1 − a 2 ) E [ ∥ Z ∥ 2 ] . E[\lVert U\rVert^2] = a^2 E[\lVert X\rVert^2] + 2a\sqrt{1-a^2}\langle E[X], E[Z]\rangle+ (1-a^2)E[\lVert Z\rVert^2]
= a^2 E[\lVert X\rVert^2] + (1-a^2)E[\lVert Z\rVert^2]. E [∥ U ∥ 2 ] = a 2 E [∥ X ∥ 2 ] + 2 a 1 − a 2 ⟨ E [ X ] , E [ Z ]⟩ + ( 1 − a 2 ) E [∥ Z ∥ 2 ] = a 2 E [∥ X ∥ 2 ] + ( 1 − a 2 ) E [∥ Z ∥ 2 ] . Since the state space of X X X { − 1 , 1 } d + 1 \{-1,1\}^{d+1} { − 1 , 1 } d + 1 ∥ X ∥ 2 = d + 1 \lVert X\rVert^2 = d+1 ∥ X ∥ 2 = d + 1 Z Z Z E [ ∥ Z ∥ 2 ] = d + 1 E[\lVert Z\rVert^2] = d+1 E [∥ Z ∥ 2 ] = d + 1 d + 1 d+1 d + 1 a a a
E [ ( ∥ U ∥ 2 − ( d + 1 ) ) 2 ] = ∑ i , j = 1 d + 1 E [ ( U i 2 − 1 ) ( U j 2 − 1 ) ] . E[(\lVert U\rVert^2 - (d+1))^2] = \sum_{i,j=1}^{d+1} E[(U_i^2 - 1)(U_j^2 - 1)]. E [(∥ U ∥ 2 − ( d + 1 ) ) 2 ] = i , j = 1 ∑ d + 1 E [( U i 2 − 1 ) ( U j 2 − 1 )] . Note that X i 2 = 1 X_i^2 = 1 X i 2 = 1
U i 2 − 1 = 2 a 1 − a 2 X i Z i + ( 1 − a 2 ) ( Z i 2 − 1 ) . U_i^2 - 1 = 2a\sqrt{1-a^2}X_iZ_i + (1-a^2)(Z_i^2 - 1). U i 2 − 1 = 2 a 1 − a 2 X i Z i + ( 1 − a 2 ) ( Z i 2 − 1 ) . Consider first the diagonal terms i = j i=j i = j
E [ ( U i 2 − 1 ) 2 ] = 4 a 2 ( 1 − a 2 ) E [ Z i 2 ] + ( 1 − a 2 ) 2 E [ ( Z i 2 − 1 ) 2 ] + 4 a ( 1 − a 2 ) 3 / 2 E [ X i ] E [ Z i 3 − Z i ] = 4 a 2 ( 1 − a 2 ) + 2 ( 1 − a 2 ) 2 = 2 ( 1 − a 4 ) . \begin{aligned}
E[(U_i^2 - 1)^2] &= 4a^2(1-a^2)E[Z_i^2] + (1-a^2)^2E[(Z_i^2 - 1)^2] + 4a(1-a^2)^{3/2}E[X_i]E[Z_i^3-Z_i]\\
&= 4a^2(1-a^2) + 2(1-a^2)^2\\
&= 2(1-a^4).
\end{aligned} E [( U i 2 − 1 ) 2 ] = 4 a 2 ( 1 − a 2 ) E [ Z i 2 ] + ( 1 − a 2 ) 2 E [( Z i 2 − 1 ) 2 ] + 4 a ( 1 − a 2 ) 3/2 E [ X i ] E [ Z i 3 − Z i ] = 4 a 2 ( 1 − a 2 ) + 2 ( 1 − a 2 ) 2 = 2 ( 1 − a 4 ) . Then the cross terms i ≠ j i\ne j i = j E [ ( U i 2 − 1 ) ( U j 2 − 1 ) ] E[(U_i^2-1)(U_j^2-1)] E [( U i 2 − 1 ) ( U j 2 − 1 )] Z i ⊥ Z j Z_i \perp Z_j Z i ⊥ Z j X X X
E [ X i X j Z i Z j ] = E [ X i X j ] E [ Z i ] E [ Z j ] = 0 , E[X_iX_jZ_iZ_j] = E[X_iX_j]E[Z_i]E[Z_j] = 0, E [ X i X j Z i Z j ] = E [ X i X j ] E [ Z i ] E [ Z j ] = 0 , E [ X i Z i ( Z j 2 − 1 ) ] = E [ X i ] E [ Z i ] E [ Z j 2 − 1 ] = 0 , E[X_iZ_i(Z_j^2-1)] = E[X_i]E[Z_i]E[Z_j^2-1] = 0, E [ X i Z i ( Z j 2 − 1 )] = E [ X i ] E [ Z i ] E [ Z j 2 − 1 ] = 0 , E [ ( Z i 2 − 1 ) ( Z j 2 − 1 ) ] = E [ Z i 2 − 1 ] E [ Z j 2 − 1 ] = 0. E[(Z_i^2-1)(Z_j^2-1)] = E[Z_i^2-1]E[Z_j^2-1] = 0. E [( Z i 2 − 1 ) ( Z j 2 − 1 )] = E [ Z i 2 − 1 ] E [ Z j 2 − 1 ] = 0. Therefore Var ∥ U ∥ 2 = ∑ i Var U i 2 = 2 ( 1 − a 4 ) ( d + 1 ) \operatorname{Var}\lVert U\rVert^2 = \sum_i \operatorname{Var}U_i^2 = 2(1-a^4)(d+1) Var ∥ U ∥ 2 = ∑ i Var U i 2 = 2 ( 1 − a 4 ) ( d + 1 ) d + 1 d+1 d + 1
P ( ∣ ∥ U ∥ 2 d + 1 − 1 ∣ > ϵ ) = P ( ∣ ∥ U ∥ 2 − ( d + 1 ) ∣ > ϵ ( d + 1 ) ) ≤ 2 ( 1 − a 4 ) ( d + 1 ) ϵ 2 ( d + 1 ) 2 = 2 ( 1 − a 4 ) ϵ 2 ( d + 1 ) . P\left(\left|\frac{\lVert U\rVert^2}{d+1}-1\right|>\epsilon\right) = P\left(\left|\lVert U\rVert^2-(d+1)\right|>\epsilon(d+1)\right) \le\frac{2(1-a^4)(d+1)}{\epsilon^2(d+1)^2} = \frac{2(1-a^4)}{\epsilon^2(d+1)}. P ( d + 1 ∥ U ∥ 2 − 1 > ϵ ) = P ( ∥ U ∥ 2 − ( d + 1 ) > ϵ ( d + 1 ) ) ≤ ϵ 2 ( d + 1 ) 2 2 ( 1 − a 4 ) ( d + 1 ) = ϵ 2 ( d + 1 ) 2 ( 1 − a 4 ) . Condition (2):
Condition on X = x ∈ { − 1 , 1 } d + 1 X=x\in\{-1,1\}^{d+1} X = x ∈ { − 1 , 1 } d + 1 ∣ x i ∣ = 1 |x_i|=1 ∣ x i ∣ = 1
∣ U i ∣ = ∣ a x i + 1 − a 2 Z i ∣ = ∣ a + 1 − a 2 x i Z i ∣ . |U_i| = |ax_i+\sqrt{1-a^2}Z_i| = |a+\sqrt{1-a^2}x_iZ_i|. ∣ U i ∣ = ∣ a x i + 1 − a 2 Z i ∣ = ∣ a + 1 − a 2 x i Z i ∣. Coordinate-wise sign flips preserve N ( 0 , I ) N(0,I) N ( 0 , I ) ( x 1 Z 1 , … , x d + 1 Z d + 1 ) (x_1Z_1,\ldots,x_{d+1}Z_{d+1}) ( x 1 Z 1 , … , x d + 1 Z d + 1 ) Z Z Z X = x X=x X = x ( ∣ U 1 ∣ , … , ∣ U d + 1 ∣ ) (|U_1|,\ldots,|U_{d+1}|) ( ∣ U 1 ∣ , … , ∣ U d + 1 ∣ ) ( ∣ a + 1 − a 2 Z 1 ∣ , … , ∣ a + 1 − a 2 Z d + 1 ∣ ) (|a+\sqrt{1-a^2}Z_1|,\ldots,|a+\sqrt{1-a^2}Z_{d+1}|) ( ∣ a + 1 − a 2 Z 1 ∣ , … , ∣ a + 1 − a 2 Z d + 1 ∣ ) ∣ N ( a , 1 − a 2 ) ∣ |N(a,1-a^2)| ∣ N ( a , 1 − a 2 ) ∣ x x x π \pi π
Proof of Lemma 2
Proof. Let q ( x ) ∝ π ( x ) e τ ⟨ x , v ⟩ q(x) \propto\pi(x)e^{\tau\langle x,v\rangle} q ( x ) ∝ π ( x ) e τ ⟨ x , v ⟩ X t X_t X t
X t = t X + t Z , p ( x t ∣ x ) ∼ N ( t x , t I ) . X_t=tX+\sqrt{t}Z,\ p(x_t|x)\sim N(tx,tI). X t = tX + t Z , p ( x t ∣ x ) ∼ N ( t x , t I ) . The conditional density for p ( x ∣ x t ) p(x|x_t) p ( x ∣ x t )
p ( x ∣ x t ) ∝ q ( x ) e − 1 2 t ∥ x t − t x ∥ 2 p(x|x_t)\propto q(x)e^{-\frac{1}{2t}\lVert x_t-tx\rVert^2} p ( x ∣ x t ) ∝ q ( x ) e − 2 t 1 ∥ x t − t x ∥ 2 ∝ π ( x ) e τ ⟨ x , v ⟩ e − 1 2 t ∥ x t ∥ 2 + t 2 ⟨ x , x t ⟩ − t 2 ∥ x ∥ 2 , \propto\pi(x)e^{\tau\langle x,v\rangle}e^{-\frac{1}{2t}\lVert x_t\rVert^2+\frac{t}{2}\langle x,x_t\rangle-\frac{t}{2}\lVert x\rVert^2}, ∝ π ( x ) e τ ⟨ x , v ⟩ e − 2 t 1 ∥ x t ∥ 2 + 2 t ⟨ x , x t ⟩ − 2 t ∥ x ∥ 2 , with x x x x t x_t x t ∥ x t ∥ 2 \lVert x_t\rVert^2 ∥ x t ∥ 2 x x x x x x ∥ x ∥ 2 \lVert x\rVert^2 ∥ x ∥ 2 p ( x ∣ x t ) p(x|x_t) p ( x ∣ x t ) π \pi π
p ( x ∣ x t ) ∝ π ( x ) e ⟨ τ v + x t , x ⟩ . p(x|x_t)\propto\pi(x)e^{\langle\tau v+x_t,x\rangle}. p ( x ∣ x t ) ∝ π ( x ) e ⟨ τ v + x t , x ⟩ . Recall the definition of the tilt strength function β ( a ) = a 1 − a 2 \beta(a)=\frac{a}{1-a^2} β ( a ) = 1 − a 2 a a ∈ [ 0 , 1 ] a\in[0,1] a ∈ [ 0 , 1 ] θ = 1 + 5 2 ≈ 1.618 \theta=\frac{1+\sqrt{5}}{2}\approx1.618 θ = 2 1 + 5 ≈ 1.618 θ + 1 = θ 2 \theta+1=\theta^2 θ + 1 = θ 2 φ = θ − 1 ≈ 0.618 \varphi=\theta-1\approx0.618 φ = θ − 1 ≈ 0.618
1 − φ 2 = 1 − ( θ − 1 ) 2 = − θ 2 + 2 θ = − 1 + θ = φ . 1-\varphi^2=1-(\theta-1)^2=-\theta^2+2\theta=-1+\theta=\varphi. 1 − φ 2 = 1 − ( θ − 1 ) 2 = − θ 2 + 2 θ = − 1 + θ = φ . Thus it is the case that β ( φ ) = 1 \beta(\varphi)=1 β ( φ ) = 1 G G G φ \varphi φ τ v + X t \tau v+X_t τ v + X t
E π v [ X ∣ X t = t X + t Z ] = ∑ x x π ( x ) e β ( φ ) ⟨ τ v + X t , x ⟩ E π [ e β ( φ ) ⟨ τ v + X t , x ⟩ ] = G ( τ v + X t , φ ) . E_{\pi_v}[X|X_t=tX+\sqrt{t}Z]=\sum_x x\frac{\pi(x)e^{\beta(\varphi)\langle\tau v+X_t,x\rangle}}{E_{\pi}[e^{\beta(\varphi)\langle\tau v+X_t,x\rangle}]}=G(\tau v+X_t,\varphi). E π v [ X ∣ X t = tX + t Z ] = x ∑ x E π [ e β ( φ ) ⟨ τ v + X t , x ⟩ ] π ( x ) e β ( φ ) ⟨ τ v + X t , x ⟩ = G ( τ v + X t , φ ) . Proof of Lemma 3
Proof. For the mean-tilt operator G ( v , a ) G(v,a) G ( v , a ) β ( a ) = a 1 − a 2 \beta(a)=\frac{a}{1-a^2} β ( a ) = 1 − a 2 a ⟨ v , x ⟩ \langle v,x\rangle ⟨ v , x ⟩
G ( v , a ) = ∑ x x π ( x ) e β ( a ) ⟨ v , x ⟩ ∑ x ˉ π ( x ˉ ) e β ( a ) ⟨ v , x ˉ ⟩ . G(v,a)=\sum_x x\frac{\pi(x)e^{\beta(a)\langle v,x\rangle}}{\sum_{\bar{x}}\pi(\bar{x})e^{\beta(a)\langle v,\bar{x}\rangle}}. G ( v , a ) = x ∑ x ∑ x ˉ π ( x ˉ ) e β ( a ) ⟨ v , x ˉ ⟩ π ( x ) e β ( a ) ⟨ v , x ⟩ . The operator β \beta β [ 0 , 1 ] → [ 0 , ∞ ) [0,1]\to[0,\infty) [ 0 , 1 ] → [ 0 , ∞ ) b b b β \beta β a = β − 1 ( b ) a=\beta^{-1}(b) a = β − 1 ( b ) β \beta β
a = ( 1 − a 2 ) C a=(1-a^2)C a = ( 1 − a 2 ) C C a 2 + a − C = 0 ⟹ a = − 1 ± 1 + 4 C 2 2 C , Ca^2+a-C=0\Longrightarrow a=\frac{-1\pm\sqrt{1+4C^2}}{2C}, C a 2 + a − C = 0 ⟹ a = 2 C − 1 ± 1 + 4 C 2 , only the positive root results in a ∈ [ 0 , 1 ] a \in[0,1] a ∈ [ 0 , 1 ] C ∈ [ 0 , ∞ ) C \in[0,\infty) C ∈ [ 0 , ∞ )
β − 1 ( C ) = 1 2 ( 1 C 2 + 4 − 1 C ) . \beta^{-1}(C)=\frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right). β − 1 ( C ) = 2 1 ( C 2 1 + 4 − C 1 ) . For any inputs v , a v,a v , a v v v λ \lambda λ
∑ x x π ( x ) e β ( a ) ⟨ v , x ⟩ ∑ x ˉ π ( x ˉ ) e β ( a ) ⟨ v , x ˉ ⟩ = ∑ x x π ( x ) e ( β ( a ) λ ) ⟨ v λ , x ⟩ ∑ x ˉ π ( x ˉ ) e ( β ( a ) λ ) ⟨ v λ , x ˉ ⟩ . \sum_x x\frac{\pi(x)e^{\beta(a)\langle v,x\rangle}}{\sum_{\bar{x}}\pi(\bar{x})e^{\beta(a)\langle v,\bar{x}\rangle}}=\sum_x x\frac{\pi(x)e^{(\beta(a)\lambda)\left\langle\frac{v}{\lambda},x\right\rangle}}{\sum_{\bar{x}}\pi(\bar{x})e^{(\beta(a)\lambda)\left\langle\frac{v}{\lambda},\bar{x}\right\rangle}}. x ∑ x ∑ x ˉ π ( x ˉ ) e β ( a ) ⟨ v , x ˉ ⟩ π ( x ) e β ( a ) ⟨ v , x ⟩ = x ∑ x ∑ x ˉ π ( x ˉ ) e ( β ( a ) λ ) ⟨ λ v , x ˉ ⟩ π ( x ) e ( β ( a ) λ ) ⟨ λ v , x ⟩ . Setting λ = ∥ v ∥ d + 1 \lambda=\frac{\lVert v\rVert}{\sqrt{d+1}} λ = d + 1 ∥ v ∥ v ~ = v λ \tilde{v}=\frac{v}{\lambda} v ~ = λ v d + 1 \sqrt{d+1} d + 1 a ~ \tilde{a} a ~ β − 1 ( β ( a ) λ ) \beta^{-1}(\beta(a)\lambda) β − 1 ( β ( a ) λ )
Proof of Lemma 4
Proof. In this proof we take advantage of each coordinate of x x x { − 1 , 1 } \{-1,1\} { − 1 , 1 } x i 2 = 1 x_i^2=1 x i 2 = 1 i i i { − 1 , 1 } d \{-1,1\}^d { − 1 , 1 } d ∥ x ∥ 2 = d \lVert x\rVert^2=d ∥ x ∥ 2 = d U ∣ X U|X U ∣ X N ( a Λ X , ( 1 − a 2 ) I ) N(a\Lambda X,(1-a^2)I) N ( a Λ X , ( 1 − a 2 ) I )
p ( x ∣ u , a , Λ ) ∝ π ( x ) p ( u ∣ x , a , Λ ) p(x|u,a,\Lambda)\propto\pi(x)p(u|x,a,\Lambda) p ( x ∣ u , a , Λ ) ∝ π ( x ) p ( u ∣ x , a , Λ ) ∝ π ( x ) e − 1 2 ( 1 − a 2 ) ∥ u − a Λ x ∥ 2 \propto\pi(x)e^{-\frac{1}{2(1-a^2)}\lVert u-a\Lambda x\rVert^2} ∝ π ( x ) e − 2 ( 1 − a 2 ) 1 ∥ u − a Λ x ∥ 2 ∝ π ( x ) e a 1 − a 2 ⟨ Λ u , x ⟩ e − 1 2 ( 1 − a 2 ) ∥ u ∥ 2 − a 2 2 ( 1 − a 2 ) ∑ i = 1 d + 1 Λ i , i 2 x i 2 . \propto\pi(x)e^{\frac{a}{1-a^2}\langle\Lambda u,x\rangle}e^{-\frac{1}{2(1-a^2)}\lVert u\rVert^2-\frac{a^2}{2(1-a^2)}\sum_{i=1}^{d+1}\Lambda_{i,i}^2x_i^2}. ∝ π ( x ) e 1 − a 2 a ⟨ Λ u , x ⟩ e − 2 ( 1 − a 2 ) 1 ∥ u ∥ 2 − 2 ( 1 − a 2 ) a 2 ∑ i = 1 d + 1 Λ i , i 2 x i 2 . X X X u u u ∥ u ∥ 2 \lVert u\rVert^2 ∥ u ∥ 2 x i 2 = 1 x_i^2=1 x i 2 = 1 Λ i , i \Lambda_{i,i} Λ i , i x x x
Proof of Corollary 1
Proof. In this proof we collect the previous denoiser symmetry results of Lemma 3 and 4 the specify the form of the base denoiser under Λ \Lambda Λ d + 1 \sqrt{d+1} d + 1 a ~ \tilde{a} a ~ Λ \Lambda Λ G G G Λ U \Lambda U Λ U a a a
E π [ X ∣ U = a Λ X + 1 − a 2 Z ] = G ( Λ u , a ) , E_{\pi}[X|U=a\Lambda X+\sqrt{1-a^2}Z]=G(\Lambda u,a), E π [ X ∣ U = a Λ X + 1 − a 2 Z ] = G ( Λ u , a ) , we can then renormalize to u ~ = d + 1 Λ u ∥ Λ u ∥ \tilde{u}=\sqrt{d+1}\frac{\Lambda u}{\lVert\Lambda u\rVert} u ~ = d + 1 ∥ Λ u ∥ Λ u a ~ \tilde{a} a ~
Proof of Corollary 2
Proof. In this proof we collect the previous denoiser symmetry results of Lemma 2 and 3 to specify the form of the tilted denoiser under a stochastic localization path as a d + 1 \sqrt{d+1} d + 1 a ~ \tilde{a} a ~ τ v \tau v τ v φ ≈ 0.618 \varphi\approx0.618 φ ≈ 0.618
E π v [ X ∣ X t = t X + t Z ] = G ( τ v + X t , φ ) , E_{\pi_v}[X|X_t=tX+\sqrt{t}Z]=G(\tau v+X_t,\varphi), E π v [ X ∣ X t = tX + t Z ] = G ( τ v + X t , φ ) , we can then renormalize to u ~ = d + 1 τ v + X t ∥ τ v + X t ∥ \tilde{u}=\sqrt{d+1}\frac{\tau v+X_t}{\lVert\tau v+X_t\rVert} u ~ = d + 1 ∥ τ v + X t ∥ τ v + X t a ~ \tilde{a} a ~
Proof of Lemma 5
Proof. The loss function here is set up as an inverse problem, where for any two random variables with a joint distribution, the minimizer of the expected squared loss of one variable minus a function of the other is the conditional mean. The previous Lemmas have established that various conditional means can be expressed as evaluations of G G G d + 1 \sqrt{d+1} d + 1
Throughout the proof let ( a , Λ ) ∼ μ (a,\Lambda)\sim\mu ( a , Λ ) ∼ μ X ∼ π X\sim\pi X ∼ π Z ∼ N ( 0 , I ) Z\sim N(0,I) Z ∼ N ( 0 , I ) U = a Λ X + 1 − a 2 Z U=a\Lambda X+\sqrt{1-a^2}Z U = a Λ X + 1 − a 2 Z g g g
L ( g ) = E ( a , Λ ) ∼ μ , X ∼ π , Z ∥ X − g ( U , a , Λ ) ∥ 2 . L(g)=\mathbb{E}_{(a,\Lambda)\sim\mu,\,X\sim\pi,\,Z}\left\lVert X-g(U,a,\Lambda)\right\rVert^2. L ( g ) = E ( a , Λ ) ∼ μ , X ∼ π , Z ∥ X − g ( U , a , Λ ) ∥ 2 . The proof proceeds in three steps: (i) for fixed ( a , Λ ) (a,\Lambda) ( a , Λ ) X X X U U U ( a , Λ ) ∼ μ (a,\Lambda)\sim\mu ( a , Λ ) ∼ μ ( a , Λ ) (a,\Lambda) ( a , Λ ) X X X G ( Λ u , a ) G(\Lambda u,a) G ( Λ u , a ) G ( H ( Λ u , a ) ) G(H(\Lambda u,a)) G ( H ( Λ u , a ))
Step 1: fixed ( a , Λ ) (a,\Lambda) ( a , Λ ) Fix a ∈ [ 0 , 1 ] a\in[0,1] a ∈ [ 0 , 1 ] Λ \Lambda Λ U a , Λ = a Λ X + 1 − a 2 Z U_{a,\Lambda}=a\Lambda X+\sqrt{1-a^2}Z U a , Λ = a Λ X + 1 − a 2 Z X ∼ π X\sim\pi X ∼ π Z ∼ N ( 0 , I ) Z\sim N(0,I) Z ∼ N ( 0 , I ) f f f
E X , Z ∥ X − f ( U a , Λ ) ∥ 2 = E X , Z ∥ X − E [ X ∣ U a , Λ ] ∥ 2 + E X , Z ∥ E [ X ∣ U a , Λ ] − f ( U a , Λ ) ∥ 2 . \mathbb{E}_{X,Z}\left\lVert X-f(U_{a,\Lambda})\right\rVert^2
=
\mathbb{E}_{X,Z}\left\lVert X-\mathbb{E}[X\mid U_{a,\Lambda}]\right\rVert^2
+
\mathbb{E}_{X,Z}\left\lVert\mathbb{E}[X\mid U_{a,\Lambda}]-f(U_{a,\Lambda})\right\rVert^2. E X , Z ∥ X − f ( U a , Λ ) ∥ 2 = E X , Z ∥ X − E [ X ∣ U a , Λ ] ∥ 2 + E X , Z ∥ E [ X ∣ U a , Λ ] − f ( U a , Λ ) ∥ 2 . It is then a standard fact that the minimizing function is the conditional mean
g a , Λ ( u ) : = E X ∼ π [ X ∣ U a , Λ = u ] g_{a,\Lambda}(u):=\mathbb{E}_{X\sim\pi}[X\mid U_{a,\Lambda}=u] g a , Λ ( u ) := E X ∼ π [ X ∣ U a , Λ = u ] satisfies
g a , Λ = argmin f E X , Z ∥ X − f ( U a , Λ ) ∥ 2 . g_{a,\Lambda}=\underset{f}{\operatorname{argmin}}\ \mathbb{E}_{X,Z}\left\lVert X-f(U_{a,\Lambda})\right\rVert^2. g a , Λ = f argmin E X , Z ∥ X − f ( U a , Λ ) ∥ 2 . By Lemma 4, the posterior of X X X U a , Λ = u U_{a,\Lambda}=u U a , Λ = u π ( x ) e β ( a ) ⟨ Λ u , x ⟩ \pi(x)e^{\beta(a)\langle\Lambda u,x\rangle} π ( x ) e β ( a ) ⟨ Λ u , x ⟩ g a , Λ ( u ) = G ( Λ u , a ) g_{a,\Lambda}(u)=G(\Lambda u,a) g a , Λ ( u ) = G ( Λ u , a )
Step 2: random ( a , Λ ) ∼ μ (a,\Lambda)\sim\mu ( a , Λ ) ∼ μ Condition on ( a , Λ ) (a,\Lambda) ( a , Λ ) ( a , Λ ) (a,\Lambda) ( a , Λ )
L ( g ) = E ( a , Λ ) ∼ μ [ E [ ∥ X − g ( U , a , Λ ) ∥ 2 ∣ a , Λ ] ] . L(g)=\mathbb{E}_{(a,\Lambda)\sim\mu}\left[\mathbb{E}\left[\left\lVert X-g(U,a,\Lambda)\right\rVert^2\mid a,\Lambda\right]\right]. L ( g ) = E ( a , Λ ) ∼ μ [ E [ ∥ X − g ( U , a , Λ ) ∥ 2 ∣ a , Λ ] ] . Because ( a , Λ ) (a,\Lambda) ( a , Λ ) ( X , Z ) (X,Z) ( X , Z ) ( X , U ) (X,U) ( X , U ) ( a , Λ ) (a,\Lambda) ( a , Λ ) ( X , U a , Λ ) (X,U_{a,\Lambda}) ( X , U a , Λ )
p ( x , u ∣ a , Λ ) = p ( u ∣ x , a , Λ ) p ( x ∣ a , Λ ) = p ( u ∣ x , a , Λ ) π ( x ) , p(x,u\mid a,\Lambda)=p(u\mid x,a,\Lambda)p(x\mid a,\Lambda)=p(u\mid x,a,\Lambda)\pi(x), p ( x , u ∣ a , Λ ) = p ( u ∣ x , a , Λ ) p ( x ∣ a , Λ ) = p ( u ∣ x , a , Λ ) π ( x ) , i.e. X X X ( a , Λ ) (a,\Lambda) ( a , Λ ) ( a , Λ ) (a,\Lambda) ( a , Λ ) p ( u ∣ x , a , Λ ) p(u\mid x,a,\Lambda) p ( u ∣ x , a , Λ )
E [ ∥ X − g ( U , a , Λ ) ∥ 2 ∣ a , Λ ] = E X ∼ π , Z ∥ X − g ( U a , Λ , a , Λ ) ∥ 2 , \mathbb{E}\left[\left\lVert X-g(U,a,\Lambda)\right\rVert^2\mid a,\Lambda\right]
=
\mathbb{E}_{X\sim\pi,\,Z}\left\lVert X-g(U_{a,\Lambda},a,\Lambda)\right\rVert^2, E [ ∥ X − g ( U , a , Λ ) ∥ 2 ∣ a , Λ ] = E X ∼ π , Z ∥ X − g ( U a , Λ , a , Λ ) ∥ 2 , which, for each fixed ( a , Λ ) (a,\Lambda) ( a , Λ ) f = g ( ⋅ , a , Λ ) f=g(\cdot,a,\Lambda) f = g ( ⋅ , a , Λ ) g g g ( a , Λ ) (a,\Lambda) ( a , Λ ) g ( ⋅ , a , Λ ) g(\cdot,a,\Lambda) g ( ⋅ , a , Λ ) ( a , Λ ) (a,\Lambda) ( a , Λ ) ( a , Λ ) (a,\Lambda) ( a , Λ )
min g L ( g ) = E ( a , Λ ) ∼ μ [ min f E X , Z ∥ X − f ( U a , Λ ) ∥ 2 ] , g ⋆ ( u , a , Λ ) = g a , Λ ( u ) = G ( Λ u , a ) . \min_g L(g)=\mathbb{E}_{(a,\Lambda)\sim\mu}\left[\min_f\mathbb{E}_{X,Z}\left\lVert X-f(U_{a,\Lambda})\right\rVert^2\right],
\qquad
g^\star(u,a,\Lambda)=g_{a,\Lambda}(u)=G(\Lambda u,a). g min L ( g ) = E ( a , Λ ) ∼ μ [ f min E X , Z ∥ X − f ( U a , Λ ) ∥ 2 ] , g ⋆ ( u , a , Λ ) = g a , Λ ( u ) = G ( Λ u , a ) . Step 3: renormalization. By Lemma 3 , for any v ∈ R d + 1 ∖ { 0 } v \in\mathbb{R}^{d+1} \setminus\{0\} v ∈ R d + 1 ∖ { 0 } a ∈ [ 0 , 1 ] a \in[0,1] a ∈ [ 0 , 1 ]
G ( v , a ) = G ( d + 1 v ∥ v ∥ , a ~ ) = G ( H ( v , a ) ) , a ~ = β − 1 ( β ( a ) ∥ v ∥ d + 1 ) . G(v,a) = G\left(\sqrt{d+1}\frac{v}{\lVert v\rVert},\tilde{a}\right) = G(H(v,a)), \qquad\tilde{a} = \beta^{-1}\left(\beta(a)\frac{\lVert v\rVert}{\sqrt{d+1}}\right). G ( v , a ) = G ( d + 1 ∥ v ∥ v , a ~ ) = G ( H ( v , a )) , a ~ = β − 1 ( β ( a ) d + 1 ∥ v ∥ ) . so g ⋆ ( u , a , Λ ) = G ( H ( Λ u , a ) ) g^{\star}(u,a,\Lambda) = G(H(\Lambda u,a)) g ⋆ ( u , a , Λ ) = G ( H ( Λ u , a )) ( U , a , Λ ) (U,a,\Lambda) ( U , a , Λ ) H ( Λ U , a ) H(\Lambda U,a) H ( Λ U , a ) { g ∘ H ( Λ ⋅ , ⋅ ) : g measurable on S × [ 0 , 1 ] } \{g \circ H(\Lambda\cdot,\cdot):g\text{ measurable on }S \times[0,1]\} { g ∘ H ( Λ ⋅ , ⋅ ) : g measurable on S × [ 0 , 1 ]} ( U , a , Λ ) (U,a,\Lambda) ( U , a , Λ )
G = argmin g E ∥ X − g ( H ( Λ U , a ) ) ∥ 2 , G = \underset{g}{\operatorname{argmin}}\ \mathbb{E}\left\lVert X-g(H(\Lambda U,a))\right\rVert^{2}, G = g argmin E ∥ X − g ( H ( Λ U , a )) ∥ 2 , uniquely up to sets of measure zero under the law of H ( Λ U , a ) H(\Lambda U,a) H ( Λ U , a ) □ \square □
Additional Discussion
Visualization of Denoiser Symmetry
Figure 5. Sketch of overlap between renormalized training set via Λ focusing and points shifted by tilt along stochastic localization path.
Stochastic Localization
Stochastic localization can be viewed as a different time scaling and spatial scaling than the usual score-based diffusion derived from the forwards flow Ornstein-Uhlenbeck process as outlined in [23 ] . We re-derive the equivalence here for an audience more familiar with the score-based construction. The score-based diffusion is variance preserving, initialized at a standard normal draw, and uses the score as its drift. The stochastic localization process is not variance preserving (it is growing in magnitude over time), is initialized at a Dirac mass at the origin, and uses the denoiser as its drift term. Both can be derived from the same reverse diffusion process theory.
Time reversal. Both constructions are reverse-time implementations of a forward diffusion producing a sequence of marginal densities that are equivalent to convolutions with Gaussian random variables of various mixture amounts. By the time-reversal result of Anderson [1 ] , popularized in modern literature by Song et al. [32 ] , Lipman et al. [19 ] , consider any target density π \pi π X 0 ∼ π X_0 \sim\pi X 0 ∼ π
d X t = f ( X t , t ) d t + g ( t ) d B t , t ∈ [ 0 , T ] , X 0 ∼ π , dX_t = f(X_t,t)\,dt + g(t)\,dB_t,\qquad t \in[0,T],\ X_0 \sim\pi, d X t = f ( X t , t ) d t + g ( t ) d B t , t ∈ [ 0 , T ] , X 0 ∼ π , and let p t p_t p t X t X_t X t s t = ∇ log p t s_t = \nabla\log p_t s t = ∇ log p t Y 0 ∼ p T Y_0 \sim p_T Y 0 ∼ p T
d Y t = − [ f ( Y t , T − t ) − g ( T − t ) 2 s T − t ( Y t ) ] d t + g ( T − t ) d B t , (1) dY_t = -\left[f(Y_t,T-t)-g(T-t)^2s_{T-t}(Y_t)\right]dt + g(T-t)\,dB_t, \tag*{(1)} d Y t = − [ f ( Y t , T − t ) − g ( T − t ) 2 s T − t ( Y t ) ] d t + g ( T − t ) d B t , (1) satisfies Y t ∼ p T − t Y_t \sim p_{T-t} Y t ∼ p T − t f f f g g g p t p_t p t s t s_t s t
Linear forward processes. When the forward drift f ( X t , t ) f(X_t,t) f ( X t , t ) X t X_t X t [25 ] . Let α , σ : [ 0 , T ] → R + \alpha,\sigma: [0,T] \to\mathbb{R}_+ α , σ : [ 0 , T ] → R + σ / α \sigma/\alpha σ / α
κ ( t ) : = d d t log σ ( t ) α ( t ) = σ ˙ ( t ) σ ( t ) − α ˙ ( t ) α ( t ) ≥ 0. \kappa(t) := \frac{d}{dt}\log\frac{\sigma(t)}{\alpha(t)} = \frac{\dot{\sigma}(t)}{\sigma(t)}-\frac{\dot{\alpha}(t)}{\alpha(t)} \ge0. κ ( t ) := d t d log α ( t ) σ ( t ) = σ ( t ) σ ˙ ( t ) − α ( t ) α ˙ ( t ) ≥ 0. Consider the linear forward process
d X t = α ˙ ( t ) α ( t ) X t d t + σ ( t ) 2 κ ( t ) d B t , X 0 = α ( 0 ) X + σ ( 0 ) Z , (2) dX_t = \frac{\dot{\alpha}(t)}{\alpha(t)}X_t\,dt+\sigma(t)\sqrt{2\kappa(t)}\,dB_t,\qquad X_0=\alpha(0)X+\sigma(0)Z, \tag*{(2)} d X t = α ( t ) α ˙ ( t ) X t d t + σ ( t ) 2 κ ( t ) d B t , X 0 = α ( 0 ) X + σ ( 0 ) Z , (2) with X ∼ π X\sim\pi X ∼ π Z ∼ N ( 0 , I ) Z\sim N(0,I) Z ∼ N ( 0 , I )
X t = d α ( t ) X + σ ( t ) Z , i.e. p t ( x t ∣ X = x 0 ) = N ( x t ; α ( t ) x 0 , σ ( t ) 2 I ) . X_t \mathrel{\overset{d}{=}} \alpha(t)X+\sigma(t)Z,\qquad\text{i.e.}\qquad p_t(x_t\mid X=x_0)=N\left(x_t;\alpha(t)x_0,\sigma(t)^2I\right). X t = d α ( t ) X + σ ( t ) Z , i.e. p t ( x t ∣ X = x 0 ) = N ( x t ; α ( t ) x 0 , σ ( t ) 2 I ) . Writing m t ( x ) : = E [ X ∣ X t = x ] m_t(x):=\mathbb{E}[X\mid X_t=x] m t ( x ) := E [ X ∣ X t = x ] [25 ] gives the score as a posterior mean under a quadratic tilt of π \pi π
s t ( x t ) = ∫ − x t − α ( t ) x 0 σ ( t ) 2 π ( x 0 ) e − 1 2 σ ( t ) 2 ∥ x t − α ( t ) x 0 ∥ 2 ∫ π ( x 0 ′ ) e − 1 2 σ ( t ) 2 ∥ x t − α ( t ) x 0 ′ ∥ 2 d x 0 ′ d x 0 = α ( t ) m t ( x t ) − x t σ ( t ) 2 . s_t(x_t)=\int-\frac{x_t-\alpha(t)x_0}{\sigma(t)^2}\frac{\pi(x_0)e^{-\frac{1}{2\sigma(t)^2}\lVert x_t-\alpha(t)x_0\rVert^2}}{\int\pi(x_0')e^{-\frac{1}{2\sigma(t)^2}\lVert x_t-\alpha(t)x_0'\rVert^2}\,dx_0'}\,dx_0=\frac{\alpha(t)m_t(x_t)-x_t}{\sigma(t)^2}. s t ( x t ) = ∫ − σ ( t ) 2 x t − α ( t ) x 0 ∫ π ( x 0 ′ ) e − 2 σ ( t ) 2 1 ∥ x t − α ( t ) x 0 ′ ∥ 2 d x 0 ′ π ( x 0 ) e − 2 σ ( t ) 2 1 ∥ x t − α ( t ) x 0 ∥ 2 d x 0 = σ ( t ) 2 α ( t ) m t ( x t ) − x t . Substituting into (1), and abbreviating α = α ( T − t ) \alpha=\alpha(T-t) α = α ( T − t ) σ = σ ( T − t ) \sigma=\sigma(T-t) σ = σ ( T − t ) κ = κ ( T − t ) \kappa=\kappa(T-t) κ = κ ( T − t )
d Y t = [ − α ˙ α Y t + 2 σ 2 κ s T − t ( Y t ) ] d t + σ 2 κ d B t . (3) dY_t=\left[-\frac{\dot{\alpha}}{\alpha}Y_t+2\sigma^2\kappa s_{T-t}(Y_t)\right]dt+\sigma\sqrt{2\kappa}\,dB_t. \tag*{(3)} d Y t = [ − α α ˙ Y t + 2 σ 2 κ s T − t ( Y t ) ] d t + σ 2 κ d B t . (3) Equivalently stated in terms of the denoiser instead of the score,
d Y t = [ − ( α ˙ α + 2 κ ) Y t + 2 κ α m T − t ( Y t ) ] d t + σ 2 κ d B t . (4) dY_t=\left[-\left(\frac{\dot{\alpha}}{\alpha}+2\kappa\right)Y_t+2\kappa\alpha m_{T-t}(Y_t)\right]dt+\sigma\sqrt{2\kappa}\,dB_t. \tag*{(4)} d Y t = [ − ( α α ˙ + 2 κ ) Y t + 2 κ α m T − t ( Y t ) ] d t + σ 2 κ d B t . (4) Signal level and growth. A linear combination α ( t ) X + σ ( t ) Z \alpha(t)X+\sigma(t)Z α ( t ) X + σ ( t ) Z α ( t ) 2 + σ ( t ) 2 = 1 \alpha(t)^2+\sigma(t)^2=1 α ( t ) 2 + σ ( t ) 2 = 1 t t t λ \lambda λ τ \tau τ
λ ( t ) 2 : = α ( t ) 2 α ( t ) 2 + σ ( t ) 2 , τ ( t ) : = α ( t ) 2 + σ ( t ) 2 , \lambda(t)^2 := \frac{\alpha(t)^2}{\alpha(t)^2+\sigma(t)^2}, \qquad\tau(t):=\sqrt{\alpha(t)^2+\sigma(t)^2}, λ ( t ) 2 := α ( t ) 2 + σ ( t ) 2 α ( t ) 2 , τ ( t ) := α ( t ) 2 + σ ( t ) 2 , so that
α ( t ) X + σ ( t ) Z = τ ( t ) ( λ ( t ) X + 1 − λ ( t ) 2 Z ) . \alpha(t)X+\sigma(t)Z=\tau(t)\left(\lambda(t)X+\sqrt{1-\lambda(t)^2}Z\right). α ( t ) X + σ ( t ) Z = τ ( t ) ( λ ( t ) X + 1 − λ ( t ) 2 Z ) . The VP case is τ = 1 \tau=1 τ = 1
Score-based diffusion and stochastic localization are defined by two different choices of α \alpha α σ \sigma σ
Score-Based Diffusion from the Ornstein–Uhlenbeck Process
Taking α ( t ) = e − t \alpha(t)=e^{-t} α ( t ) = e − t σ ( t ) = 1 − e − 2 t \sigma(t)=\sqrt{1-e^{-2t}} σ ( t ) = 1 − e − 2 t α ˙ / α = − 1 \dot{\alpha}/\alpha=-1 α ˙ / α = − 1 2 σ 2 κ = 2 2\sigma^2\kappa=2 2 σ 2 κ = 2 X t X_t X t
d X t = − X t d t + 2 d B t , X 0 ∼ π , \mathrm{d}X_t=-X_t\,\mathrm{d}t+\sqrt{2}\,\mathrm{d}B_t,\qquad X_0\sim\pi, d X t = − X t d t + 2 d B t , X 0 ∼ π , d Y t = ( Y t + 2 s T − t ( Y t ) ) d t + 2 d B t . \mathrm{d}Y_t=\left(Y_t+2s_{T-t}(Y_t)\right)\mathrm{d}t+\sqrt{2}\,\mathrm{d}B_t. d Y t = ( Y t + 2 s T − t ( Y t ) ) d t + 2 d B t . This is a VP combination, τ ( t ) ≡ 1 \tau(t)\equiv1 τ ( t ) ≡ 1 λ ( t ) 2 = e − 2 t \lambda(t)^2=e^{-2t} λ ( t ) 2 = e − 2 t
The Stochastic Localization Time Scaling
Fix T > 0 T>0 T > 0 α ( t ) = T − t \alpha(t)=T-t α ( t ) = T − t σ ( t ) = T − t \sigma(t)=\sqrt{T-t} σ ( t ) = T − t t ∈ [ 0 , T ] t\in[0,T] t ∈ [ 0 , T ]
α ˙ ( t ) α ( t ) = − 1 T − t , σ ˙ ( t ) σ ( t ) = − 1 2 ( T − t ) , κ ( t ) = 1 2 ( T − t ) , σ ( t ) 2 2 κ ( t ) = 1 , \frac{\dot{\alpha}(t)}{\alpha(t)}=-\frac{1}{T-t},\qquad\frac{\dot{\sigma}(t)}{\sigma(t)}=-\frac{1}{2(T-t)},\qquad\kappa(t)=\frac{1}{2(T-t)},\qquad\sigma(t)^2 2\kappa(t)=1, α ( t ) α ˙ ( t ) = − T − t 1 , σ ( t ) σ ˙ ( t ) = − 2 ( T − t ) 1 , κ ( t ) = 2 ( T − t ) 1 , σ ( t ) 2 2 κ ( t ) = 1 , so the forward process (2) is
d X t = − 1 T − t X t d t + d B t , X 0 = T X + T Z , \mathrm{d}X_t=-\frac{1}{T-t}X_t\,\mathrm{d}t+\mathrm{d}B_t,\qquad X_0=TX+\sqrt{T}Z, d X t = − T − t 1 X t d t + d B t , X 0 = T X + T Z , which is driven to X T = 0 X_T=0 X T = 0 Y t Y_t Y t α ˙ / α + 2 κ = 0 \dot{\alpha}/\alpha+2\kappa=0 α ˙ / α + 2 κ = 0 2 κ α = 1 2\kappa\alpha=1 2 κ α = 1
d Y t = E [ X ∣ t X + t Z = Y t ] d t + d B t , Y 0 = 0. \mathrm{d}Y_t=E[X\mid tX+\sqrt{t}Z=Y_t]\,\mathrm{d}t+\mathrm{d}B_t,\qquad Y_0=0. d Y t = E [ X ∣ tX + t Z = Y t ] d t + d B t , Y 0 = 0. This is the stochastic localization process of Montanari [23 ] , with Y t = d t X + t Z Y_t\mathrel{\stackrel{\mathrm{d}}{=}}tX+\sqrt{t}Z Y t = d tX + t Z Y t / t → X Y_t/t\to X Y t / t → X t → ∞ t\to\infty t → ∞
Relation to the OU diffusion. Both processes are naturally indexed by their reverse-time clocks, along which the signal level increases. In stochastic localization time s s s
λ S L ( s ) 2 = s 1 + s , τ S L ( s ) = s ( 1 + s ) . \lambda_{\mathrm{SL}}(s)^2 = \frac{s}{1+s}, \qquad\tau_{\mathrm{SL}}(s) = \sqrt{s(1+s)}. λ SL ( s ) 2 = 1 + s s , τ SL ( s ) = s ( 1 + s ) . For the OU diffusion on a horizon [ 0 , T ] [0,T] [ 0 , T ] Y r O U ∼ X T − r Y^{\mathrm{OU}}_r \sim X_{T-r} Y r OU ∼ X T − r r ∈ [ 0 , T ] r \in[0,T] r ∈ [ 0 , T ]
λ O U ( r ) 2 = e − 2 ( T − r ) , τ O U = 1. \lambda_{\mathrm{OU}}(r)^2 = e^{-2(T-r)}, \qquad\tau_{\mathrm{OU}} = 1. λ OU ( r ) 2 = e − 2 ( T − r ) , τ OU = 1. Matching signal levels, e − 2 ( T − r ) = s / ( 1 + s ) e^{-2(T-r)} = s/(1+s) e − 2 ( T − r ) = s / ( 1 + s )
r ( s ) = T − 1 2 log ( 1 + 1 / s ) , r(s) = T - \frac{1}{2}\log(1+1/s), r ( s ) = T − 2 1 log ( 1 + 1/ s ) , which is increasing in s s s s ∈ [ ( e 2 T − 1 ) − 1 , ∞ ) s \in[(e^{2T}-1)^{-1},\infty) s ∈ [( e 2 T − 1 ) − 1 , ∞ ) [ 0 , T ) [0,T) [ 0 , T ) T → ∞ T \to\infty T → ∞ s > 0 s>0 s > 0
Y s S L = d s ( 1 + s ) Y r ( s ) O U . Y^{\mathrm{SL}}_s \mathrel{\mathop{\mathrm{=}}\limits^{d}} \sqrt{s(1+s)}\,Y^{\mathrm{OU}}_{r(s)}. Y s SL = d s ( 1 + s ) Y r ( s ) OU . Stochastic localization is thus the OU reverse diffusion run under the time change r ( s ) r(s) r ( s ) τ S L ( s ) \tau_{\mathrm{SL}}(s) τ SL ( s ) s s s r ( s ) r(s) r ( s ) s ( 1 + s ) \sqrt{s(1+s)} s ( 1 + s ) Figure 6 to visualize the equivalence.
Figure 6. Visualization of score-based diffusion and stochastic localization reverse flow path for an asymmetric mixture of normals target density in one dimension. Marginal densities produced at each stage of the processes are equivalent under a time shift and rescaling.
Log-Odds Calculation
All tilts used in the experiments are linear tilts of the extended base density by a vector supported on the reward coordinates. A linear tilt reweights π \pi π τ \tau τ
Single reward bit. Let r : { − 1 , 1 } d → { − 1 , 1 } r:\{-1,1\}^{d} \to\{-1,1\} r : { − 1 , 1 } d → { − 1 , 1 } A = { x : r ( x ) = 1 } A=\{x:r(x)=1\} A = { x : r ( x ) = 1 } p = π ( A ) p=\pi(A) p = π ( A ) τ e ∗ \tau e_\ast τ e ∗
q τ ( x ) = π ( x ) e τ r ( x ) E π [ e τ r ( X ) ] = π ( x ∣ A ) p e τ p e τ + ( 1 − p ) e − τ + π ( x ∣ A c ) ( 1 − p ) e − τ p e τ + ( 1 − p ) e − τ . q_{\tau}(x)=\frac{\pi(x)e^{\tau r(x)}}{E_{\pi}[e^{\tau r(X)}]}=\pi(x\mid A)\frac{p e^{\tau}}{p e^{\tau}+(1-p)e^{-\tau}}+\pi(x\mid A^{c})\frac{(1-p)e^{-\tau}}{p e^{\tau}+(1-p)e^{-\tau}}. q τ ( x ) = E π [ e τ r ( X ) ] π ( x ) e τ r ( x ) = π ( x ∣ A ) p e τ + ( 1 − p ) e − τ p e τ + π ( x ∣ A c ) p e τ + ( 1 − p ) e − τ ( 1 − p ) e − τ . The conditional laws π ( ⋅ ∣ A ) \pi(\cdot\mid A) π ( ⋅ ∣ A ) π ( ⋅ ∣ A c ) \pi(\cdot\mid A^{c}) π ( ⋅ ∣ A c ) logit p = log p 1 − p \operatorname{logit}p=\log\frac{p}{1-p} logit p = log 1 − p p p τ = q τ ( A ) p_{\tau}=q_{\tau}(A) p τ = q τ ( A )
logit p τ = logit p + 2 τ . (5) \operatorname{logit}p_{\tau}=\operatorname{logit}p+2\tau. \tag*{(5)} logit p τ = logit p + 2 τ . (5) The slope 2 is the difference r ( x ) − r ( x ′ ) = 2 r(x)-r(x')=2 r ( x ) − r ( x ′ ) = 2 p p p
One-hot class block. For K K K K K K c c c r c ∈ { − 1 , 1 } K r_c\in\{-1,1\}^{K} r c ∈ { − 1 , 1 } K + 1 +1 + 1 c c c − 1 -1 − 1 h = τ r c h=\tau r_c h = τ r c k k k exp ( τ ⟨ r c , r k ⟩ ) \exp(\tau\langle r_c,r_k\rangle) exp ( τ ⟨ r c , r k ⟩) ⟨ r c , r c ⟩ = K \langle r_c,r_c\rangle=K ⟨ r c , r c ⟩ = K ⟨ r c , r k ⟩ = K − 4 \langle r_c,r_k\rangle=K-4 ⟨ r c , r k ⟩ = K − 4 k ≠ c k\ne c k = c e ( K − 4 ) τ e^{(K-4)\tau} e ( K − 4 ) τ
π τ ( c ) = π ( c ) e 4 τ π ( c ) e 4 τ + 1 − π ( c ) , π τ ( k ) = π ( k ) π ( c ) e 4 τ + 1 − π ( c ) , k ≠ c , \pi_{\tau}(c)=\frac{\pi(c)e^{4\tau}}{\pi(c)e^{4\tau}+1-\pi(c)},\qquad\pi_{\tau}(k)=\frac{\pi(k)}{\pi(c)e^{4\tau}+1-\pi(c)},\qquad k\ne c, π τ ( c ) = π ( c ) e 4 τ + 1 − π ( c ) π ( c ) e 4 τ , π τ ( k ) = π ( c ) e 4 τ + 1 − π ( c ) π ( k ) , k = c , so that
logit π τ ( c ) = logit π ( c ) + 4 τ , (6) \operatorname{logit}\pi_{\tau}(c)=\operatorname{logit}\pi(c)+4\tau, \tag*{(6)} logit π τ ( c ) = logit π ( c ) + 4 τ , (6) and π τ ( k ) / π τ ( k ′ ) = π ( k ) / π ( k ′ ) \pi_{\tau}(k)/\pi_{\tau}(k')=\pi(k)/\pi(k') π τ ( k ) / π τ ( k ′ ) = π ( k ) / π ( k ′ ) k , k ′ ≠ c k,k'\ne c k , k ′ = c ⟨ r c , r c ⟩ − ⟨ r c , r k ⟩ \langle r_c,r_c\rangle-\langle r_c,r_k\rangle ⟨ r c , r c ⟩ − ⟨ r c , r k ⟩ e c e_c e c h h h c c c k k k ⟨ h , r c − r k ⟩ \langle h,r_c-r_k\rangle ⟨ h , r c − r k ⟩ (5) and (6) are the two cases used.
Tilt grid. Inverting (5) , the tilt strength at which reward is achieved with probability p ⋆ p^\star p ⋆
τ ( p ⋆ ) = 1 2 ( logit p ⋆ − logit p ) . \tau(p^\star)=\frac{1}{2}\left(\operatorname{logit}p^\star-\operatorname{logit}p\right). τ ( p ⋆ ) = 2 1 ( logit p ⋆ − logit p ) . with 1 4 \frac{1}{4} 4 1 1 2 \frac{1}{2} 2 1 (6) . Experiments fix τ 98 = τ ( 0.98 ) \tau_{98}=\tau(0.98) τ 98 = τ ( 0.98 ) τ 95 \tau_{95} τ 95 τ 98 \tau_{98} τ 98 d + 1 \sqrt{d+1} d + 1
Estimation. At each τ \tau τ N N N N = 2000 N = 2000 N = 2000 p ^ τ \hat{p}_{\tau} p ^ τ p τ = 1 p_{\tau} = 1 p τ = 1 p ^ = 1 \hat{p} = 1 p ^ = 1 N N N logit ( 1 − 1 / N ) \operatorname{logit}(1 - 1/N) logit ( 1 − 1/ N ) N = 2000 N = 2000 N = 2000 τ \tau τ § 10 .
Clopper–Pearson Confidence Intervals. Let K ∼ Binomial ( N , p τ ) K \sim\operatorname{Binomial}(N, p_{\tau}) K ∼ Binomial ( N , p τ ) τ \tau τ k k k p ^ τ = k / N \hat{p}_{\tau} = k/N p ^ τ = k / N p p p α / 2 \alpha/2 α /2
∑ j = k N ( N j ) p L j ( 1 − p L ) N − j = α 2 , ∑ j = 0 k ( N j ) p U j ( 1 − p U ) N − j = α 2 . \sum_{j=k}^{N} \binom{N}{j}p_L^j(1-p_L)^{N-j} = \frac{\alpha}{2}, \qquad\sum_{j=0}^{k} \binom{N}{j}p_U^j(1-p_U)^{N-j} = \frac{\alpha}{2}. j = k ∑ N ( j N ) p L j ( 1 − p L ) N − j = 2 α , j = 0 ∑ k ( j N ) p U j ( 1 − p U ) N − j = 2 α . By the identity linking binomial tails to the regularized incomplete beta function, these endpoints have the closed form
p L = B − 1 ( α 2 ; k , N − k + 1 ) , p U = B − 1 ( 1 − α 2 ; k + 1 , N − k ) . p_L = B^{-1}\left(\frac{\alpha}{2}; k, N-k+1\right), \qquad p_U = B^{-1}\left(1-\frac{\alpha}{2}; k+1, N-k\right). p L = B − 1 ( 2 α ; k , N − k + 1 ) , p U = B − 1 ( 1 − 2 α ; k + 1 , N − k ) . where B − 1 ( q ; a , b ) B^{-1}(q; a,b) B − 1 ( q ; a , b ) q q q ( a , b ) (a,b) ( a , b ) p L = 0 p_L = 0 p L = 0 k = 0 k = 0 k = 0 p U = 1 p_U = 1 p U = 1 k = N k = N k = N 1 − α 1-\alpha 1 − α p τ p_{\tau} p τ N N N p τ = 1 p_{\tau} = 1 p τ = 1 1 − α 1-\alpha 1 − α [ logit ( p L ) , logit ( p U ) ] [\operatorname{logit}(p_L), \operatorname{logit}(p_U)] [ logit ( p L ) , logit ( p U )] 1 − α 1-\alpha 1 − α k = N k = N k = N [ ( α / 2 ) 1 / N , 1 ] [(\alpha/2)^{1/N},1] [( α /2 ) 1/ N , 1 ] + ∞ +\infty + ∞ logit ( ( α / 2 ) 1 / N ) ≈ 6.3 \operatorname{logit}((\alpha/2)^{1/N}) \approx6.3 logit (( α /2 ) 1/ N ) ≈ 6.3 N = 2000 N = 2000 N = 2000 α = 0.05 \alpha= 0.05 α = 0.05
Model-reported reward. Since the reward coordinate is denoised jointly with the token coordinates, the sampler also produces an internal prediction of reward: the sign of the terminal d + 1 d+1 d + 1 § 10 are distinguished by this comparison.
Self-Conditioning
Self-conditioning as originally introduced in Chen et al. [7 ] is a method to improve the sample quality of a denoising diffusion model. It was originally devised as a mostly ad hoc engineering decision, rather than a method arising from sound theory. Its interaction with standard theory for error control for diffusion models has not yet been fully understood [30 , 18 , 6 , 5 , 3 , 9 , 34 ] . In this paper we consider two forms of self-conditioning: one method we call mode “two-pass” is a Markovian form of self-conditioning, it uses only information available at a given time point of a discrete diffusion to estimate the denoiser as in Algorithm 6 in Chen et al. [7 ] . This method can be considered an architecture choice for the denoiser estimator network and is compatible with existing diffusion model error control theory which views the reverse process as a Euler-Maruyama discretization of the continuous SDE with pointwise denoiser estimation.
The second method of self-conditioning we call mode “carry” is when the denoiser reuses its previous estimate at time t − 1 t-1 t − 1 t t t
The denoiser is a conditional expectation on the origin point of a noisy observation given partial information provided by the noised observation,
F ( u , a ) = E π [ X ∣ U = a X + 1 − a 2 Z ] . F(u,a) = E_{\pi}[X \mid U = aX + \sqrt{1-a^2}Z]. F ( u , a ) = E π [ X ∣ U = a X + 1 − a 2 Z ] . If we had some additional side information W W W X X X
F ( u , w , a ) = E π [ X ∣ U = a X + 1 − a 2 Z , W ] . F(u,w,a) = E_{\pi}[X \mid U = aX + \sqrt{1-a^2}Z, W]. F ( u , w , a ) = E π [ X ∣ U = a X + 1 − a 2 Z , W ] . A self-conditioned network takes the form g θ ( u , w , a ) g_{\theta}(u,w,a) g θ ( u , w , a ) ( u , a ) (u,a) ( u , a ) w ∈ [ − 1 , 1 ] d + 1 w \in[-1,1]^{d+1} w ∈ [ − 1 , 1 ] d + 1 x x x w = 0 w=0 w = 0 { − 1 , 1 } d + 1 \{-1,1\}^{d+1} { − 1 , 1 } d + 1 g θ ( u , 0 , a ) g_{\theta}(u,0,a) g θ ( u , 0 , a ) G ( u , a ) G(u,a) G ( u , a ) ( u , a ) (u,a) ( u , a ) ( v ∗ , s ∗ ) (v^*,s^*) ( v ∗ , s ∗ )
In the training algorithm, a self-conditioned network is trained by some portion of the time being fed w = 0 w=0 w = 0 w = g θ ( u , 0 , a ) w=g_{\theta}(u,0,a) w = g θ ( u , 0 , a )
Both self-conditioning modes “two-pass” and “carry” are trained in the same way, and differ in their implementation as a sampling algorithm.
Two-pass is the simpler method, as in sampling it reuses its own estimate with no side information as a second input to a second pass evaluation. This can be viewed as an architecture choice for a larger network that is ultimately still a function of a single pair ( u , a ) (u,a) ( u , a )
g ~ ( u , a ) = g θ ( u , g θ ( u , 0 , a ) , a ) , \widetilde{g}(u,a) = g_{\theta}(u,g_{\theta}(u,0,a),a), g ( u , a ) = g θ ( u , g θ ( u , 0 , a ) , a ) , Using g ~ \widetilde{g} g g g g
The mode carry reuses the previous iteration’s denoiser estimate as its second argument in self-conditioning. Initialize G ^ − 1 = 0 \widehat{G}_{-1}=0 G − 1 = 0 G ^ k \widehat{G}_k G k k k k
G ^ k = g θ ( v k ∗ , G ^ k − 1 , s k ∗ ) , \widehat{G}_k = g_{\theta}(v_k^*,\widehat{G}_{k-1},s_k^*), G k = g θ ( v k ∗ , G k − 1 , s k ∗ ) , X t k + 1 = X t k + δ G ^ k + δ ξ k , X_{t_{k+1}} = X_{t_k} + \delta\widehat{G}_k + \sqrt{\delta}\,\xi_k, X t k + 1 = X t k + δ G k + δ ξ k , This modifies the stochastic localization discrete sampling algorithm as shown in Algorithm 4.
Algorithm 3 Batch SGD step for the self-conditioned tilt–mean network g θ ( u , w , a ) g_\theta(u,w,a) g θ ( u , w , a )
Require: Data { X i } ⊂ { − 1 , 1 } d + 1 \{X_i\} \subset\{-1,1\}^{d+1} { X i } ⊂ { − 1 , 1 } d + 1 d + 1 d+1 d + 1 N N N γ \gamma γ p s c p_{\mathrm{sc}} p sc
L ← 0 \mathcal{L} \leftarrow0 L ← 0
for n = 1 , … , N n=1,\ldots,N n = 1 , … , N do
Draw x ∼ { X i } x \sim\{X_i\} x ∼ { X i } Z ∼ N ( 0 , I d + 1 ) Z \sim\mathcal{N}(0,I_{d+1}) Z ∼ N ( 0 , I d + 1 ) β ∼ LogUnif [ β min , β max ] \beta\sim\operatorname{LogUnif}[\beta_{\min},\beta_{\max}] β ∼ LogUnif [ β m i n , β m a x ] a ← a ( β ) a \leftarrow a(\beta) a ← a ( β )
W.p. p 0 p_0 p 0 q ← 1 d q \leftarrow1_d q ← 1 d k ∼ Unif { 2 , … , d } k \sim\operatorname{Unif}\{2,\ldots,d\} k ∼ Unif { 2 , … , d } q ← d / k 1 S q \leftarrow\sqrt{d/k}\,1_S q ← d / k 1 S k k k S S S
W.p. p λ p_\lambda p λ λ ← 1 \lambda\leftarrow1 λ ← 1 λ ∼ LogUnif [ 1 , d + 1 ] \lambda\sim\operatorname{LogUnif}[1,\sqrt{d+1}] λ ∼ LogUnif [ 1 , d + 1 ]
Λ j , j ← 1 + 1 d − λ 2 d q j , j ≤ d \Lambda_{j,j} \leftarrow\sqrt{1+\frac{1}{d}-\frac{\lambda^2}{d}q_j},\ j \le d Λ j , j ← 1 + d 1 − d λ 2 q j , j ≤ d Λ d + 1 , d + 1 ← λ \Lambda_{d+1,d+1} \leftarrow\lambda Λ d + 1 , d + 1 ← λ
η ← a Λ 2 x + 1 − a 2 Λ Z \eta\leftarrow a\Lambda^2x+\sqrt{1-a^2}\,\Lambda Z η ← a Λ 2 x + 1 − a 2 Λ Z
if q = 1 d q=1_d q = 1 d and λ = 1 \lambda=1 λ = 1 then
( v ∗ , s ∗ ) ← ( η , a ) (v^*,s^*) \leftarrow(\eta,a) ( v ∗ , s ∗ ) ← ( η , a )
else
C ← β ( a ) ∥ η ∥ d + 1 , s ∗ ← 1 2 ( 1 C 2 + 4 − 1 C ) , v ∗ ← d + 1 η / ∥ η ∥ C \leftarrow\beta(a)\frac{\|\eta\|}{\sqrt{d+1}},\quad s^* \leftarrow\frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right),\quad v^* \leftarrow\sqrt{d+1}\,\eta/\|\eta\| C ← β ( a ) d + 1 ∥ η ∥ , s ∗ ← 2 1 ( C 2 1 + 4 − C 1 ) , v ∗ ← d + 1 η /∥ η ∥
W.p. 1 − p s c 1-p_{\mathrm{sc}} 1 − p sc w ← 0 w \leftarrow0 w ← 0 w ← sg [ g θ ( v ∗ , 0 , s ∗ ) ] w \leftarrow\operatorname{sg}[g_\theta(v^*,0,s^*)] w ← sg [ g θ ( v ∗ , 0 , s ∗ )]
L ← L + ∥ x − g θ ( v ∗ , w , s ∗ ) ∥ 2 \mathcal{L} \leftarrow\mathcal{L}+\|x-g_\theta(v^*,w,s^*)\|^2 L ← L + ∥ x − g θ ( v ∗ , w , s ∗ ) ∥ 2
θ ← θ − γ ∇ θ L \theta\leftarrow\theta-\gamma\nabla_\theta\mathcal{L} θ ← θ − γ ∇ θ L
Implementation
Ising Model and 4x4 Sudoku Target distributions for the Ising Model and 4x4 Sudoku model have state spaces of size 2 13 2^{13} 2 13 2 33 2^{33} 2 33
6x6 Sudoku For the 6x6 Sudoku self-conditioning mode “two-pass” was implemented and the sampler tracked performance well across log-odds levels considering the confidence intervals in Figure 18.
9x9 Sudoku The 9x9 Sudoku problem, both modes of self-conditioning were implemented and saw differing results.
Mode two-pass saturated its probability of producing proper puzzles at 60% and did not increase at higher tilt strengths. The model also saw significant differences between the model’s own prediction of if it was producing a reward in its d + 1 d+1 d + 1
Mode carry tracked the log-odds better and saturated at 90% reward-achieving solution, showing less mismatch between its prediction of reward and the actual reward of the token string produced, see Figure 20b.
MNIST Digits For the MNIST digit image generation problem, both modes of self-conditioning were implemented. In this problem mode two-pass follows the log-odds curve better than the carry mode. Mode carry over-samples the reward-achieving solutions at medium strength tilts as seen in Figure 27. As a conjecture of what may be the cause, carry self-conditioning may cause the model to become too confident in its early predictions, continuing to refine upon estimates made early in the process while two-pass is a more faithful implementation of a proper Euler-Maruyama discretization of the continuous-time SDE which defines stochastic localization.
Algorithm 4 Stochastic-localization sampling from the tilted density μ ∝ π e ⟨ ν , ⋅ ⟩ \mu\propto\pi e^{\langle\nu,\cdot\rangle} μ ∝ π e ⟨ ν , ⋅ ⟩
Require: Trained self-conditioned network g θ ( u , w , a ) g_\theta(u,w,a) g θ ( u , w , a ) ν = τ i n f i l l ( e A + − e A − ) + τ r e w a r d e ∗ \nu=\tau_{\mathrm{infill}}(e_{A_+}-e_{A_-})+\tau_{\mathrm{reward}}e_* ν = τ infill ( e A + − e A − ) + τ reward e ∗ 0 = t 0 < t 1 < ⋯ < t K = T 0=t_0<t_1<\cdots<t_K=T 0 = t 0 < t 1 < ⋯ < t K = T
X t 0 ← 0 , G ^ − 1 ← 0 X_{t_0}\leftarrow0,\quad\widehat{G}_{-1}\leftarrow0 X t 0 ← 0 , G − 1 ← 0
for k = 0 , … , K − 1 k=0,\ldots,K-1 k = 0 , … , K − 1 do
δ ← t k + 1 − t k \delta\leftarrow t_{k+1}-t_k δ ← t k + 1 − t k
η ← ν + X t k \eta\leftarrow\nu+X_{t_k} η ← ν + X t k
C ← ∥ η ∥ d + 1 , s ∗ ← 1 2 ( 1 C 2 + 4 − 1 C ) , v ∗ ← d + 1 η / ∥ η ∥ C\leftarrow\frac{\lVert\eta\rVert}{\sqrt{d+1}},\quad s^*\leftarrow\frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right),\quad v^*\leftarrow\sqrt{d+1}\,\eta/\lVert\eta\rVert C ← d + 1 ∥ η ∥ , s ∗ ← 2 1 ( C 2 1 + 4 − C 1 ) , v ∗ ← d + 1 η / ∥ η ∥
G ^ k ← g θ ( v ∗ , G ^ k − 1 , s ∗ ) \widehat{G}_k\leftarrow g_\theta(v^*,\widehat{G}_{k-1},s^*) G k ← g θ ( v ∗ , G k − 1 , s ∗ )
Draw ξ ∼ N ( 0 , I d + 1 ) \xi\sim\mathcal{N}(0,I_{d+1}) ξ ∼ N ( 0 , I d + 1 )
X t k + 1 ← X t k + δ G ^ k + δ ξ X_{t_{k+1}}\leftarrow X_{t_k}+\delta\widehat{G}_k+\sqrt{\delta}\,\xi X t k + 1 ← X t k + δ G k + δ ξ
η ← ν + X T \eta\leftarrow\nu+X_T η ← ν + X T ( v ∗ , s ∗ ) (v^*,s^*) ( v ∗ , s ∗ ) η \eta η
return x ^ ← sign ( g θ ( v ∗ , g θ ( v ∗ , 0 , s ∗ ) , s ∗ ) ) \widehat{x}\leftarrow\operatorname{sign}\left(g_\theta\left(v^*,g_\theta(v^*,0,s^*),s^*\right)\right) x ← sign ( g θ ( v ∗ , g θ ( v ∗ , 0 , s ∗ ) , s ∗ ) )
Discrete Step Size Scheduling
The stochastic localization discretization in Algorithm 2 requires the sequence of discrete time steps { t k } k = 0 K \{t_k\}_{k=0}^{K} { t k } k = 0 K [ 0 , T ] [0,T] [ 0 , T ] K K K ν \nu ν d d d t k t_k t k
The learned denoiser g θ ( v ∗ , s ∗ ) g_\theta(v^*,s^*) g θ ( v ∗ , s ∗ ) t t t s ∗ s^* s ∗ η = X t + ν \eta=X_t+\nu η = X t + ν t k t_k t k s ∗ s^* s ∗
Recall that at time t t t η = ν + X t \eta=\nu+X_t η = ν + X t
C = ∥ η ∥ d + 1 , s ∗ = β − 1 ( C ) = 1 2 ( 1 C 2 + 4 − 1 C ) , β ( s ) = s 1 − s 2 . C=\frac{\lVert\eta\rVert}{\sqrt{d+1}},\qquad s^*=\beta^{-1}(C)=\frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right),\qquad\beta(s)=\frac{s}{1-s^2}. C = d + 1 ∥ η ∥ , s ∗ = β − 1 ( C ) = 2 1 ( C 2 1 + 4 − C 1 ) , β ( s ) = 1 − s 2 s . Since ∥ η ∥ \lVert\eta\rVert ∥ η ∥ X t = t X + t Z X_t=tX+\sqrt{t}Z X t = tX + t Z ∥ X ∥ 2 = d + 1 \lVert X\rVert^2=d+1 ∥ X ∥ 2 = d + 1 E ∥ X t ∥ 2 = ( d + 1 ) ( t 2 + t ) \mathbb{E}\lVert X_t\rVert^2=(d+1)(t^2+t) E ∥ X t ∥ 2 = ( d + 1 ) ( t 2 + t ) E ∥ X t + ν ∥ 2 \mathbb{E}\lVert X_t+\nu\rVert^2 E ∥ X t + ν ∥ 2 t ⟨ ν , Z ⟩ \sqrt{t}\langle\nu,Z\rangle t ⟨ ν , Z ⟩ ⟨ ν , X ⟩ = ∥ ν ∥ \langle\nu,X\rangle=\lVert\nu\rVert ⟨ ν , X ⟩ = ∥ ν ∥
C ‾ ( t ) = ∥ ν ∥ 2 + 2 t ∥ ν ∥ + ( d + 1 ) ( t 2 + t ) d + 1 , s ‾ ( t ) = β − 1 ( C ‾ ( t ) ) . \overline{C}(t)=\sqrt{\frac{\lVert\nu\rVert^2+2t\lVert\nu\rVert+(d+1)(t^2+t)}{d+1}},\qquad\overline{s}(t)=\beta^{-1}(\overline{C}(t)). C ( t ) = d + 1 ∥ ν ∥ 2 + 2 t ∥ ν ∥ + ( d + 1 ) ( t 2 + t ) , s ( t ) = β − 1 ( C ( t )) . this is an approximation but gives a deterministic function of t t t η \eta η d + 1 \sqrt{d+1} d + 1 s ‾ \overline{s} s t t t s k = s ‾ ( 0 ) + k K ( s ‾ ( T ) − s ‾ ( 0 ) ) s_k=\overline{s}(0)+\frac{k}{K}(\overline{s}(T)-\overline{s}(0)) s k = s ( 0 ) + K k ( s ( T ) − s ( 0 )) k = 0 , … , K k=0,\ldots,K k = 0 , … , K t k = s ‾ − 1 ( s k ) t_k=\overline{s}^{-1}(s_k) t k = s − 1 ( s k ) C ‾ ( t ) 2 \overline{C}(t)^2 C ( t ) 2 t t t C k = β ( s k ) C_k=\beta(s_k) C k = β ( s k )
t k = − ( d + 1 + 2 ∥ ν ∥ ) + ( d + 1 + 2 ∥ ν ∥ ) 2 − 4 ( d + 1 ) ( ∥ ν ∥ 2 − ( d + 1 ) C k 2 ) 2 ( d + 1 ) . t_k=\frac{-(d+1+2\lVert\nu\rVert)+\sqrt{(d+1+2\lVert\nu\rVert)^2-4(d+1)(\lVert\nu\rVert^2-(d+1)C_k^2)}}{2(d+1)}. t k = 2 ( d + 1 ) − ( d + 1 + 2 ∥ ν ∥) + ( d + 1 + 2 ∥ ν ∥ ) 2 − 4 ( d + 1 ) (∥ ν ∥ 2 − ( d + 1 ) C k 2 ) . with the endpoints fixed at t 0 = 0 t_0=0 t 0 = 0 t K = T t_K=T t K = T ∥ ν ∥ \lVert\nu\rVert ∥ ν ∥ ν = 0 \nu=0 ν = 0 s ‾ ( 0 ) = 0 \overline{s}(0)=0 s ( 0 ) = 0 t = 0 t = 0 t = 0 T T T T = 20 T = 20 T = 20 K = 1000 K = 1000 K = 1000
Algorithm 5 Time grid equally spaced in the queried signal level s ∗ s^* s ∗
Require: Horizon T T T K K K d d d n = ∥ ν ∥ n = \lVert\nu\rVert n = ∥ ν ∥ n = 0 n = 0 n = 0
D ← d + 1 D \leftarrow d + 1 D ← d + 1
Define C ‾ ( t ) = ( n 2 + 2 t n + D ( t 2 + t ) ) / D \overline{C}(t) = \sqrt{\left(n^2 + 2tn + D(t^2+t)\right)/D} C ( t ) = ( n 2 + 2 t n + D ( t 2 + t ) ) / D ∥ ν + X t ∥ / d + 1 \lVert\nu+ X_t \rVert/\sqrt{d+1} ∥ ν + X t ∥ / d + 1
Define β − 1 ( C ) = 1 2 ( 1 C 2 + 4 − 1 C ) \beta^{-1}(C) = \frac{1}{2}\left(\sqrt{\frac{1}{C^2}+4}-\frac{1}{C}\right) β − 1 ( C ) = 2 1 ( C 2 1 + 4 − C 1 ) C > 0 C > 0 C > 0 β − 1 ( 0 ) = 0 \beta^{-1}(0) = 0 β − 1 ( 0 ) = 0
s l o ← β − 1 ( C ‾ ( 0 ) ) , s h i ← β − 1 ( C ‾ ( T ) ) s_{\mathrm{lo}} \leftarrow\beta^{-1}(\overline{C}(0)),\quad s_{\mathrm{hi}} \leftarrow\beta^{-1}(\overline{C}(T)) s lo ← β − 1 ( C ( 0 )) , s hi ← β − 1 ( C ( T ))
for k = 0 , … , K k = 0,\ldots,K k = 0 , … , K do
s k ← s l o + k K ( s h i − s l o ) s_k \leftarrow s_{\mathrm{lo}} + \frac{k}{K}(s_{\mathrm{hi}}-s_{\mathrm{lo}}) s k ← s lo + K k ( s hi − s lo ) s ∗ s^* s ∗
C k ← s k / ( 1 − s k 2 ) C_k \leftarrow s_k/(1-s_k^2) C k ← s k / ( 1 − s k 2 ) C k = β ( s k ) C_k = \beta(s_k) C k = β ( s k )
t k ← − ( D + 2 n ) + ( D + 2 n ) 2 − 4 D ( n 2 − D C k 2 ) 2 D t_k \leftarrow\frac{-(D+2n)+\sqrt{(D+2n)^2-4D(n^2-DC_k^2)}}{2D} t k ← 2 D − ( D + 2 n ) + ( D + 2 n ) 2 − 4 D ( n 2 − D C k 2 ) C ‾ ( t k ) = C k \overline{C}(t_k)=C_k C ( t k ) = C k
t 0 ← 0 , t K ← T t_0 \leftarrow0,\quad t_K \leftarrow T t 0 ← 0 , t K ← T
return ( t 0 , t 1 , … , t K ) (t_0,t_1,\ldots,t_K) ( t 0 , t 1 , … , t K )
ADDITIONAL EXPERIMENTS INFORMATION
Code for all experiments, including data generation, training, and evaluation is available at the the GitHub link:
https://github.com/CMcDonald-1/internal_guidance_diffusion_models
The datasets and trained checkpoints can be downloaded from the link:
https://osf.io/kewg9/
12-Dimensional Ising Model
Target Density:
The target density is on { − 1 , 1 } d \{-1,1\}^d { − 1 , 1 } d d = 12 d=12 d = 12
π ( x ) ∝ e 1 2 x T J x + ⟨ s , x ⟩ . \pi(x) \propto e^{\frac{1}{2}x^T Jx+\langle s,x\rangle}. π ( x ) ∝ e 2 1 x T J x + ⟨ s , x ⟩ . J J J d × d d \times d d × d J 0 / d J_0/\sqrt{d} J 0 / d J 0 = 1.5 J_0=1.5 J 0 = 1.5 J = J T J=J^T J = J T s s s σ = 0.3 \sigma=0.3 σ = 0.3
The reward is the non-linear indicator function
r ( x ) = { 1 exactly 10 of the 12 coordinates positive − 1 else r(x)=
\begin{cases}
1 & \text{exactly 10 of the 12 coordinates positive}\\
-1 & \text{else}
\end{cases} r ( x ) = { 1 − 1 exactly 10 of the 12 coordinates positive else equivalently r ( x ) = + 1 r(x)=+1 r ( x ) = + 1 ⟨ 1 , x ⟩ = 8 \langle\mathbf{1},x\rangle=8 ⟨ 1 , x ⟩ = 8 { − 1 , 1 } 13 \{-1,1\}^{13} { − 1 , 1 } 13 q τ ∝ π e τ r ( x ) q_\tau\propto\pi e^{\tau r(x)} q τ ∝ π e τ r ( x )
q τ ∝ π ( x , x ∗ ) e τ x ∗ q_\tau\propto\pi(x,x^*)e^{\tau x^*} q τ ∝ π ( x , x ∗ ) e τ x ∗
Training Set:
The training set is 100,000 samples from the base distribution with reward information. 4.41% achieve reward, which is the base probability of reward under the target distribution.
Model Architecture: The operator network is a residual MLP with a mean-field skip connection,
G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) , G_{\theta}(u,a)=\tanh\left(\beta(a)u+g_{\theta}(a)\odot r_{\theta}(u,a)\right), G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) , where β ( a ) u \beta(a)u β ( a ) u [ − 1 , 1 ] 13 [-1,1]^{13} [ − 1 , 1 ] 13 e ( a ) = [ a , log ( 1 + β ( a ) ) , sin ( 2 π a f ) , cos ( 2 π a f ) ] e(a)=[a,\log(1+\beta(a)),\sin(2\pi af),\cos(2\pi af)] e ( a ) = [ a , log ( 1 + β ( a )) , sin ( 2 π a f ) , cos ( 2 π a f )] f ∈ { 1 , 2 , 4 , 8 } f\in\{1,2,4,8\} f ∈ { 1 , 2 , 4 , 8 } u ∈ R 13 u\in\mathbb{R}^{13} u ∈ R 13 1 / d 1/\sqrt{d} 1/ d e ( a ) e(a) e ( a ) ( 1 + g k ) LN ( h ) + b k (1+g_k)\operatorname{LN}(h)+b_k ( 1 + g k ) LN ( h ) + b k ( g k , b k ) (g_k,b_k) ( g k , b k ) e ( a ) e(a) e ( a ) ( 10 → 32 → 13 (10\to32\to13 ( 10 → 32 → 13 g θ ( a ) g_{\theta}(a) g θ ( a )
Training Run: The model was trained for 30,000 Adam steps with batch size 1024, learning rate 10 − 3 10^{-3} 1 0 − 3 β \beta β [ 0.02 , 50 ] [0.02,50] [ 0.02 , 50 ] Λ \Lambda Λ λ ∼ Unif ( 1 , 3.5 ) \lambda\sim\operatorname{Unif}(1,3.5) λ ∼ Unif ( 1 , 3.5 ) Λ \Lambda Λ λ \lambda λ k k k E ∥ X − G θ ( V , a ) ∥ 2 \mathbb{E}\lVert X-G_{\theta}(V,a)\rVert^2 E ∥ X − G θ ( V , a ) ∥ 2
Training ran on a single NVIDIA GeForce RTX 5070 Ti Laptop GPU (Blackwell, 5888 CUDA cores, 12.8 GB, CUDA 12.9) and took 10.7 minutes of wall-clock time (0.18 GPU-hours). Processing 30,000 × 1024 ≈ 3.07 × 10 7 30{,}000\times1024\approx3.07\times10^7 30 , 000 × 1024 ≈ 3.07 × 1 0 7 1.6 × 10 13 1.6\times10^{13} 1.6 × 1 0 13 6 N D 6ND 6 N D
Extended Results For all models in question, τ \tau τ d + 1 \sqrt{d}+1 d + 1 Λ \Lambda Λ
For various signal levels a a a τ \tau τ ( u , a ) (u,a) ( u , a ) Λ \Lambda Λ τ \tau τ
The probability of reward and log-odds plot is seen in Figure 9 across a range of τ \tau τ d + 1 \sqrt{d}+1 d + 1 τ \tau τ
The 12-dimensional Ising model is a small toy model with exact enumeration of the ground truth denoiser to demonstrate highly accurate sampling across tilts and denoiser learning when a ground truth is available.
Figure 7. Ising model covariance matrix ground truth and network sampled.
4x4 Sudoku
Target Density
4 × 4 4 \times4 4 × 4 2 32 2^{32} 2 32 1 1 1 4 4 4
Figure 10. Example valid 4x4 grid.
The base distribution is a weak generator sampling each row independently. With probability 1 − ρ 1-\rho 1 − ρ 1 1 1 4 4 4 ρ \rho ρ ρ = 0.02 \rho= 0.02 ρ = 0.02
A single row is therefore valid with probability ( 1 − ρ ) + ρ ⋅ 24 256 = 98.19 % (1-\rho) + \rho\cdot\frac{24}{256} = 98.19\% ( 1 − ρ ) + ρ ⋅ 256 24 = 98.19%
The reward is the indicator of a valid puzzle appended as a 33rd coordinate, giving the extended state space { − 1 , 1 } 33 \{-1,1\}^{33} { − 1 , 1 } 33
Training Set
Because a proper puzzle arises by chance only 0.081% of the time, the base distribution is mixed with an injection component: each sample is drawn from the uniform distribution over the 288 valid solutions with probability p i n j e c t = 0.05 p_{\mathrm{inject}} = 0.05 p inject = 0.05
Model Architecture
The operator network is the same residual MLP with mean-field skip connection used for the Ising experiment,
G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) G_\theta(u,a) = \tanh\left(\beta(a)u + g_\theta(a) \odot r_\theta(u,a)\right) G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) scaled up for the larger state space. The signal level enters through a 10-dimensional level embedding e ( a ) = [ a , log ( 1 + β ( a ) ) , sin ( 2 π a f ) , cos ( 2 π a f ) ] e(a) = [a,\ \log(1+\beta(a)),\ \sin(2\pi af),\ \cos(2\pi af)] e ( a ) = [ a , log ( 1 + β ( a )) , sin ( 2 π a f ) , cos ( 2 π a f )] f ∈ { 1 , 2 , 4 , 8 } f \in\{1, 2, 4, 8\} f ∈ { 1 , 2 , 4 , 8 } u ∈ R 33 u \in\mathbb{R}^{33} u ∈ R 33 1 / d 1/\sqrt{d} 1/ d e ( a ) e(a) e ( a ) ( 1 + g k ) LN ( h ) + b k (1 + g_k)\operatorname{LN}(h) + b_k ( 1 + g k ) LN ( h ) + b k ( g k , b k ) (g_k,b_k) ( g k , b k ) e ( a ) e(a) e ( a ) g θ ( a ) g_\theta(a) g θ ( a )
Training Run
The model was trained for 30,000 Adam steps with batch size 8192, learning rate 10 − 3 10^{-3} 1 0 − 3 β \beta β [ 0.02 , 50 ] [0.02, 50] [ 0.02 , 50 ] k k k { 2 , … , 32 } \{2,\ldots,32\} { 2 , … , 32 } k k k 32 / k \sqrt{32/k} 32/ k λ \lambda λ [ 1 , 5 ] [1,5] [ 1 , 5 ] ∑ i Λ i 2 = d \sum_i \Lambda_i^2 = d ∑ i Λ i 2 = d E ∥ X − G θ ( V , a ) ∥ 2 \mathbb{E}\lVert X - G_\theta(V,a)\rVert^2 E ∥ X − G θ ( V , a ) ∥ 2
Training ran on a single NVIDIA GeForce RTX 5070 Ti Laptop GPU (Blackwell, 5888 CUDA cores, 12.8 GB, CUDA 12.9) and took 67.3 minutes of wall-clock time (1.12 GPU-hours). Processing 30,000 × 8192 ≈ 2.46 × 10 8 2.46 \times10^8 2.46 × 1 0 8 2.5 × 10 16 2.5 \times10^{16} 2.5 × 1 0 16 6 N D 6ND 6 N D
Extended Results
For the 4×4 Sudoku problem, the reward set is only 288 valid grids, and the weak generator decomposes into a product measure across rows. As such, the ground truth denoiser can be computed via enumeration at any reward tilt strength. The estimated L 2 L^2 L 2 Figure 11 . The denoiser error does increase with τ \tau τ
We show the actual 33 values of the base denoiser E [ X i ∣ u , a ] \mathbb{E}[X_i|u,a] E [ X i ∣ u , a ] u u u τ = 0 \tau= 0 τ = 0 τ = 4 \tau= 4 τ = 4 Figure 14 . Expected values are all between −1 and 1. Note the τ = 0 \tau= 0 τ = 0 τ = 4 \tau= 4 τ = 4
Figure 14. Example ground truth denoiser and network denoiser in 33 coordinates at sampled points. Samples in a) are untilted sampled and b) tilted samples.
The untitled τ = 0 \tau= 0 τ = 0 τ = 4 \tau= 4 τ = 4 Figure 11 .
For the 4 × 4 4 \times4 4 × 4
To track the log-odds in Figure 12 , we compute the τ \tau τ τ \tau τ d + 1 = 5.74 \sqrt{d+1} = 5.74 d + 1 = 5.74
The localization samples using the ground truth denoiser tracks the log-odds across all levels, even past 98% when considering the confidence intervals from 2000 samples. The network based sampler tracks well up to the 98% level, but then does strictly under-sample past this high bar with the confidence intervals not overlapping the theoretical log-odds. Note when the network based sampler fails to exactly track the log-odds at this high tilt it is still generating over 98% reward-achieving puzzles, but still under-sampling its theoretical value at this tilt.
As an added bonus to the stochastic localization path, the denoiser E [ X i ∣ X t ] E[X_i|X_t] E [ X i ∣ X t ] ∣ E [ X i ∣ X t ] ∣ > ℓ |E[X_i|X_t]| > \ell ∣ E [ X i ∣ X t ] ∣ > ℓ
One of the arguments against continuous diffusion for discrete targets is that the trajectory remains a continuous value that cannot be interpreted in discrete token space until the terminal decoding step, unlike masked diffusion or autoregressive models which build partial discrete token strings along their trajectory.
By using the confidence based bit decoding, we can achieve a similar effect to see what discrete set of tokens the puzzle was considering at each moment in the trajectory and soft decode at any time.
When both bits in a cell are fixed we record the cell value, when one bit is fixed such as ( ∗ , − 1 ) (\ast, -1) ( ∗ , − 1 ) ( 1 , − 1 ) (1, -1) ( 1 , − 1 ) ( − 1 − 1 ) (-1 -1) ( − 1 − 1 )
We note the network is not “solving” the puzzle “logically” as a human would by applying row, column, and box constraints in sequence. For example, in the second frozen grid in Figure 13 , there is a 1 fixed in the bottom left box. This would rule out all 1’s in that column, but the upper box cells in the same column have not yet ruled out a 1 in that column since the bits for those two cells have not yet reached their own confidence level. This is therefore a sequence of most likely puzzles tracking the most confident bits at each moment in the trajectory.
We call this path a “fuzzy trajectory” tracking the denoiser coordinate values across the signal levels, noting when a coordinate achieves a confidence level more than 0.75 in magnitude, and recording events when a cell gets fixed in Figure 13 . Note some cells do become confident but ultimately flip their token value by the end of the run. Note the row 2 column 4 cell flips from a 2 to a 4 by the end of sampling. We run our stochastic localization up to a T T T T = 3 T = 3 T = 3
Note there are 288 valid 4 × 4 4 \times4 4 × 4
We generate 3000 valid puzzles with the stochastic localization path using the network denoiser and the ground truth and perform a chi-squared test for uniformity. The network based sampler visits all puzzles except for 1, and its maximum visit to one particular puzzle is 25. The ground truth denoiser visits every solution, with minimum visits being 4 and maximum being 21. Visually, they are not that distinct in their visit histogram in Figure 15 . However, the ground truth sampler passes a chi-squared test for uniformity, while the network based sampler fails the test. As such, it has learned some preference for certain solutions, although it is not a strong preference.
6x6 Sudoku
Target Density
6 × 6 Sudoku has 6 rows, 6 columns, and six non-overlapping 2 × 3 2 \times3 2 × 3 36 × 3 = 108 36 \times3 = 108 36 × 3 = 108 2 108 2^{108} 2 108 4 × 4 4 \times4 4 × 4 dead code words : three bits admit eight values, of which only six are legal, so two codes per cell decode to the invalid values 7 and 8. The generator never emits them, but the model operates on the full 2 108 2^{108} 2 108
Figure 16. Example valid 6x6 grid.
There are 28,200,960 proper puzzles, vastly more than the 4x4 puzzle. The training set does not contain all proper puzzles, thus the network’s generations should be novel making new proper puzzles not seen in the training set.
The base distribution is again a weak generator sampling each row independently. With probability 1 − ρ 1-\rho 1 − ρ ρ \rho ρ ρ = 0.02 \rho= 0.02 ρ = 0.02 ( 1 − ρ ) + ρ ⋅ 720 6 6 = 98.03 % (1-\rho) + \rho\cdot\frac{720}{6^6} = 98.03\% ( 1 − ρ ) + ρ ⋅ 6 6 720 = 98.03%
Because the generator enforces no column or box constraints, and the grid is much larger than in the 4 × 4 4 \times4 4 × 4 1.8 × 10 − 10 1.8 \times10^{-10} 1.8 × 1 0 − 10
Training Set
Since the weak generator essentially never produces a proper puzzle on its own, reward-achieving examples must be injected: each sample is drawn uniformly from the set of valid solutions with probability p i n j e c t = 0.05 p_{\mathrm{inject}} = 0.05 p inject = 0.05 4 × 4 4 \times4 4 × 4
Model Architecture
The target distributions and models were trained in order from smallest to largest state space starting with the Ising Model, then 4x4, 6x6, 9x9 Sudoku, then the MNIST problem. As the space scaled up in size, larger data sets were needed to learn the more complicated density and tilt, and more expressive model architectures.
The previous Ising and 4x4 Sudoku problems used MLP models with no attention mechanism. For 6x6, 9x9 Sudoku and MNIST we use a transformer architecture with attention.
This problem uses a transformer featurizer rather than the residual MLP, since the row, column and box constraints are joint combinatorial structure that a flat trunk must discover entirely from data, whereas attention can relate cells directly. The 108 core bits are split into 6 tokens of 18 bits (one per row), and the reward coordinate is given its own token rather than folded into the level conditioning, so attention can relate x ∗ x_* x ∗ d m o d e l = 256 d_{\mathrm{model}} = 256 d model = 256 e ( a ) e(a) e ( a )
The output head is unchanged, preserving the mean-field skip
G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) , G_\theta(u,a)=\tanh\bigl(\beta(a)u+g_\theta(a)\odot r_\theta(u,a)\bigr), G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) , with every residual branch zero-initialized so that at initialization the network is exactly the mean-field denoiser. The network is additionally self-conditioned, as discussed in Appendix § 9 .§ 4 , the previous estimate of x 0 x_0 x 0 4 × 4 4 \times4 4 × 4
Training Run
The model was trained for 180,000 Adam steps with batch size 8192, learning rate 10 − 3 10^{-3} 1 0 − 3 β \beta β β \beta β k k k { 2 , … , 108 } \{2,\ldots,108\} { 2 , … , 108 } k k k 108 / k \sqrt{108/k} 108/ k λ \lambda λ [1 , 10 ] . The objective is the unweighted per-coordinate squared error summed over the 109 coordinates; the final training loss was 53.92.
Training ran on a single NVIDIA H100 80GB HBM3 GPU (Hopper, CUDA 12.6) on the PSC Bridges-2 cluster and took 3.70 hours. Processing 180,000 × 8192 ≈ 1.47 × 10 9 \approx1.47 \times10^9 ≈ 1.47 × 1 0 9 2.9 × 10 16 2.9 \times10^{16} 2.9 × 1 0 16 6 N D 6ND 6 N D
Extended Results
For the 6x6 Sudoku problem, the weak generator samples each row independently thus the ground truth denoiser expectation can be enumerated for the untilted density. With a much larger set of valid puzzles we do not enumerate the denoiser with reward tilting applied. The L 2 L^2 L 2 τ = 0 \tau=0 τ = 0 Figure 17 .
Figure 17. L 2 L^2 L 2
For the log-odds ratio, the τ \tau τ d \sqrt{d} d
At τ 98 = 3.42 \tau_{98}=3.42 τ 98 = 3.42
Figure 18. 6x6 Sudoku log-odds across tilt strengths.
valid 6 × 6 6 \times6 6 × 6 p = 0.29 p = 0.29 p = 0.29 ≥ 10 \ge10 ≥ 10
9x9 Sudoku
Target Density
Full-size 9 × 9 9 \times9 9 × 9 3 × 3 3 \times3 3 × 3 { 1 , … , 9 } \{1,\ldots,9\} { 1 , … , 9 } ≈ 6.67 × 10 21 \approx6.67 \times10^{21} ≈ 6.67 × 1 0 21
Figure 19. Example valid 9x9 grid.
Each cell is a 4-bit encoding of the numbers 1–9, giving 81 × 4 = 324 81 \times4 = 324 81 × 4 = 324 2 324 2^{324} 2 324 6 × 6 6 \times6 6 × 6
The base distribution is again a weak generator sampling each row independently: with probability 1 − ρ 1-\rho 1 − ρ ρ \rho ρ ρ = 0.02 \rho= 0.02 ρ = 0.02 5.1 × 10 − 29 5.1 \times10^{-29} 5.1 × 1 0 − 29
The reward is the validity indicator, appended as a 325th coordinate to give the extended state space { − 1 , 1 } 325 \{-1,1\}^{325} { − 1 , 1 } 325
Training Set
The training set consists of 10,000,000 samples, drawn from the weak generator with probability 1 − p i n j e c t 1-p_{\mathrm{inject}} 1 − p inject p i n j e c t = 0.05 p_{\mathrm{inject}} = 0.05 p inject = 0.05 6 × 6 6 \times6 6 × 6 6 × 6 6 \times6 6 × 6 1 , … , 6 1,\ldots,6 1 , … , 6 9 × 9 9 \times9 9 × 9 1.84 × 10 16 1.84 \times10^{16} 1.84 × 1 0 16 ∣ G ∣ = 1 , 218 , 998 , 108 , 160 |G| = 1,218,998,108,160 ∣ G ∣ = 1 , 218 , 998 , 108 , 160
The realized training set contains 500,883 valid grids (5.01%), all of them distinct, so each is seen roughly 147 times over the training run. Proper sampling should essentially have 0 probability of recreating a training example, thus the model should generate 100% novel valid puzzles and never recreate a valid puzzle seen in the training set.
Model Architecture
The architecture is the same self-conditioned transformer used at 6 × 6 6 \times6 6 × 6 d m o d e l = 512 d_{\mathrm{model}} = 512 d model = 512
G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) , G_\theta(u,a) = \tanh\left(\beta(a)u + g_\theta(a) \odot r_\theta(u,a)\right), G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) , so that at initialization the network is exactly the mean-field denoiser. Self-conditioning concatenates the previous x 0 x_0 x 0
Training Run
The model was trained for 180,000 Adam steps with batch size 8192, learning rate 10 − 3 10^{-3} 1 0 − 3 [ 0.02 , 50 ] [0.02, 50] [ 0.02 , 50 ] [ 0.5 , 2.5 ] [0.5, 2.5] [ 0.5 , 2.5 ] k k k { 2 , … , 324 } \{2,\ldots,324\} { 2 , … , 324 } k k k 324 / k \sqrt{324/k} 324/ k λ \lambda λ [ 1 , 17 ] [1,17] [ 1 , 17 ]
Training ran on a single NVIDIA H100 80GB HBM3 GPU (Hopper, CUDA 12.6) on the PSC Bridges-2 cluster and took 22.5 hours. Processing 180,000 × 8192 ≈ 1.47 × 10 9 180{,}000 \times8192 \approx1.47 \times10^9 180 , 000 × 8192 ≈ 1.47 × 1 0 9 2.3 × 10 17 2.3 \times10^{17} 2.3 × 1 0 17 6 N D 6ND 6 N D
Extended Results
For the 9 × 9 9 \times9 9 × 9
For the 9 × 9 9 \times9 9 × 9
G ~ t = G ( X t , G ( X t , 0 , a ) , a ) . \widetilde{G}_t = G(X_t, G(X_t, 0, a), a). G t = G ( X t , G ( X t , 0 , a ) , a ) . Mode “carry” uses the previous time point’s estimate as input, which breaks the Markov structure and the sampling path is no longer a pure Euler-Maruyama discretization of the continuous SDE with pointwise denoiser estimation,
G ~ t = G ( X t , G ~ t − 1 , a ) . \widetilde{G}_t = G(X_t, \widetilde{G}_{t-1}, a). G t = G ( X t , G t − 1 , a ) . For the 9 × 9 9 \times9 9 × 9
Two-pass fails to properly track the log-odds. It plateaus around 60% reward-achieving solutions, and is already lagging behind at lower tilt strengths. Note the model has its own prediction of if it thinks it is producing a reward-achieving solution in its d + 1 d+1 d + 1
Figure 20. (a) Self-conditioning mode two-pass. (b) Self-conditioning mode carry. 9x9 Sudoku log-odds across tilt strengths for both self-conditioning modes.
The carry mode does much better at tracking the log-odds, plateauing at 90% reward-achieving. At higher tilt strengths, it also experiences some discrepancy between the solutions it thinks are reward-achieving via its d + 1 d + 1 d + 1
To consider why this may be the case, it is possible the few missing percent of samples that were “supposed” to terminate at a valid puzzle mis-sampled a few bits due to errors in the learned denoiser accumulating over the sample path, which would get a -1 reward from the sharp correct/incorrect puzzle reward function for any cell violations.
Using a CP-SAT optimization problem we can find for each generated puzzle the nearest valid solution in Hamming distance. We plot the CDF of how many produced puzzles in 300 samples are valid (distance 0), 1 bit away from correct, 2, and so on from a valid puzzle in Figure 21 . We see an example puzzle in Figure 22 that was generated by two-pass self-conditioning at a tilt strength of 3 and is 8 bit flips away from a valid puzzle. We see in Figure 21 for the two-pass mode, while only 60% of puzzles are valid, 90% are within 10 bits out of 324 from a valid puzzle. Two-pass and carry have comparable proportions of puzzles produced within 10 bits of a valid puzzle.
Figure 21. Cumulative distribution function of Hamming distance to nearest valid puzzle for 300 samples.
Figure 22. Generated invalid puzzle (right) with constraint violations and nearest valid puzzle (left) in Hamming distance. Nearest valid grid would require 8 bit flips to the binary encoding. The bit flips are not all restricted to the invalid bits of the produced puzzle, but previously valid bits may be flipped to get to a valid puzzle.
Here we summarize the infilling performance of the model, using an infill strength and a reward tilt strength of τ = 3 \tau= 3 τ = 3
Table 1. Sudoku 9 × 9 infilling as a function of the number of given cells. Each row is based on 100 uniformly random valid solutions × 20 draws per puzzle (n = 2000 n = 2000 n = 2000 τ = 3 \tau= 3 τ = 3
A natural concern for a sampler that reaches high validity is that it has simply memorized its training set. We test this on Sudoku 9 × 9 with the self-conditioned transformer (carry, τ = 3 \tau= 3 τ = 3 N = 1000 N = 1000 N = 1000 6.67 × 10 21 6.67 \times10^{21} 6.67 × 1 0 21 Figure 23 ). The two distributions are almost identical. Generated grids are therefore no closer to the training data than new samples from the data distribution would be, and the closest generation (39 cells away) still differs from every reference grid in almost half its cells. The model has learned the Sudoku constraints rather than copying its training grids.
Figure 23. CDF and density of nearest distance reward-achieving gird in training set for network generated samples and held out set. Nearly identical distributions, so generated puzzles no more similar to training set than a held out portion. No puzzles are exact match, closet puzzle differ in 39 cells
MNIST Conditional Image Generation
Target Density
MNIST consists of 28 × 28 grayscale images of handwritten digits 0–9, with pixel intensities in { 0 , … , 255 } \{0,\ldots,255\} { 0 , … , 255 } 784 × 8 = 6,272 784 \times8 = 6{,}272 784 × 8 = 6 , 272 2 6272 2^{6272} 2 6272
Since there are 10 digit classes, the reward is a 10-bit block rather than a single coordinate. The class label c c c r c ∈ { − 1 , 1 } 10 \mathbf{r}_c \in\{-1,1\}^{10} r c ∈ { − 1 , 1 } 10 c c c { − 1 , 1 } 6282 \{-1,1\}^{6282} { − 1 , 1 } 6282 c c c h = τ r c \mathbf{h} = \tau\mathbf{r}_c h = τ r c e c \mathbf{e}_c e c
logit π τ ( c ) = logit π ( c ) + 4 τ , \operatorname{logit}\pi_{\tau}(c) = \operatorname{logit}\pi(c) + 4\tau, logit π τ ( c ) = logit π ( c ) + 4 τ , where π ( c ) \pi(c) π ( c ) c c c
Training Set
The training set is the standard 60,000 MNIST training images and their labels, each encoded as above. The official 10,000-image test split is held out and never used for training. It is used for three things: (i) a fixed batch of 4,096 held-out images, noised with the training recipe, is scored every 1,000 steps to monitor overfitting; (ii) it provides the real-digit comparison in Figure 28 ; and (iii) it is used to measure the error rates of the classifier that scores generated samples (see Extended Results).
(a) 25 class-0 images from the held-out set.
(b) 25 generated class-0 images, carry (non-Markovian) self-conditioning.
(c) 25 generated class-0 images, two-pass (Markovian) self-conditioning.
Figure 28. 25 class-0 images from the held-out set (a), and network samples under both self-conditioning modes with the same random seed (b, c). Panel titles give the classifier’s label (green if it matches the requested class, red if not). The samples are similar in quality and variety to the held-out digits, with no sign of mode collapse. τ = 1.289 \tau= 1.289 τ = 1.289
Model Architecture
We use the same self-conditioned transformer as for 6 × 6 and 9 × 9 Sudoku, with tokens built from image patches. The image is divided into a 7 × 7 grid of non-overlapping 4 × 4 pixel patches. Each patch becomes one token of 16 × 8 = 128 16 \times8 = 128 16 × 8 = 128 d m o d e l = 256 d_{\mathrm{model}} = 256 d model = 256 e ( a ) e(a) e ( a )
G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a ) ) . G_\theta(u,a) = \tanh(\beta(a)u + g_\theta(a) \odot r_\theta(u,a)). G θ ( u , a ) = tanh ( β ( a ) u + g θ ( a ) ⊙ r θ ( u , a )) . Self-conditioning concatenates the previous x 0 x_0 x 0
Training Run
The model was trained for 30,000 Adam steps with batch size 2048, learning rate 10 − 3 10^{-3} 1 0 − 3 β ∈ [ 0.02 , 50 ] \beta\in[0.02, 50] β ∈ [ 0.02 , 50 ] Λ \Lambda Λ λ \lambda λ [ 1 , 24 ] [1, 24] [ 1 , 24 ] λ < d / 10 = 25.06 \lambda< \sqrt{d/10} = 25.06 λ < d /10 = 25.06 d \sqrt{d} d
30,000 × 2048 ≈ 6.1 × 10 7 30{,}000 \times2048 \approx6.1 \times10^7 30 , 000 × 2048 ≈ 6.1 × 1 0 7
Training ran on a single NVIDIA H100 80GB GPU on the PSC Bridges-2 cluster for 1 hour and 15 minutes. Processing 6.1 × 10 7 6.1 \times10^7 6.1 × 1 0 7 1.3 × 10 15 1.3 \times10^{15} 1.3 × 1 0 15 6 N D 6ND 6 N D
Extended Results
To detect if a generated image is of a given class, we train an independent score CNN classifier. The scorer is a two-layer CNN (1.2M parameters) trained for 3 epochs with Adam on the 60,000 MNIST training images, reaching 98.74% accuracy on the held-out test split.
The classifier is an imperfect instrument and cannot detect very strong tilt accuracies beyond its own classification error. Its error is also not constant across classes. As such, our ability to detect tilt strength saturates at the classifier accuracy.
Near the top of the tilt range its error is the dominant source of uncertainty. Let s c s_c s c c c c c c c e c e_c e c c c c c c c p p p
m = p s c + ( 1 − p ) e c . m = p s_c + (1-p)e_c. m = p s c + ( 1 − p ) e c . We measure s c s_c s c e c e_c e c s c s_c s c e c ≤ 0.0023 e_c \le0.0023 e c ≤ 0.0023 s c s_c s c τ ≥ 1.69 \tau\ge1.69 τ ≥ 1.69
We report two rates from the same samples: the classifier-scored rate, and the class the network itself writes in its 10-bit reward block (argmax of the block). The second is read directly from the sample and involves no classifier. The τ \tau τ [ 0 , τ 95 ] [0,\tau_{95}] [ 0 , τ 95 ] τ 95 = 1.285 \tau_{95} = 1.285 τ 95 = 1.285 ( τ 95 , 2.5 ] (\tau_{95}, 2.5] ( τ 95 , 2.5 ] τ \tau τ τ \tau τ Table 2 and Table 3 give the pooled results, and Figure 24 –Figure 27 show each class separately.
Table 2: MNIST pooled class-conditional log-odds, two-pass self-conditioning, n = 2 , 560 n = 2,560 n = 2 , 560 τ \tau τ
Table 2.
Table 3: As Table 2 , for carry self-conditioning, with the same τ \tau τ
Table 3.
Figure 24. MNIST class-conditional log-odds against tilt strength τ \tau τ
Figure 27. MNIST log-odds pooled over all ten classes, carry self-conditioning (Table 3 ).
The self-conditioning modes two-pass (Markovian) and carry (non-Markovian) behave in the opposite way to 9 × 9 9 \times9 9 × 9
With two-pass, the classifier-scored rate matches the corrected exact rate within its confidence interval, or nearly so, from τ = 0 \tau= 0 τ = 0 τ ≈ 1 \tau\approx1 τ ≈ 1 overshoots over the whole rising region. At τ = 0.51 \tau= 0.51 τ = 0.51 τ \tau τ
Figure 25. MNIST log-odds pooled over all ten classes, two-pass self-conditioning (Table 2 ).
Figure 26. As Figure 24 , for carry self-conditioning. The overshoot over 0.26 ≤ τ ≤ 1.03 0.26 \le\tau\le1.03 0.26 ≤ τ ≤ 1.03
The log-odds check does not consider the sample quality or variety of digits produced. We sample 25 class 0 images from the hold out set, 25 samples images at the carry mode, and 25 samples at two-pass mode at a tilt τ 0.95 \tau_{0.95} τ 0.95 Figure 28 . Digits generated by the network show the same variety and quality as true hold out images. Note one of the generated samples was mis-sampled as a class 8 image, but this is to be expected with a 5% non-reward rate at this tilt strength. Shared seed and randomness was used in the two sampling runs which results in similarity in the images generated in pane b) and c).
Figure 29. Example digit generations tilted towards each class. Leftmost image per digit is a sample of 25 images from hold out set network was not trained on, second image is from sample mode carry and third is sample mode two-pass. Middle and right grids share the same random seed in their sampling, resulting in often similar samples which is not a form of memorization. Tilt strength is calibrated to 95% reward-achieving in log-odds theory, thus it is reasonable to see a few mis-samples in repeated 25 size sets.
