Blink of an eye: a simple theory for feature localization in generative models
Authors: Marvin Li
Aayush Karan
Sitan Chen
Venue: ICML 2025 oral (2025)
Keywords: stochastic localization
theory of diffusion
large language models
interpretability
reasoning
Primary area: theory->probabilistic_methods
TL;DR: A simple, general, and unifying theory for feature localization in language and diffusion models

Abstract: Large language models can exhibit unexpected behavior in the blink of an eye. In a recent computer use demo, a language model switched from coding to Googling pictures of Yellowstone, and these sudden shifts in behavior have also been observed in reasoning patterns and jailbreaks. This phenomenon is not unique to autoregressive models: in diffusion models, key features of the final output are decided in narrow ``critical windows'' of the generation process. In this work we develop a simple, unifying theory to explain this phenomenon. Using the formalism of stochastic localization for generative models, we show that it emerges generically as the generation process localizes to a sub-population of the distribution it models.  While critical windows have been studied at length in diffusion models, existing theory heavily relies on strong distributional assumptions and the particulars of Gaussian diffusion. In contrast to existing work our theory (1) applies to autoregressive and diffusion models; (2) makes very few distributional assumptions; (3) quantitatively improves previous bounds even when specialized to diffusions; and (4) requires basic mathematical tools. Finally, we validate our predictions empirically for LLMs and find that critical windows often coincide with failures in problem solving for various math and reasoning benchmarks.

Blink of an eye: a simple theory for feature localization in generative models
Marvin Li
Aayush Karan 2
Sitan Chen
Abstract
Reasoning
Jailbreaks
Diffusion
Large language models can exhibit unexpected behavior in the blink of an eye. In a recent com- puter demo, language model switched from coding to Googling pictures of Yellowstone, and these sudden shifts in behavior have also been ob- served in reasoning patterns and jailbreaks. This phenomenon is not unique to autoregressive mod- els: in diffusion models, key features of the final output are decided in narrow "critical windows" of the generation process. In this work we de- velop a simple, unifying theory to explain this phenomenon. Using the formalism of stochastic localization for generative models, we show that it emerges generically as the generation process localizes to a sub- sub-population of the distribution it models. While critical windows have been studied at length in diffusion models, existing theory heav- ily relies on strong distributional assumptions and the particulars of Gaussian diffusion. In contrast to existing work our theory (1) applies to autore- gressive and diffusion models; (2) makes very few distributional assumptions; (3) quantitatively improves previous bounds even when specialized to diffusions; and (4) requires basic mathematical tools. Finally, validate predictions empir- ically for LLMs and find that critical windows often coincide with failures in problem solving for various math and reasoning benchmarks.
Picorrdd
0.5
Token position Abdin et. 2024
Hatza Labs 2024
Georgiev et 2023
Figure Examples of critical windows different data modal- ities and samplers, including reasoning (Abdin et al., 2024; Qi et al., 2024) and certain jailbreaks (Haize Labs, 2024b) for lan- guage models and image class for diffusion models (Georgiev et al., 2023).
The Phi-4 team reported that the probability of correctly answering a math problem can plummet with a single token (Abdin et al., 2024; Lin et al., 2024).
Gemini abruptly threatened a student who was using it to study (Gemini, 2024).
Llama models can be jailbroken by manipulating the first handful of generated tokens (Qi al., 2024; Haize Labs, 2024b).
These abrupt shifts are not unique to autoregressive mod- els. In diffusion models, it has been observed that certain properties like the presence of an object in the background or the image class emerge in narrow time intervals, some- times called critical windows, of the generation process (Ho et al., 2020; Meng et al., 2022; Choi et al., 2022; Raya Ambrogioni, 2023; Georgiev al., 2023; Sclocchi al., 2024; 2025; Biroli al., 2024; Li & Chen, 2024).
1. Introduction
In large language models (LLMs), undesirable behavior can often emerge very suddenly. For example,
Critical windows, more broadly characterizable as a few steps of the sampling procedure during which features of the final output appear, arise many different generative models and data modalities (Figure 1). They are extremely useful from an interpretability perspective as they represent the steps of the sampler responsible for given property of the output (Georgiev et al., 2023; Qi et al., 2024), and have also been used to provide richer stepwise rewards for preference optimization and finetuning (Abdin et al., 2024; Lin et al., 2024; Qi et al., 2024). As the applications of gen- erative models proliferate, it is crucial from interpretability. safety, and capability perspectives to understand how and
Claude transitioned from coding to browsing pictures of Yellowstone while using a computer (Anthropic, 2024).
'Harvard College, Cambridge, MA, USA JJhn A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA, USA. Correspondence to: Marvin Li <mar- cali@coollgecharvaaa
Proceedings of the 4dnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Blink of an eye: a simple theory for feature localization in generative models
why these critical windows emerge.
beyond the Markovianity of the observation process. Below we highlight our main contributions:
Recently, this phenomenon has received significant attention within the theoretical literature on diffusion models (Raya & Ambrogioni, 2023: Sclocchi et al., 2024; 2025; Biroli et al., 2024; Li & Chen, 2024). While existing works do offer pre- dictive theory in the diffusions setting, they either (A) make strong distributional assumptions or (B) rely heavily on the particulars of diffusion, which do not straightforwarddy extend to autoregressive models.
1. Generality: In comparison to existing work, our theory (Theorem 3.1) makes very few distributional assump- tions and requires no statistical physics or stochastic calculus machinery, relying only on simple mateleaaaaaaai mathemat- ical tools. The simplicity and generality of our theory yield concrete improvements to rigorously characterizing the location of critical windows for well-studied models of data. For example, in contrast (Li & Chen, 2024), who could only identify an analogue to Theorem 3.1 with bounds that grow at a polynomial factor with di- mension, our improved theory allows us to obviate this dimensional dependence and replace it with a constant.
Works in the former category carry out non-rigorous statis- tical physics calculations tailored to specific toy models of data like mixtures of Gaussians or context-free grammars with random production rules (Sclocchi et al., 2025; 2024; Biroli et al., 2024; Raya & Ambrogioni, 2023). Works in the latter category derive rigorous bounds in settings without explicit parametric structure, e.g. mixtures of strongly log- concave distributions (Li & Chen, 2024), but they rely on tools like Girsanov's ' theorem which are specific to Gaussian diffusion. Additionally, the bounds in the latter are generally cruder, losing dimension-dependert factors. We ask:
2. Diverse instantiations: Another blessing of the gener- ality of our framework is that we can characterize the location of critical windows for many different models of data and generative models, while previous works were restricted one particular form of data or model. To illustrate the flexibility of our bounds, we explicitly com- pute the locations and widths of these windows for dif- ferent generative models and data modalities (Section 4), One such example we provide elucidates a new connec- tion between critical windows for in-context learning and the all-or-nothing phenomenon in statistical inference.
Is there a simple, general theory that can explain critical windows across all generative modeling paradigms and data modalities?
1.1. Our contributions
3. Insights into hierarchical data: We instantiate our bounds for hierarchically structured models of data, sig- nificantly generalizing results of (Li & Chen, 2024) which only applied diffusions and Gaussian mixtures (Section 5). This allows to show that the hierarchy for generative model may resemble the hierarchy of the true data generating process if both come from the same kind of sampler, but in general may differ. We also use our theory to argue that autoregressive models can support deeper hierarchies than diffusion models.
In this work, we develop simple theoretical framework that characterizes critical windows in both diffusion models and autoregressive models. Our theory is fully rigorous and show that such windows arise generically when the model localizes from a larger sub-population to a smaller one, distilling the phenomenon of critical windows very general facts about the data distribution. In particular, we show that the critical window is parameterized by the last step of the generative model at which the larger and smaller sub-populations are indistinguishable and the first step of the generative model at which that smaller and larger sub- populations have negligible overlap between each other. We will formalize these notions in Section 3.
4. Experimental results: Finally, we show that our the- ory makes accurate predictions for the location of crit- ical windows for LLMs in toy synthetic setting. We also empirically demonstrate critical windows for gener- ations from LLAMA-3.1 1--8-BBB---B----- Instruct, 10l-3-7B-Insttree and Qwen-2.5-778 2----.--------trree on 7 different math and rea- soning benchmarks. Concurrently with (Abdin et al., 2024; Lin et al., 2024), we observe that critical windows occur during important mistakes in the reasoning patterns of I.LMs.
The key insight of our theory is to apply to the formalism of stochastic localization samplers (see Section 2.1 for a for- mal description) to understand critical windows (Montanari, 2023b; Chen & Eldan, 2022). Roughly speaking, a stochas- tic localization scheme is any generative model given by a time-reversal of a Markovian degradation process which takes a sample from the target distribution and generates progressively less informative "observations" of it. In diffu- sion models, the degradation is a convolution of the original sample with larger and larger amounts of Gaussian noise. In autoregressive models, the degradation is the masking of entries from right to left. Importantly, our theory does not use anything about the specific structure of the sampler
Our theory provides valuable insights for practitioners. For instance, in Example 4.3 we provide model for critical win- dows in jailbreaks and the Yellowstone example (A (Anthropic, 2024; Qi et al., 2024). and argue that training on correc-
'The code to reproduce the experiments can be found at
https: / /gttbbb.tttut/gg///g coom/mrrvini critical aidcal-windii wdows- 1m.
2

---

Blink of an eye: a simple theory for feature localization in generative models
tions from critical windows can enable models to recover from these 'bad' modes of behavior. This provides rigorous theoretical justification for (Qi et al., 2024)'s approach for deepening safety alignment through finetuning.
Stochastic localization. (El Alaoui et al., 2022; Monta- nari & Wu, 2023; Alaoui al., 2023; Montanari, 2023a; Huang al., 2024) applied Eldan's 's stochastic localization method (Eldan, 2013; 2020) to develop new sampling algo- rithms for distributions inspired by statistical physics. Our work applies the stochastic localization framework (Mon- tanari, 2023b) to understand an empirical phenomenon ap- pearing among different oocalization-based samplers widely used in practice.
1.2. Related work
We briefly overview some related work here and defer our discussion of other relevant literature to Appendix A.
2. Technical preliminaries
Theory of critical windows in diffusion. Several recent works have studied critical windows in the context of diffu- sion models, using either statistical physics methods (Raya & Ambrogioni, 2023; Sclocchi et al., 2024; 2025; Biroli et al., 2024) or Girsanov's 's theorem (Li & Chen, 2024). The statistical physies papers assume an explicit functional form for the data and use accurate and non-rigorous statistical physics methods to compute critical windows. For instance, (Biroli et al., 2024) computes the critical time at which the reverse process specializes to one component for a mixture of two spherical Gaussians using Landau-type perturba- tive calculation, and (Sclocchi et al., 2025; 2024) passed through a mean-field approximation to compute the critical windows for random hierarchy model (Petrini et al., 2023), a multi-level context-free grammar with random production rules. Our work is most similar to (Li & Chen, 2024), which derives rigorous, aon-asympooiic bounds analogous to our Theorem 3.1 for mixtures of log-concave distributions with Girsanov's theorem (Chen et al., 2023).
Probability notation. Given distributions P, Q defined on (0, with a base measure Je, the total variation distance is defined as TV(P,Q) - 2 I[PP .Q dp. For random variables X, YY, we will also use TV(X, Y), as shorthand to denote the TV of the measures of X,,Y. Let supp P) {or E (x)) > 0} denote the support. We will also use the following well-known relationship.
Lemma EPE
2.1.
For probability
measures
P,Q.
dQ JP-dQ
S
To study critical windows in diffusion and autoregressive models, consider forward-reverss experiment. A forward-reverse experiment considers the amount of "noise" one would need to add to generation so that running the generative model starting from the noised generation would still yield a sample with the same feature. For a diffusion model, this could mean taking an image of a cat, adding Gaussian noise, and resampling to see if the result is still a cat. For language model, it could mean truncating story about a cat and resampling to check if the story remains about cat. Now, we will use the language of stochastic lo- calization to place these analogous experiments for diffusion and language models within the same framework.
In contrast to existing work, our theory applies to all localization-based samplers, including diffusion and autore- gressive language models, and imposes no functional form or log-concavity assumptions on the distribution. We also improve upon the main theorem of (Li & Chen, 2024) by
obtaining dimension-indee ierensionnindeeeedeent error bounds. Using our improved theorem, we can extend the definition of hier- archy of critical windows from (Li & Chen, 2024) all localization- samplers and, for continuous diffusions, to distributions beyond mixtures of Gaussians.
2.1. Stochastic localization samplers
We formally define the framework for stochastic localiza- tion samplers, following (Montanari, 2023b). Xet ~ p be a random variable over Rd 2 We consider a sequence of random variables (Y:)teT with a compact index set I c [0, oo) U(oo). As / increases, Y, becomes less infor- mative and degrades the original information about X (Defi- nition 2.2). As in (Montanari, 2023b), we will only consider complete observation processes, where information about the path (Ye)ter uniquely identifies X: for any measurable set A C R", we require P(X E (YY))r) E (0, 1}. For the sake of simplicity, we will assume 0, oc E I and Yso is totally uninformative about X.
Forward-reverse experiment. Here we study critical win- dows with the forward-reverse experiment, where we noise and denoise samples with a given attribute to understand critical windows. This was also explored in (Li & Chen, 2024; Sclocchi et al., 2025; 2024). This approach is very similar the framework in which one imagines re-running the reverse process at intermediate point Y, (Georgiev et al., 2023; Biroli et al., 2024; Raya& & Ambrogioni, 2023). Both perspectives provide rigorous frameworks to under- stand critical windows, and in the case where the forward process is deterministic, i.c autoregressive language mod- els, these frameworks are equivalent.
Definition 2.2. (Yi)ee is an observation process with
"These definitions are easily carried over the setting where X lives a discrete space.
3

---

Blink of an eye: a simple theory for feature localization in generative models
respect to X if for any positive integer k and sequence h 122 .... l I, the sequence X + Yr, - YY²- + + Yi forms Markov chain.
2.2. Forward-reverse experiment
Now we describe the general formalism under which we will study critical windows. Fixing some t E I and 5 + we start with some X ~ p², sample Y; X from the observation process conditioning on X, and finally take X from the stochastic localization sampler conditioning on Y,. This can be understood as a generalization of the forward- reverse experiment in diffusions, originally studied in (Sclocchi et al., 2025; 2024; Li & Chen, 2024), to arbitrary stochastic localization samplers.
Because X + Y, + ... - Y1, is a Markov chain, its ... - Vn - X is also a Markov chain. To any such observation process one can thus associate a generative model as follows:
reverse Vt, ....
Definition 2.3. Given observation process (YI) and times
h <.... Im = oc in I, the associated stochastic localiza- tion sampler is the algorithm that generates sample for X by first sampling Yt,, and then, for k m - 1, m 2.....0 sampling from the posterior on Yt, conditioned on Yo...
Forward process. For any t E I, define the forward Markov transition kernel P.- (A(X) = P(Y, E A X). Note the forward Markov transition kernel does not depend on the distribution of X. The fact that the forward process is agnostic to the specifics of the original distribution shared by the most widely used stochastic localization samplers. For example, in diffusion and flow-matching models, the forward transition is convolution of X with a Gaussian; in autoregressive language models, it is masking of the last remaining token in the sequence.
by taking one step in the reverse Markov chain above, and finally sampling X conditioned on Yi.-
In Appendix B, we formally verify that diffusion and au- toregressive models are special cases of this framework. In practice, one does not have access to the true posteriors of the data distribution and must learn approximations to the posterior from data. This issue of learning the true distribu- tion is orthogonal to our work, and thus we define X ~ p to be the sampler's 's distribution. Furthermore, it is more natural to study the sampler's distribution for applications such as interpretability or jailbreaks.
For any I E and S C O, we let p denote the law of YP, where we sample X ~ ps and then sample Y,S P. (- We omit the θ in p.¹
Reverse process. For any E I and initial distribu- tion X ~ p. we define the posterior of X given Y, by (A(r)) = Px~p(X E A|YY),, that is, the distribution of X given by starting the sampling process at E I and Y, instead of oo and Yoo. We will also use this notation for the probability density.
Features, mixtures, and sub- mixtures. To capture the notion of feature of the generation, we assume that the dis- tribution X ~ p is a mixture model. Consider a discrete set = (0) 0x] } with non-negative weights WT w....., Wx summing to 1. Each 0, E Ô is associated with a probability density function pa R >> R20, To generate a sample X p, we first draw 0- Cat(O, {w:) K
Now, we are ready describe the main forward-reverse experiment that we will study throughout the paper.
(==) and return X p". This yields an overall density of p _ Loce wop". For any non-empty S C O, we also define the uu--mixture p by p & =EEES Eoes wo.
Definition 2.5 (Forward-reverse experiment (Sclocchi et al., 2025; 2024; Li & Chen, 2024)). For nonempty S c O and i E I, let ft be the distribution of X defined by the following procedure:"
wip
Remark 2.4. Note that the definition of θ is extremely flexible and can be tailored to the particular data modality or task. For example, θ could be (cat, dog) for image diffusion models; (right, wrong} for math and reasoning tasks; {unsafe, safe} for jailbreaks. The ability to attach these labels and partition into subpopulations is broadly applicable the datasets, which we also make use of in our experiments.
1. Sample Y ~ psJ - i.e run the forward process for time f starting at the sub-mixture ps.
2. Sample X XS. P (-12)) cess starting at Y to sample from the posterior on X.
- i.e. run the reverse pro-
We emphasize that in the second step, we run the reverse process with the prior on X given by the entire distribution rather than the sub-mixture pa. If we did the latter, the marginal distribution of the result would simply be p². In- stead, the marginal distribution of X S is some distribution whose relation to p and sub-mixtures thereof is a priori unclear. Intuitively, as T + 0, this distribution converges
Here we study a family of observation processes correspond- ing observation processes for different initial distributions of X S for S C θ. To ensure that we can meaning- fully compare the observation processes within this family, we will assume that the degradation procedure is fixe.. To formalize this intuition, we borrow the language from diffu- sion models of a forward process, which degrades X, and a reverse process, which takes a degraded V, and produces X.
NNte that this equips 2 with the structure of poset, i.e. A C B if and only if there exists some E I such that running the forward-reverse experiment up to t from pd yields p
4

---

Blink of an eye: a simple theory for feature localization in generative models
to p², and T + oo, this distribution converges to p. The essence of work is to understand the transition between these two regimes as one varies T.
Piis cat)
TC after
T before
2.2.1. INSTANTIATION FOR LANGUAGE MODELS
Consider an autoregressive language model, which produces stories of cats or dogs. We have a discrete set of tokens A and X E A representing length-7 sequences. Letting I= (0 1, ...,. ,. the observation process is defined with Y, A Yo _ X. and Y; being the first T elements of Y for = I...., T. It is easy to see that this is Markovian and the samples become less informative about the original X as I. + 0o.
T
Cats
critical window
Cats and dogs
For the stochastic localization sampler, the posterior Y2|X++1 for adjacent E, + 1 E I the conditional distribution of
Figure 2. Illustration of a critical window for cat feature with an image diffusion model.
T I-length sequences given a prefix of length T - (1 + 1), exactly the task of next-token prediction.
of being a cat) is determined by the diffusion, yielding the critical window the transition period [Ag, B3] in which we specialize to the more specific version of the feature. We can characterize both of these intervals [A1, A2), [B], B2] via the forward-reverse experiment.
Now apply the forward-reverss experiment to a story of a cat. In this case, this means masking the last T tokens of a sample and then resampling with the same model. If T is small, the story will likely still mention a cat and resampling will yield a story about a cat. If T is large, then the first appearance of cat may be truncated, so resampling could produce story about a dog as well.
More generally, let Sinit C ④ denote some sub-mixture, corresponding to uub-population of p that possesses a certain property. Let Starg C Smit denote some sub-mixture containing Smit. For instance, Starg might correspond to all possible responses to the math question, including incorrect ones. We are interested in the following question: if we run the forward-reverse experiment for time 7 starting from S there some range of times for which the resulting distribution is close to Stare? That is, can we characterize the T for which
3. Characterization of critical windows
For expositional clarity, we will introduce our theory along with toy setting to provide some useful intuition: consider an image diffusion model that outputs pictures of orange cats, brown cats, and dogs and a hypothetical critical win- dow in which we transition from sampling from both cats and dogs to cats with certainty (Figure 2). In this case, the critical window simply refers to the dramatic increase in the probability of sampling a cat between time interval [TC (Taber: -- Note that the intuition heretofore will apply to all localization-basse based samplers with a forward and reverse process indexed by time E R+
( Smme of Stare)
is small?
Suppose one could prove that the range of T for which this is the case is some interval [To. Tı]. This would mean that if the stochastic localization sampler runs for time T and ends up at a sample from Srtn then from time T Ti time T of the generation process, the sampler has not yet localized the features that distinguish priui from the larger sub-mixture y However, the sampler has localized the features that distinguish pStwes from 9- When there is a shift from localizing the features Starg to the features Stnit we say there is a critical window. We now formally state and prove our main result.
Importantly, the critical window is situated between and in some sense defined by the two intervals which surround it: the left interval indicates that we have committed to only sampling from the distribution of cats indexed by Safter and have a high probability of producing a cat as the final image, and the right interval, indexed by Sbefore. indicates we can sample images of cats dogs. Thus critical windows simply represent the speciation from sampling from larger sub- population of the distribution given by Sbefore to a smaller sub-population given by Safter- In other words, we want to understand at what interval t E [B1, By] is some feature (herein the feature of being a cat or dog) determined by the diffusion process and then at what interval [A1, Ag] the more specific version of the feature (herein the feature
3.1. Main result
We will parameterize the interval by the total variation dis- tance between sub-mixtures Smit and Stary along the for-
5

---

Blink of an eye: a simple theory for feature localization in generative models
ward process. For an error parameter 0 e < define
localization samplers, not just diffusions. Secondly, Li & Chen (2024) needed to assume that the components of p were strongly log-concave and that the score, i.e. gradient of the log-density, of Pt was Lipschitz and momenttboundee for all t. Thirdly, their final bound includes a polynomial dependence on the moments of the score, which scale with the dimension d; in contrast, our final bound is independent of d.
(e(()) E supft E I: TV Stare @-Stnrk)
(pr .Pt
2 =1-E2 se).
To (e) E inffl E I: (p Smit Starx)
.P
When the value of e is understood, we abbreviate the above with T and T eud- Our main result is that in T E In st (Tend: T the distance TV (y Tf -SSuar)) is small:
Theorem 3.1. Let Sinit C Starg C O and W Star "i Eveguar
With Theorem 3.1 in place, we are ready to formally define critical windows. These capture the moments where we transition from sampling from sub-mixture subset of that
T.l. then TV T .pStur)) c + max (1, WVV2..
For 0, if E In IT I
Definition 3.3. Define Safter C &Sbefore C +. For Shefore.
Intuitively, Tu represents the largest for which there still separation between Starg and @ - Starg and T repre- sents the smallest for which samples from Smitt Starg are indistinguishable. Thus, running it for T E In [Tend T:s erases the differences between samples from Siit and Starg but preserves the difference between Starg and θ Searg yielding samples looking like y yStare.
we define Thetore = inf(l E : TV(p Satter Statorre .Pi
S
and TV Shefere .Pr (-Sbaforr) > 1 2 (Sinit A
Safter' Starg Sbefore). For Safter, consider T sup(t E I (Stit. Starg _ Safter). A critical window is the interval [T Teeforels where there is a transition from sampling from Sbefore to the smaller subset Safter-
after
TV -P: p.--Satie:) 2 11-1 e²}
Finally, note that this definition does not directly explain why when critical window is sharp, which depends on the properties of the model of the data and forward process. In Section 4, we analytically compute critical windows in different toy settings and demonstrate that they are sharp in a natural sense. These examples also provide general intuition where critical windows are sharp, i.e. for autoregressive LLMs, critical windows may occur when a few tokens are very unlikely under one mode than other in Example 4.3. In general, we intuit that sharp critical windows occur when it only takes few steps from the forward process to erase the differences between Sbefore and Safter This could happen if the data has multi-scale hierarchical structure, where a feature is decided in a narrow intermediate band of the tree (Section 5).
Low noise
Hcsk
HosA
Hcs
Medium nodiu noise
Hcsh
HCS
High noise
HCS
Figure 3. Intuition with the forward-reverse experi- ment with Smit = forange cat} and Starg (orange cat, brown cat). In the low-noise regime, the forward-reverse experiment does not modify the image non-trivially, meaning we still sample from Sinit; in the high-noise regime, the forward-reverse experiment destroys all information in the image and thus could yield a dog, lying outside of Starg. Thus there is "sweet spot" where the differences between orange and brown cats collapse but those between cats and dogs persist.
We finally define some helpful notation for our proof. For nonempty S C θ and E I, we define p- P- to be the posterior of X with the prior X ~ p.. We similarly define (- and (- (Y S) to be the posterior of
o conditioning on Y, with X ~ p or X ~ pp, respecttii. When S = (i), we exclude the braces.
3.2. Proof outline
A priori it is not clear why T'S should be smaller than st and indeed in general it need not be and our bound would be vacuous. In Section however, we show that in many natural settings for diffusion models and autoregressive models, the relation does hold.
Crucially, our proof relies in several places on the Markov property of stochastic localization samplers, together with the data processing inequality.
Proof of Theorem 3.1. By ity,
the triangle Sma. pSiee)
inequal-
Remark 3.2. Note a similar bound was shown in the context of diffusions by Li & Chen (2024) (see Theorem 7 therein). Our result is a strict improvement of that bound along sev- eral important axes. First, our results apply to all stochastic
write T
we
can
Sp .f p ppraan T + TV
9ran T pp Suar)) )
Sann
1
(
6

---

Blink of an eye: a simple theory for feature localization in generative models
and pSuux.?? are the laws of the posterior P (-I-) but applied to Yp with distributions P7 Sinss and PT Suoe
T TC after then TV -pt) Se e. When t 2 The before" TV .P) Se.
Using the Markov property of localization-based sam- plers (Definition 2.2), we apply the data processing inequality and the definition of T to bound (I) via TV )) S TV pir
For an isotropic Gaussian mixture model with randomly selected means, the critical window between sampling from one component to the entire mixture is also narrow. Note that we derive dimension- free widths in Example 4.2, an improvement over (Li & Chen, 2024) who had a In d dependence on dimension for isotropic Gaussians.
Sigit Stare)
-Pr E.
end
To bound (II) we apply the definition of TV and a coupling argument. By the law of total probability we can express p (x) (z) = (x))(x)X?)) and pStre (x) EIP (a|Yz,, Starg)] for Y7 ~ Pf
Example 4.2. [Random mean spherical Gaussians] We first sample bi ~ NN Id) for i E [K] i.i.d. and let We let Sbefore = O and Sante = {ug]. Then, we can compute T before = max In || Myll + In(1/e) and Tater = minze In lIke - uyll r.
Stacx as these observation
θ = (N(,,,d))EEi
processes have the same distribution at index T. Thus, TV pp - _ I ()swe))(s) -6))-p0 -pStrr (x) de- I (x1y,]] - -B(P--[PA -Ep"((xXY Starg)] dz.
1/ 8 In K, Furthermore, with high probability over the se- lection of the means, Tefore -TC after = (1,e(1) as d + co.
By Jensen's inequality and Fubini's theorem. we bring &
the expectation outside the integral, S p p JSE P Yp) - P Yp. Starg) dx J PT (xl)p) - PT (zr) (x, Starg) ) dx].
Example D.5 of discrete diffusion model, where T after
T + 0, is deferred to Appendix D.5.
C.3 Appendix C), we find TV o pSux)
Combining Lemmas C.2 and
(proved
in
4.2. Autoregression
We first present a theoretical model for important critical windows in LLMs, e.g., jailbrekss that occur over the first few tokens in the generation and the Yellowstone exam- plc (Anthropic, 2024; Qi et al., 2024).
max (1, W)E LE e- Ssors (pp) Paaes (Yp))
Then, finally
applying Lemma 2.1, we are able to bound the total vari-
ation in terms of E, obtaining (V T p pSaxx)
Example 4.3. ["Critical Tokens" for Jailbreaks and Yel- lowstone (Qi et al., 2024; Anthropic, 2024)] Again con- sider an autoregressive language model, with A denot- ing the vocabulary, dexed by I = {0. 1, 2,.... T), T tokens of X. Let θ {GGoogling Yellowstone: feoding H). We assume that these two modes do not differ until some T T' E I. Between T - T' and T - T' - k, the distributions become nearly disjoint, P Pa-, Phurntul (I E supp(p?"m--.) S e. In the "---pr jailbreaking example, T" = 0 and they are disjoint because the first tokens generated in the safe mode is always some form of refusal. In the Yellowstone example, they are dis- joint the first time the agent decides to Google Yellowstone pictures. Then, on component Gharmful we have the critical window The hefore = T T' and TCC Tatter T T k.
max (1, W) /1- 1 (p
2 Stare StorE)
.PT
2 max (1, W).. Combining our bounds on (I) and (II) achieves the desired result.
E A", a forward process in- and V, to be the first {Oharmfull Ssate] (or
4. Examples of critical windows
In this section, we analytically compute Th after for diverse stochastic localization samplers and models of data, including diffusion and autoregression processes. In these natural settings, the critical window is small in the sense of having a size which shrinks or does not depend on the dimen- sion or context length. We shall also connect our framework to in-context learning and the all-or-nothing phenomenon."
4.1. Diffusion
We first consider two examples of Gaussian Mixture Models and a diffusion model. We show that with two isotropic Gaussians, the critical window appears around single point, In [|и]], with width independent of the dimension.
This prediction that the critical windows are narrow in these settings is verified by experiments from (Haize Labs, 2024b; Qi et al., 2024), which show that one consequently ob- serves critical windows for jailbreaks. Moreover, notice that one can actually mitigate the effect of these critical windows by finetuning on examples of corrections to in-
[Two Isotropic Gaussians] Let θ = {+1), p+1 N(u,,d) p 1 WV(-u,d). Then, we have critical window transitioning from sampling from both components to the component +1 between TC In [|l||| In 1/e and 1 = In || In In 2,. When
Example 4.1.
(x E supp(pTT This explains the effectiveness of finetuning on corrections in (Qi et al., 2024).
crease P Ptamnna T'
Remark 4.4. The quantity that measures probability of mode-switching p /p, suggests using a likelihood ra-
Prooff are deferred to Appendix D.
7

---

Blink of an eye: a simple theory for feature localization in generative models
tio to distinguish between harmful and benign prompts. In App. F.2, we test class of likelihood ratio methods that obtain recall 5-10x the false positive rate for 5 different types of jailbreaks (Table 2).
ment of the original distribution p into increasingly specific components. To formalize this, we introduce the concept of e-mixture tree, which decomposes p into hierarchical structure.
Definition 5.1 (Informal, see Definition E.1). For an error term c and mixture model p. an c-mixture tree is tuple (T, (P-((-))) O (p")eee Subset, Subsett Noise T (V, F) is a tree associated with function Subset: V - (00, which maps vertices to subsets of Q and AoiuuAmount: - R which characterizes the noises levels for different aggregations in the mixture tree. If u is a parent of v. Subset(v) C Subset(u) for u V: NoiseAmount(u)) is defined such that all PNoiseAmouna(()) for O E Subset(u) overlap greatly and for Subset(u) P (NoiscAmoun(()) (u)
We defer critical window for stylized model of solving math problem as random walk (Example D.6) and a critical window for an autoregressive model which expresses the outputs as emissions from a random walk of an underlying concept variable, akin to the model in (Arora et al., 2019), to Appendix D.2.
4.2.1. IN-CONTEXT LEARNING
Autoregressive critical windows can also be applied to de- scribe in-context learning. In particular, we can capture the idea that with sufficiently many in-context examples, we learn the 8" E θ that generated the transitions for in-context examples, with a sample complexity in terms of TC after
u P (NoiscAmoun(()))
e Subset(u)
have negligible overlap.
(orange cat, brown cat, dog)
Example 4.5 (Informal, see Example D.16). Con- sider an in-context learning setup, where the context [21, MI, (o...,,IT+1.. 1 consists of question-answer pairs (24, delimiters O, and x; + 1/s sampled from p" for some 8" E O. In the forward-reverse experi- ment, we truncate it to [21, 0,, IT+ and then re- sample with p to produce [æ1..... 2T11, 0T+1, o). The total variation between the sequences [x1, M1, o, x1] and [21..... +1: 0r+1, o| can be viewed as the average-case error of the in-context learner and can be bounded within our critical windows framework. We have after = 37 + 3 - (1). with O₂ (1) independent of T (Sater & {8*)). Note that T is the order of how many samples that can be erased so that we still are able to learn 8" O
forange cat, brown cat}
f
{orange cat)
bbronn cat}
Figure 4. Example of an e-mixture tree diffusion model that produces images of orange cats, brown cats, and dogs.
We emphasize that this framework is highly general, solely defined with the initial distribution p and the forward pro- cess. It strictly expands the definition in (Li & Chen, 2024), which focused on hierarchies of isotropic Gaussians, to all localization- baased samplers and mixture models. We can also relate it to the sequences of critical windows we observe in Figure 8, capturing the idea that each critical window rep- resents the refinement into smaller subpopulations of p.
One might ask if there is The for in-context learning, threshold such that it is impossible to distinguish between Satter Shefore with that many samples. In Appendix D.4, we will provide an example of a Thetore with the all-or-nohhing phase transition from statistical inference.
Corollary 5.2. Consider an e-mixture tree For 0g O, consider the path un U2 U3...., urr E V where un is the leaf node with O, E Subset(u1) ) and uH is the root. There is a sequence of times Tı < T2 < ... < The with (Vp11) ,.Subber(tr)) Sw c.
5. Hierarchies in stochastic localization samplers
Thus far we have largely focused on a specific critical win- dow, which would correspond to the formation of specific feature in the generated output. In some cases (Figure 8), there could be sequence of critical windows, motivating theory of hierarchical sampling within our critical windows framework.' 6 This is naturally represented as a tree, where the vertices correspond to the cluster at some resolution of the forward process, the root is the entire data distribution, and the leaves are raw components of the distribution. A path from the root to a leaf captures the progressive refine-
We first observe that the hierarchy of two samplers with the same forward process are identical if the samplers agree on nub-populations. Assume we have {p") jeee (e.g. the true distribution) and (4 Jese (e.g. generative model), where qa R p" across all A E θ with the same (wo]eee-
Corollary 5.3. Consider (T (P ((,)),,, (nr (n")nees, "nees,, NoiseAmount). Suppose we have another distribution {q") such that TV(p", ,q²) S 8/2 for all 0 E O. Then we have C + Vo-mixture tree given by (P.(P jeee, Subset, NoiseAmount).
e-mixture tree
Details in Appendix E.
8

---

Blink of an eye: a simple theory for feature localization in generative models
This similarity does not hold generally. We can define ar- bitrary hierarchies by choosing the right forward process.
Table I. Differences between Accuracy (Ace) without versus with critical windows and frequency of critical windows (CW) when the original generation is wrong versus right.
Example 5.4. Consider set of alphabets {A.) and de- fine ④ = ={(a){= Vi E [d]. az E A,) and ph: 50,- Let d]. and for any permutation i1, f2, , Rd of [d]. define a forward process such that at t E I, we mask all iæ, id-.,..., dd-t- This constructs a hierarchy where the values for i,, i2. ...,.., id are decided in that order.
Dataset
AAcc
ACW
AQUA-RAT
0.42
0.20
I=
ARC Challenge
0.53
0.22
ARC Easy
0.73
0.26
LogiQA
0.15
0.11
MATH
0.41
0.33
Finally, we note that hierarchies of diffusions are in general shallower than for autoregressive models. The hierarchy for
MMLU
0.45
0.24
TruthfulQA
0.44
0.20
a mixture of Gaussians has depth O(1) (Example 4.2), as the
forward process simultaneously contracts all distances with
MATH (Phi-3-7B. I-Inntuuutt)
similar dependence on d. For autoregressive models, depth
Q: If the first tenn of an ar arithmetic sequence equals the third minus the secund, and founth torm
can scale linearly with the context length (Example 5.4). While they refer to different modeling approaches, it is worthwhile to highlight this vast difference in hierarchy depth can vary with the dimension. We speculate that this could mean autoregressive models can learn more complex feature hierarchies than diffusions.
tems? A: know that aritiientic sequence has difference, which we will denoto as "d"", 2. Let's denotc the tomas sequence as Ti. T2. Ta. and Ta the fins, second, third, and fourth terms, respectivelyy 3. We're given the information that 71 = T3- T2. Using the general formula for an anishmetic sequence,
T2 Ty + 22 - Tg. which simpliiies to I3 2d. 4. We're also given T4
wrile
The general formula sequence tells us that T3 = Ti 3d. So we cun write 8- Ti + 3d. Now we have system of two oquations: 1) Ti 2d 2) 8 Ti 3d 6. We can substitute equation (1) into equation (2), which gives us a 2d + 3d od. Solving d. get
arithmetic
6. Experiments
As many authors? have empirically studied critical windows for diffusion and verified different theoretical predictions
8/5.
d. we can find Ty. Using the equation from step I,
Ti 2d 2- (8/5) 1G/5. Tharafore, the fisst tena
with experiments for this model, we focus on experiments on critical windows for LLMs. In Appendix G, we validate the predictions of our theory on outputs with hierarchical structure, showing strong agreement with Section 5.
T1 16/5.
Figure S. Example of critical window (bolded).
Then we identify critical windows for LLAMA-3.1-88 Instruct. Phi-3-7B-Instruct, and Qwen-2.5- tiven-2.5-77 7-.----BB 15-7B-Innttuu Instruct in the correctness of real-world math and reasoning bench- m In our forward-reverse experiments to identify crit- ical windows, we take the original generation, truncate a fraction of tokens, and check if resampling yields the same answer. We say that a particular generation has critical window if there is a >> 0.5 jump in probability of obtaining the same answer in consecutive truncation fractions.
portant for the final answer, resulting in no critical window. In fact, we find that generations with critical windows are significantly less accurate than generations without critical windows across all datasets and models. For ILAMA-3. 1- IB-tnstu critical windows result in up to 73% lower accuracy, and incorrect responses contain 11 11%-3% % more critical windows (Table 1).
7. Discussion
We find that critical windows appear prominently across all benchmarks (Figure 10 to 12); for MATH, they occur in 42.2%, 16.6%, 30.2% of generations from LLAMA-3.1 8B- SB-Instruct, Qwen-2.5 2.5--B--BBBBBBB In--tuutt and Phi-3-7B --B-3-BB-Ao Instrrcct These jumps seem to occur during important steps of reason- ing, as certain parts of the solution may be more significant to the formation of the final answer. For example, in Fig- ure 5, the model restates the correct formula to apply to the problem within the critical window. For other generations, could be that no particular part of the text is especially im-
In this work, we developed simple yet general theory for critical windows for stochastic localization samplers like diffusion and autoregressive models. Already, practitioners have applied critical windows to make LLMs safer (Qi et al., 2024) and reason better (Abdin et al., 2024; Lin et al., 2024). Our theory significantly streamlines our understanding of critical windows and provides concrete insights for practi- tioners. While we instantiate our theory in many settings (Section 4), an important future direction is connect it to statistical physics predictions for other models of data, e.g., the random hierarchy model (Sclocchi et al., 2024). An- other interesting direction would be to explore why critical windows exist in different settings in practice.
"(Ho et al., 2020; Meng et al., 2022; Choi et al., 2022; Raya & Ambrogioni, 2023; Georgiev et al., 2023)
*See Appendix H for full results across models and datasets and a discussion on the effect temperature on critical windows.

---

Blink of an eye: a simple theory for feature localization in generative models
Acknowledgments
Wann__Shot_Jailrreeaee 2 0936.pdf.
2024_04_02
ML would like to thank Cynthia Dwork for regular meetings and kind mentorship which were of great help throughout this research project, Cem Anil, Michael Li, and Eric Ze- likman for insightful conversations which inspired some of the experiments, and Max Simchowitz, Michael Albergo, Seth Neel, and Sahil Kuchlous for thoughtful discussions regarding the framing of this work.
Anthropic.
Developing computer use model.
http://wwwww..ppccccccoeeeeeeeee con/news develloping-coopptee--u-ses 2024.
Arora, A., Jurafsky, D., Potts, C., and Goodman, N.
Bayesian scaling laws for in-context learning, 2024. URL https: /ttp:///aaxvvvvee org/ abs 2410. 16531.
Impact statement
Arora, S., Li, Y., Liang, Y. Ma, T., and Risteski, A. A latent variable model approach to pmi-based word em- beddings, 2019. URL http:///axvvv...eeveeeee
While this paper is largely theoretical in nature, it does de- scribe theory for jailbreaks which could impact model safety in the future. We hope the insights in this manuscript about jailbreaks lead to better alignment strategies and train- ing methods.
Bailey, L., Serrano, A., Sheshadri, A., Seleznyov, M., Taylor,
J., Jenner, E., Hilton, J., Casper, S., Guestrin, C., and Emmons, S. Obfuscated activations bypass Ilm latent- space defenses, 2024. URL https: http::://aaxxvvorrr abs 24..2412...0 09565.
References
Abdin, M., Aneja, J., Behl, H., Bubeck, S., Eldan, R., Gu-
nasekar, S., Harrison, M., Hewett, R. J., Javaheripi, M.,
Barbier, J., Macris, N., and Rush, C. All-or-nothing statisti- cal and computational phase transitions in sparse spiked matrix estimation, 2020. URL https: //arxiv.
Kauffmann, P., Lee, J. R., Lee, T., Li, Y., Liu, W., Mendes, C. C. T., Nguyen, A., Price, E., de Rosa, G., Saarikivi, O., Salim, A., Shah, S., Wang, X., Ward, R., Wu, Y., Yu, D., Zhang, C., and Zhang, Y. Phi-4 techni- cal report, 2024. httph:t//aaaaa org/ abs/ 2412.08900 08905.
org/ abs/2000555 079771
Biroli, G., Bonnaire, T., de Bortoli, V., and Mézard, M. Dynamical regimes of diffusion models, 2024.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learning? investigations with linear models, 2023. URL https: //arxiv. org/ abs/ 22112221111 15661.
Chao, P., Robey, A., Dobriban, E., Hassani, H., Pappas, J. and Wong, E. Jailbreaking black box large language models twenty queries, 2023.
Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., and Zhang, A. Sampling is easy as learning the score: the- ory for diffusion models with minimal data assump- tions. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, OpenReview.net, 2023. URL https: / //openreview net
Alaoui, A. E., Montanari, A., and Sellke, M. Sampling from mean-field gibbs measures via diffusion processes. arXiv preprint wrXiv:2310.08912. 2023.
Alon, G. and Kamfonas, M. Detecting language model attacks with perplexity, 2023. URL org/ brgg/bs/// abs/ 2308. 14132.
Chen, Y. and Eldan, R. Localization schemes: A frame- work for proving mixing bounds for markov chains. In 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science (FOCS), pp. 110-122. IEBE, 2022.
Anil, C., Durmus, E., Sharma, M., Benton, J., Kundu, S., Batson, J., Rimsky, N., Tong, M., Mu, J., Ford, D., Mosconi, F., Agrawal, R., Schaeffer, R., Bashkansky,
N., Svenningsen, S., Lambert, M., Radhakrishnan,
Choi, J., Lee, J., Shin, C., Kim, S., Kim, H., and Yoon, S. Perception prioritized training of diffusion models. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11462-1144 11471, 2022. doi: 10.1109/CVPR52268..8 202.22266688.2211111.1
A., Denison, C., Hubinger, E, J., Bai, Y., Bricken,
T., Maxwell, T. Schiefer, N., Sully, J.. Tamkin, A., Lanham, T. Nguyen, K., Korbak, T.. Kaplan, J.. Ganguli, D., Bowman, S. R., Perez, E.. Grosse, R., and Duvenaud, D. Many-shot jailbreaking. In Advances in the Thirth-Eighth Annual Conference on Neural Information Processing Systems, 2024. URL sttps://wwwccde can ant com/ 42533333eeeebbbb2222222333333333
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1803.05457. 2018.
10

---

Blink of an eye: a simple theory for feature localization in generative models
Coja-Oghlan, A., Gebhard, O., Hahn-Klimroth, M., Wein, A. S., and Zadik, I. Statistical and computational phase transitions in group testing. In Loh, P.-L. and Raginsky, M. (eds.), Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pp. 4764-4781, PMLR, 02-05 Jul 2022. URL htt.L:tt//////mm /ttpttpes//ppocceediing patessmrr.pp..... p.rsss/ v1 #278/coja--ghlan222 html.
Haize Labs. Automated multi-turn red-teaming with cas- cade. https: //lllll...... ha meeizeeabs wtw.aaaa..oooo.oo..oo..m.m.....a com/ coompostts/ cascade, 2024a.
Haize Labs. A trivial jailbreak against Ilama 3. https: github com/ htt//hiiaee haizelabbe hatss/a1mmaaaaaaaaaaaa labbeeak,, 2024b.
He, L., Xia, M., and Henderson, What in your safe data? identifying benign data that breaks safety, 2024. URL https: aaaaaaiivv.... abs 01099.
Ding, N., Chen, Y. Xu, B., Qin, Y., Zheng, Z., Hu, S., Liu, Z., Sun, M., and Zhou, B. Enhancing chat lan- guage models by scaling high-quality instructional con- versations, 2023. URL https: //axxxxivvv..///////////////avva org/ abs/ 2305. 14233.
Hendryeks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, Measuring massive multitask language understanding. In International Conference Learning Representations, 2021a. URL https:// openreview net forum? i d=d7kB jm 3mI3Gmç.
Dong, Q. Li, L., Dai, D., Zheng, C., Ma, J., Li, R., Xia, H., Xu, J., Wu, Z., Liu, T., Chang. B., Sun, X., Li, L., and Sui, Z. A survey on in-context learning, 2024. URL https: / /aarivvvv.. www.:bgg///ooo.ooooooo....s..s htt.s///31.0000 00234.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring math- ematical problem solving with the math dataset, 2021b. URL https::////...m arxiv org/ abs 2203. 03874 3874.
El Alaoui, A., Montanari, A., and Sellke, M. Sampling from the Bneerington-kinkpatrii gibbs measure via algorithmic stochastic localization. In 2022 IEEE 63rd Annual Sym- posium on Foundations of Computer Science (FOCS), pp- 323-334. IEEE, 2022.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. Advances in Neural Information Pro- cessing Systems, 33:6840-6851, 2020. URL https: / //arxiv. wwwp:i.org/abb//00000 11239.
Eldan, R. Thin shell implies spectral gap up polylog via stochastic localization scheme. Geometric and Func- tional Analysis, 23(2):532-56) 2013.
Huang, B., Montanari, A., and Pham, H. T. Sampling from spherical spin glasses in total variation via algorithmic stochastic localization. arXiv preprint arXiv:2404. 15651, 2024.
Eldan, R. Taming correlations through entropy-efficient efficient measure decompositions with applications to field approximation. Probability Theory and Related Fields, 176(3-4):777-755, 2020.
Lanham, T., Chen, A., Radhakrishnan, A., Steiner, B., Deni- son, C., Hernandez, D., Li, D., Durmus, E., Hubinger, E., Kernion, J., Lukošiūté, K., Nguyen, K., Cheng, N., Joseph, N., Schiefer, N., Rausch, O., Larson, R. McCan- dlish, S., Kundu, S., Kadavath, S., Yang, S., Henighan, T., Maxwell, T., Telleen-Lawton, T., Hume, T., Hatfield- Dodds, Z., Kaplan. J., Brauner, J.. Bowman, S. R., and Perez, E. Measuring faithfulness in ofain-of-thought reasoning, 2023. URL https: http://aaxvvv...//ee/e.eeaaa org.eiivv.rrroooooooooooooo abs/ 13702.
Gamarnik, D. and Zadik, I. High-dimensional regression with binary coefficients. estimating squared error and phase transition, 2019. URL https: /ataaaiivvog............./.....a abs/1701. 04455.
Garg, S., Tsipras, D., Liang, P., and Valiant, G. What can transformers learn in-context? case study of simple function classes, 2023. URI. htw.tttpp///// http:::://axxv..orrr abs/2208. 01066.
Li, M. and Chen, S. Critical windows: non-asymptotic theory for feature emergence in diffusion models, 2024. URL https:ttpt.........t hawpiviiiv..... org.
Gemini. Challenges and solutions for aging adults. https://gemini: google. cwmcco/////c//oe share/ 6d141b742a13, 2024.
/abp///b//////////a/////aa// abs 2403. 01 633.
Li, N., Han, Z., Steneker, I., Primack, W., Goodside, R., Zhang, H., Wang, Z., Menghini, C., and Yue, S. Llm defenses are not robust to multi-turn human jailbreaks yet, 2024. URL htL.htts::/ httpp:s/////xvoooo//o// abs//40//. 15221.
Georgiev, K., Vendrow, J., Salman, H., Park, S. M., and Madry, A. The journey, not the destination: How data guides diffusion models. arXiv preprint arXiv:2312.06205, 2023.
grimjim. Jailbroken laama-3, 3.1 1-8b-instruct via lora. ht https: / //huggingface.cc hmtp///mmmm///aaaaaaaaaaa
Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let's verify step by step. arXiv preprint arXiv:2305. 20050, 2023.
1- 8--8BIInstruce www.laaalliieeeeeee via adapter, 2024.

---

Blink of an eye: a simple theory for feature localization in generative models
Lin, B. Y., Ravichander, A., Lu, X., Dziri, N., Sclar, M., Chandu, K., Bhagavatula, C., and Choi, Y. The unlocking spell on base Ilms: Rethinking alignment via in-context learning, 2023. URL https:/ http://axxvvoooo org/ abs/
Montanari, A. Sampling, diffusions, and stochastic local- ization, 2023b. URL httpss://aaaiiiiie abs/
2305. 10690.
2312. 01552.
Montanari, A. and Wu, Y. Posterior sampling from the spiked models via diffusion processes. arXiv preprint arXiv:2304. 1.X4:22304111449.22 2023.
Lin, S., Hilton, J., and Evans, O. TruthfulQA: Measuring how models mimic human falschoods. In Proceedings of the 60th Annual Meeting of the Association for Compu- tational Linguistics (Volume I: Long Pupers), pp. 3214- 3252, Dublin, Ireland, May 2022. Association for Com- putational Linguistics. doi: 10. 10.11663/w1/20222aa--onggg 229. URL https: http://aaaaantoo.oor/o///// acl-long 229.
Mossel, E., Niles-Weed, J., Sohn, Y., Sun, N., and Zadik, I. Sharp thresholds in inference of planted subgraphs. Neu, G. and Rosasco, L. (eds.), Pro- ceedings of Thirty Sixth Conference on Learning The- ory, volume 195 of Proceedings of Machine Learn- ing Research, pp. 5573-5577. PMLR, 12-15 Jul 2023. URL httphhttss//pppeeeeeee http:/ppeenn..mmr.r... mll. pr..s.sss../p..////... v
Lin, Z., Liang, T., Xu, J., Wang, X., Luo, R., Shi, C., Li, S., Yang. Y., and Tu, Z. Critical tokens matter: Token-level contrastive estimation enhances Ilm's reasoning capabil- ity, 2024. URL https: //tp////aaiv./i//ii org/ httgg/abb////// 19943.
html.
Nasr, M., Carlini, N., Hayase, J., Jagielski, M., Cooper, A. F., Ippolito, D., Choquette-Choo, C. A., Wallace, E., Tramèr, F,, and Lee, K. Scalable extraction of training data from (production) language models, 2023. URL org/ hbag:/ab//2330 17035.
Ling, W., Yogatama, D., Dyer, C., and Blunsom, P. Program induction by rationale generation: Learning to solve and explain algebraic word problems. In Proceedings of the 55th Annual Meeting of the Association for Computa- tional Linguistics (Volume I: Long Papers), pp. 158-167, Vancouver, Canada, July 2017. Association for Compu- tational Linguisties. doi: 10. 1155//1111111-1115. URL https: htt////aaaatahhtt -ttp:/athhhoooo..oooo/1111011110
https: /ttsss//arxvvv.a www.s..oo//..........
Wiles-Weed J. and Zadik, I. The all-or-nothing phe- nomenon in sparse tensor pea. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 17674-17684, Curran Associates, Inc., 2020. URL
cc 0a000/accc022ccc3322227222223333
Liu, J., Cui, L., Liu, H., Huang, D., Wang. Y., and Zhang, Y. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. In Bessiere, C. (ed.), Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20, pp. 3622- 3628. International Joint Conferences on Artificial Intel- ligence Organization, 7 2020. doi: 10.24963/ijcai.2000/ 501. URL https: / //oiiiioo////ooo..o org/ 10. 29..00.4.222eee 2020 /501. Main track.
pdf.
Niles-Weed,1 J. and 7adik, I. It was "all" for "nothing""" sharp phase transitions for noiseless discrete channels, 2023, URL /tttps://arxitn wwwpeiiiiirrrooooooooooooo httpg//aaaa/aa.///// /2t.2.22.2.22222
Petrini, L., Cagnetta, F., Tomasini, U. M., Favero, A., and Wyart, M. How deep neural networks learn composi- tional data: The random hierarchy model. arXiv preprint arXiv:230.002129, 2023.
Liu, X., Xu, N., Chen, M., and Xiao, C. Autodan: Gen- erating stealthy jailbreak prompts on aligned large lan- guage models. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: //openreview.o httpeevwww...weeeooooeeeeeeeeeoe forum? www.w: qKkb.
Qi, X., Panda, A., Lyu, K., Ma. X., Roy, S., Beirami, A., Mittal, P., and Henderson, P. Safety alignment should be made more than just few tokens deep, 2024. URL https http://axxvvvvoo/ wwwpiiivorr/bbb/////// 05946.
Lou, A., Meng, C., and Ermon, S. Discrete diffusion model- ing by estimating the ratios of the data distribution, 2024.
Raya, G. and Ambrogioni, L. Spontaneous symmetry break- ing in generative diffusion models. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https: https///pppeeeieeeeeeeeeee net forum? id=1xGFGMMSVL
URL https:: /taps://aaee /ttp///aavvv///////////ee//aaav/ org/ abs 2310 6834.
Meng, C., He, Y., Song, Y., Song, J., Wu, J., Zhu, J.,Y., and Ermon, SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference Learning Representations, 2022.
Reeves, G., Xu, J., and Zadik, I. The all-or-nothing phenomenon in sparse linear regression, 2019. URL
Montanari, A. Sampling, diffusions, and stochastic localiza- tion. arXiv preprint orXiv:2305. 10690, 2023a.
12

---

Blink of an eye: a simple theory for feature localization in generative models
Roos, B. Binomial approximation the poisson bino- mial distribution: The krawtchouk expansion. The- ory of Probability & Its Applications. 45(2):258-272, 2001. doi: 10.1 10.1137/0040088X9777XXX URL https: //doi.ore org / htt.//1.1111111011111111 1111111/000000000077777770000000
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning implicit bayesian inference, 2022. URL 2111. 02080.
Zhang, X. and Wu, Dissecting learning and forget- ting in language model finetuning. In The Twelfth In- ternational Conference Learning Representations, 2024. URL https://openree 1istmsgb6Wp1...
Röttger, P., Kirk, H. R., Vidgen, B., Attanasio, G., Bianchi, F. and Hovy, D. Xstest: A test suite for identifying exag- gerated safety behaviours in large language models, 2024. URL wwt.tttps//.cc. /ttttp:://aiivi org/ abs 2308. 01263.
net forum?
Zhang, Y., Zhang, F., Yang, Z., and Wang. Z. What and how does in-context learning learn? bayesian model aver- aging, parameterization, and generalization, 2023, URL wwwpor////////eoeeaeeeeeeeeaeeee
Scarlett, J. and Cevher, Limits on support recov- ery with probabilistic models: An information theoretic framework, 2016. htttttts:://aatt wwwpeii..ooaaaaeeeee abs/ 1501.07440.
httpss//arxxvv.rrg
abs / org/ 2305 19420.
Zhao, Z., Zhang, X., Xu, K., Hu, X., Zhang, R., Du, Z. Guo, Q., and Chen, Y. Adversarial contrastive decoding: Boosting safety alignment of large language models via opposite prompt optimization, 2024. URL https:/ arxiv. azxiv.org/ abs//ab///// 2406..442 0.406.16743.. 16743.
Sclocchi, A., Favero, A., Levi, N. I., and Wyart, M. Probing the latent hierarchical structure of data via diffusion mod- els, 2024. URL https: 13770.
Sclocchi, A., Favero, A., and Wyart, M. A phase transition in diffusion models reveals the hierarchical nature data. Proceedings of the National Academy of Sciences, 122 (1)xe240879912,, 2025.
Zou, A., Wang, Z., Carlini, N., Nasr, M., Kolter, J. Z., and Fredrikson, M. Universal and transferable adversarial attacks on aligned language models, 2023. URL https: //arxiv. org. abs 2200 15043.
Souly, A., Lu, Q., Bowen, D., Trinh, T., Hsich, E., Pandey, S., Abbeel, P., Svegliato, J., Emmons, S., Watkins, O., and Toyer, A strongreject for empty jailbreaks, 2024. URL httpt:t//eeee http://aaxvvvv..
Zou, A., Phan, L.,, Wang, J., Duenas, D., Lin, M., An- driushchenko, M., Wang, R., Kolter, Z., Fredrikson, M., and Hendrycks, D. Improving alignment and ro- bustness with circuit breakers, 2024. URL https: 04313.
abs 24022 10260. 10260.
/ //arxiv.org/abo..
Truong, L. V. and Scarlett, J. Support recovery in the phase retrieval model: coormation-theoretie fundamental lim- its, 2020. URL https::/arviv..rr///eeee// 1901. 10647.
Truong, L. V., Aldridge, M., and Scarlett, J. On the all-or- nothing behavior of bernoulli group testing, 2021. URL https: //tpxavaivv///////////e///aaaav/ org/ httg/abs/2000 10137.
van Handel, R.
Probability in high dimension. 2016. URL https::/ /we.. wwt...aatthh.e.eeeee princeton.. edu/ www. pdf.
Vershynin, R. Digh--Dimessioal Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press. ISBN
978-1-108-4151999 doi: 10.101779781100223159 URL
books/highdine bhoo:ss/hiinddieeesssoooaooooooo T9970666023772222222344322220000
Wei, A., Haghtalab, N., and Steinhardt, Jailbroken: How does Ilm safety training fail?, 2023. URL https:// arxiv. www.orsg.bb/....... 2307. 02483.
13

---

Blink of an eye: a simple theory for feature localization in generative models
A. Additional related work
Theory of in-context learning for language models. With respect to theory for language models, our results are most closely related to the Bayesian framework for in-context learning (Xie et al., 2022; Akyürek et al., 2023; Garg et al., 2023; Zhang al., 2023; Arora et al., 2024). For example, (Xie al., 2022) also considered mixture model of topics and showed that language models can learn the underlying class despite in-context learning and training distribution mismatch. We view this manuscript as connecting the Bayesian framework for in-context learning to other empirical phenomena observed in language models and diffusion models and the all-or-nothing phenomenon.
Chain of thought. (Lin et al., 2024; Abdin et al., 2024) also observed that the presence of critical windows in the chain of thought of math and reasoning tasks and their significance in leading the model to incorrect outputs, concurrent with our results in Figure 5. They then used them to provide rewards or data for a preference optimization algorithm to improve reasoning performance. (Lin et al., 2024) called them critical tokens and utilized a contrastive estimation algorithm to identify critical windows and provide token-level rewards. The Phi-4 Technical report called them pivotal tokens, developed binary-search based algorithm to identify the location of critical windows, and used them to produce contrasting pairs for preference optimization (Abdin et al., 2024). Using our broad theoretical perspective, we provide new insight into critical windows of these kinds and view our work as corroborating and extending these empirical works.
Jailbreaks. Existing work on jailbreaks has studied the appearance of critical windows in the first few generated tokens (Qi et al., 2024; Zhang & Wu, 2024; He et al., 2024; Lin et al., 2023). Our theory provides a simple explanation for when jailbreaks occur: when the unaligned component assigns a much higher probability to the current text than the aligned component, then the model is jailbroken. This generalizes the explanation from (Qi al., 2024) (see Example 4.3 for our particular formalism of their insights). It also explains the success of perplexity- Kamfonas, 2023), which monitor for low probability of the context and generation. We view our work as providing rigorous mathematical framework for jailbreaks, as well as highlighting the important role off-distribution contexts play in eliciting harmful behaviors; we also develop a novel jailbreak from our framework (Section F.2) ) similar to the adversarial contrast decoding method proposed by (Zhao et al., 2024), which also uses a likelihood ratio between an unaligned and an aligned model. However, we use a jailbroken and non-jailbroken pair of models instead of two versions of the model with different prompts.
monitors for jailbreaks (Alon &
B. Examples of stochastic localization
In this section, we present several kinds of generative models within the stochastic localization framework and their forward and reverse processes.
Example B.1 (Continuous Diffusion Models (Li & Chen, 2024)). For continuous diffusion models, the forward process progressively degrades samples X p into pure Gaussian noise through scaling and convolution with Gaussian noise. It is the
process, stochastic process (Xt)120 given by the following stochastic differential equation (SDE),
dX2- = = dt + v2dB,,
Xo ~pp
where (Br)s20 is a standard Brownian motion. Let q: & law((,) for t 2 O, and observe that as + co, qe converges exponentially quickly the standard Gaussian distribution (O, Id). Assume we end the forward process at time T 2 0. For the reverse process (XT (0,7), employ the reversal of the Omsteen- -eatein-Uhheeneckk SDE, given by
ax;- = (X;- +
V2dB,, X-~ ~ qT
where here (Bi)2o is also a Brownian motion. Defining (Yz) = (Xt)ueI. we see that the forward process satisfies the Markov property in Definition 2.2, and the information from the original sample Xo is degraded by more steps in the SDE. Furthermore, the reverse SDE with parameterized by the score function In (X,) can be viewed as successively sampling from the posteriors via Tweedie's 's formula.
Example B.2 (Discrete Diffusion Models (Lou et al., 2024)). Consider set A denoting the vocabulary and let p E AA, and consider a forward process with index set = j0, K] U {oo). Yu = X, and Y, E AT defined in the limit as follows.
p(Y+ = a Y2 = b) = dab + Qe(b, At O(AP),
14

---

Blink of an eye: a simple theory for feature localization in generative models
where Qi E ER diffusion matrices with nonnegative non-diagonal entries and columns which also Markov chain and as & -> o0, Y, is degraded until it is eventually uninformative about the original sample Yo- Example B.3 (Autoregressive Language Models (?)). Consider set A denoting the vocabulary and let E AA", and consider a forward process with index set I = {0, 1.2 .T. Yo = X, and Y; E A For t I, we let Y equal the last first T t tokens of X. Clearly this is a Markov Chain, and the reverse process is equivalent to next-token prediction.
0. (Y))er is
Remark B.4. We note if the reverse process is deterministic, e.g. ODE-based diffusion models language models at temperature 0, then there is no notion of a critical window under our framework. The initial position at the start of sampling completely characterizes the final image. For example, given a fixed piece of text and language model, truncating it anywhere in the model's 8 response and resampling at temperature will yield the original piece of text.
C. Deferred details from Section
Remark C.1 (Technicality in defining T55 Teud). For general stochastic localization schemes,
only ask that T e Stark) S {) instead of sup, inf like (Li & Siit ,Pt
{r E I: (V Sture ,Pr
8 Starx) 2 <²) and T (6) E (1 E Ty
Chen, 2024), because the sets {1 E I TV Stark @ - 2 1 e² E TV .Pr ()aak) S e) may not be closed for observation processes which are discontinuous. For autoregressive language models and continuous diffusion, the observation process is continuous, so we will clide these technicalities.
.Pt
For Theorem 3, employ the following two helper lemmas.
Lemma C.2. By applying the law of total probability and Bayes' rule, we can show for Y E supp(pz Stare)
f (IP'x)))7 P-
P=(x-'' Yx. Starg) dr 2
Σ
(017)
BEO-S
Proof. We can rewrite p* (=((x(( P- Ly, Starg) using the law of total probability and Bayes' rule.
p' (x(())
(O
8Ee
(2) Eoesing (0 (x:V)) ExEsuex (0(()
pr (z|\g,, Starg)
Σ
0EStar
Note that the second equality on the second line follows from the fact that for all O E Searg. the posteriors (-18) Pi (-1y,, Starg) by the same normalization constant. Therefore the difference can be written as
/ P (xl()) P (xlYg, Starg) dx
EoEsue (0)) (x)V 0)
=
Σ (0 (x) BEE
dx
EoEsum (0)())
Σ (= (2) + BEStur
Σ (0)) (x)Ys, 0) dz. Star
Σ EBesmes (0
BEB-
If LoEe-Swn Starg ( = 0, then the above is equal to and we are done. If it is non-zero, we can factor out Epco-Soar SSNE P. (0 term, which allows us write everything in terms of posteriors with respect to θ Starg and Starg:
1
Σ
(io(() 0)+
Σ
P: (x|Y 0)
ExESum (()7))
BES
0EB- Seace
Σ (())())))))))))))))))))))))))))
/IP
(x/ Store) - p (x| θ Searg) dx
8E0 StanStarg
15

---

Blink of an eye: simple theory for feature localization in generative models
Employing the trivial observation that
/ (x Starg) - -P- θ -Starg) dr
(o)s Starg) P- P=p(=/VV (g) θ Starg) $ 2,
we have
/ (x((()))) PT
(x P= Vg. Starg) do S2
Σ
((1)))
1E0 Stare
Lemma C.3. By Bayes's 's rule, we can derive for Yy E supp(pp).
e Stare (Y
PT
Σ
(()(()))))) I max (1, W)
lE0 Steck
PT
e Seax
(Yy) + PT
Stare (YT)
Proof. We obtain through Bayes' rule,
Eoeo Stare (Yp
Σ
Pr (0
Enco wop (V
Starx
We divide by the same normalizing constant Epce Senex
Wy to obtain
Esco-Burx (bP4((F) Eneo Sewng
Eoco S Stare (V7)
EoEo Wop (Y ()p)
Eoco Stary (Yp)
Locoror wi Enco- 80
EoEsINE
w
Σoco -Sanx Starx WD
Pî
θ Sears (Yp)
S max
1,
Eoesuay Starx W8
PT
Staze
(Vr) + +P (Yr)
Swun
D. Omitted details from Section
D.1. Diffusions
Here use alternative f f-divergence to characterize the critical windows, the squared Hellinger distance, defined as (P, Q) & f(vPP- 2 because there are explicit computations for the Hellinger distance for mixtures of Gaussians. We similarly exploit the following ratio inequality akin to Lemma 2.1,
Lemma D.1. For probability measures P, Q
dQ dP + dQ
1
1 2
(((.Q))).
ExnP
< 5
We apply the following well-known formula for the Hellinger distance between two Gaussians.
Lemma D.2. We have
Epl4//
EQ
1/4
EEp + Eq
(up -ol
HP ((),
(up [
1/2 exp
2
Σx+Eq
2
16

---

Blink of an eye: a simple theory for feature localization in generative models
Example D.3. [Two Isotropic Gaussians] Let θ {+1),
pt1 (u, Id), p (-pp,, Id). Then, have critical
window transitioning from sampling from both components to the component +-1 between T before In || + In 2 In 1/c
and T after = In |||||| - In 2c When T < Tater then TV(pt1.2 ,ptt) & When T T before TV +1.7 PP) S
Proof. The proof for hefore" a simple application of Pinsker's inequality, can be found in Appendix B.1 of (Li & Chen, 2024). Directly applying the new Master Theorem 3.1 to T Tatter we need only show that
v2
))
V2 exp( 1 211ulpe 21 Sc 2
1 - 2
H2 (ptt Pro after
2
after
2
Example D.4. [Random mean spherical Gaussians] We first sample 14 ~ (O, Id) for E [K] i.i.d. and let O [N(,,, Id) le Sbefore = θ and Safter {p1}. Then, we can compute Th = ma In [|ا: and TC min] In ||p Myl 2 8n c. Furthermore, with high probability over the selection of the means, before Th after = (1) as d + co.
aater
Th
Proof. The proof for The can be found in Section 5.2 of (Li & Chen, 2024). We need to slightly modify the proof of Theorem 3.1 so that we can write the desired bound for TC We use the same notation. By convexity, we can bound
after in terms of the Hellinger distance of individual components.
Eoco-suo (Y7)
(
E By
Yp
Seare
Eoco wop
We
Σ WolLy
WalLy
warl, (YG) + wop
OE Starg
S Starr
(1 11 See
< K
max l Star-SSarg SES
when 7 < after- To conclude the second part of the theorem, observe that by concentration of measure (e.g.. Theorem 3.1.1 from (Vershynin)) and union bound, there exists constant T independent of d such that || E [vd - T, vd T also assume that there
for all; [K] with high probability. Furthermore, by known Gaussian Suprema inequalities, exists constant T' independent of d such that ((4:H1) S T' ||s||| (Lemma 5.1 from (van Handel, 2016)). Thus, we can conclude that
max |1 - 14y12 2d 4T V +2T Odd).
JE K
max I|u4 - jE(K]
1111 1ina-11122
4T 4T 2T 2 Q(d).
The difference in log scale thus constant,
1
max || - M1l13 In
min JE(K],
||4 11p--1911 -111(²) = (1,)())
2
JE[K]
Example D.5. [Two Dirae delta functions with a random masking procedure] Let p E {+1)², and consider forward For E I, we let all the value at index i [T] be
process with index set I [0, 1). Yo X, and Y; E [+1, set to [MASKED] with probability i independently. For a mixture of two Dirac delta functions, we can express the critical window in terms of the Hamming distance between the corresponding strings. Let θ {0+1). L+1 E {+1)), pf= 1 Then, on component we have the critical window
de De W11 2.
In(1 -c) du (oe, Se_.)
In dd
The exp
TC after = exp
When T < T.C. then TV p²) S.. When T 2 TCC Tefore T. .p) S e. For sufficiently large du (ōe, 8t_,), the window size Thetore TC 0 d (( If du (8e, )e_,) increases with T, then the width of the critical window compared to the width of the index set becomes negligible.
after
17

---

Blink of an eye: a simple theory for feature localization in generative models
Proof. To prove TV p²)) & e when P S afters observe that when î S
T after' the probability that all the differing
elements between &,, 1-1 are masked is exactly Tidwlo,, .6__) I 2.2 That means that there exists a set A with Ps ! (A) 2
1 e and p7 (A) = O, so by the definition of total variation, TV .P7') 2 - è². Obviously, TV . = 0 as
well, so by Theorem 1. we obtain TV(p'.7 .p') 2 To prove that TV 1. .p(=)))= S c when T 2 To
+- before. we need only show that " .P) I By Lemma 15 of (Li & Chen, 2024), it suffices to show that ( PT S E by a simple
triangle inequality argument. Consider the set C (+1, [MASKE]]* o such that TV ( .pp - Pj (A) - PT 2 (A). Consider the set B supp(r?) n supp(pz") ) For any E B, we know P7 (x) - P7 (xx because the same number of tokens need be masked from f+1. This means we have (B) P7 (B) 2 (Tduletie)) 2 1-16 6. Because
1 1 (A B), we have p (A) p (A) S p ((=1, (ASKEE/A - B) S 0 Pf +1 (A) P (A) p (A B) pj
1.0
Prob consistent with original
0.9
0.8
0.7
0.6
Tattrl 0.2)
Toetore
0.5
10-2
10-1
10
t
Figure 6, Plot of critical window for discrete diffusion on delta distributions with beforer T after plotted.
D.2. Autoregression
Example D.6. [Math problem-solving as a random walk] We model solving a math problem as taking a random walk on Z with stepsize of length T. If the random walk hits +A, then it has 'solved' the problem; if the random walk hits -A, then it has obtained an incorrect solution. Assume that we have two modes: strong problem solving mode (denoted +1), which takes a +1 step with probability 0.5 + A, and weak problem solving mode (denoted -)), which takes a +1 step with ln(2/e2)
probability 0.5 - A. Assuming that 24 < A and e² < 10 (0.5 A) + A), there is a critical window for the p
strong problem solving window of T before = T - 2 independent T.
A + 2 and after T In(2/e²) 2A
Note the critical window has width
Proof. Because only the direction of steps matter, we can model the critical window for this random walk as observing a sequence of 11 with an autoregressive language model. Let p E (+1]T, and consider forward process with index set I I= (0, TT, Yo- = X, and Y; E {+1, [MASKED]]" For t E I, we let the last t tokens of Y, be deterministically set to [MASKED]. We generate data as a mixture of biased coins with separation 2A < 0.01. For a mixture of two biased coins, with probabilities of 0.5 = A (0+1 respectively) of yielding 1, we can compute the critical window and show that it tightly clusters around ((1/2)). Let θ (0_1), p (Bern(=21))87 w=1 = 2. We also assume
2 10 3 (0.5- (0.5 A). Then, component 1 we have the critical window TC Tefore = T x= + and after then TV 1. ,r') S e. When I 2 TC Toe TV ((pp"..))s Sp)
TC T In(2/e") When 7 < TC
after
24
Note that the number of +1 is sufficient for disambiguating 011. To prove the bounds To we show that with only samples the total variation between 0.5 A and 0.5 + A is negligible. Using (Roos, 2001), we find
2
A=
T T-T2 after +2 2(0.5 ()(55++))
2A
TV .Po )
2 S 3c.
after
after
T-TC +2 2(0.5- ()00.++a
2A
18

---

Blink of an eye: a simple theory for feature localization in generative models
For TC
after We compute how many samples it takes for p=t have only overlap in total variation using Hoeffding's
inequality. If we have n samples, the mean X of the n samples of +1 for p 0.5 + A satisfies the concentration
inequality P([X- 2A > A) < 2 2exp(-2xdi (furthermore we can ignore the stopping condition by our requirement
that In(2/c2) 24
< A). We find P([X 2 > A) S 2 for T T Tatte
after samples, proving that the total variation is at least
1
Now, we consider a model for autoregressive data similar to the one presented in (Arora et al., 2019), Each word a vector E Rd and the context length is T E z>0 The original samples are x RTxd Let θ {u. e), where u, v E Sa-1 We define the distribution p" for 0 @ as follows. We generate the path of discourse vector (C (E(uxe.))) RRd with the
reverse SDE Distein-Uhlenbeck process such that Ch (O, Id) and C8 -V(0, Id a00 for some a > O. We let of be the law of CJ for & O. We let I = 10. T), and for t E I, draw samples w/ I where impose normal Gaussian prior and have w.ICy, & exp((C), -)). Then we return the corpus {w, hhe as an output.
Theorem D.7 (Autoregressive with a mixture of two Gaussians as the concept distribution ). We assume that (u.v)///) Then, on component u we have In
exp(-T) - log(1 + aa S e. Let o(x)
(2-z(1+ 2-8
av VI (2, v)2
a after 2 In
Telore 2
0 (c2)
When T < TC after then TV(ptu) pput) S. When P2 The
TV(plu). p) S..
This proof will require Theorem D.8 and Lemmas D.2 and D.9, which is stated below:
Theorem D.8 (Section 5.2 of (Chen et al., 2023)). Let (Yo)eetor) and (Y))eEjo.)) denote the solutions to dy. = (V,)) dt + V2d AY: K,(Y)) dt VdB.,
Yo ~
Yo ~9
Let q and o' denote the laws of Yr and Yf respectively. If be, br satisfy that So EQQ (Vb) - 1((Y,)l12d < co, then KL(glla') S So EQ (Y,) - (F
Lemma D.9. Let u, v E 8d-1, Then
| -wv 1lop S /1 1 (u, u)2 AuuuJ T + vv -v= =1 + (u, z))z)).
Proof. There exists r E Sd-1 such that v == (u, u)u + /I - (u, ur r and t 1 r. We find that
ww - vu = (1 (u, u²²)uu" - (u, )rr - (u, v) v11 - (u, u) + tru
We can explicitly compute the eigenvalues of uu" T To using the discriminant and find that they are equal to + #vI (u, u)2, By similar derivation, we can write
un + uou (1 (u, v))))u + (1 (u, 2Tr + (u, 1) V (u, u) tru
which gives us eigenvalues for wu + vu of 1 + (u, v).
Proof. To compute the after bounds, we compare the difference in Hellinger distance of the distribution of words words at generated at index T. w... By the data processing inequality
(p S 1 - 1 !1 (w w so it suffices to
show 1 - 1/2 2 (w wog)) $ c. Because the Gaussian is its own conjugate prior and w CP x exp( 2 1 - CfIl²). we can compute w C NCC,, Id) and w NO 21d + de 2u Applying Lemmas D.2 and D.9, compute
explicitly
2 ae 29
27
1
= H2 H2(
a w S V (2+ ae -2 27 21 (u, w))/22(2 + ae- -f u))/2) S V Volae-2t) S
2
19

---

Blink of an eye: simple theory for feature localization in generative models
To compute The beforee first the data processing inequality to reduce the difference in the emitted tokens the difference in the paths of the context vectors, and then apply the approximation emror bounds from Theorem D.8 bound the differences in path measures. When f 2 before we can use the triangle inequality to write TV .p) = TV Pf (u))) S
TV P.). Note that P/ is the distribution the first T - T tokens generated by the model under 6. Note that p is function of (Ct)ee ( and P.J is function of (Ci |igin(f .7. By the data processing inequality, we can bound the difference in terms of the distributions over the tokens in terms of the law of the process of the discourse vectors,
(eennt) (Ct TV P7) S TV((CT) Leintyy
TV((CT) (e) S
)ee (CT) (E).))
Note that for 0 E O, (Ot) [e(T, .r] is generated by the following reverse time SDE,
dc. = {C) 2V In (C,) di + v2dB,, I Le (T, ~ sh.
Now we define (C, hE(T .T] to be the reverse SDE defined by initializing at o7 but with the score of q.,
dC (or- + 2V [n dt + v2dB..
t [T, C 9%
By the triangle inequality, we have TV((CT))e(ft) (CT)ee(fm)) S TV((CT) (e (Ct"") (
TV
((C
leff my (Ci) (C))e)(.)))
(I)
To bound (I), observe that the SDEs have different scores but the same initializations. TV((CT) be) (e( and obtain
We apply Theorem D.8 to
TV((CT) (Cr"")e S (Ll((")) ")EE(TT))
V Ex-c11VVlptt V (x)
We simplify the inner expectation by using the V In (X) (lnpp(())))) 11 [[d + ar -21 "u 1) (Id ae -2 - xel
=I[(Ta Id + ae 21 un
-2
+ ae 21 vo ) x1
ae -21 21
Id
vv S ae 2t 21luuu
llop
II
span(as,)
ae 21 V1 u)2
(Lemma D.9)
We can upper bound Exworl II o) X|I2 by considering right-triangular L such that L L Id + auu" T. and II [span(u.)) X = IT [ppaanuu.)) LY, where Y ~ (O, Id). The operator norm of II span(u..) L is
II LL
< |bllop I 11 vvI=
(1,
Ispan(uu LP is also rank 2 and (Ispaniu.e)) ((0, Id2). Thus we have
)b o Y = (Isspn(u.o)) o1) o sp))(
I Y. where II span(uu,)) Y
27
2
By~N(U.1a)
Ll Y
E Ey'~A(o,dd))
II Ll Y (uan(u.))
S 1 + a²
span(u,.))
20

---

Blink of an eye: a simple theory for feature localization in generative models
Combining this information together, we are able to compute.
TV((C)) (Cr )ettrr) (Cr"") 1/ J
Ex~c; IV V In pi
0² - 27
To bound (II), we observe that both are run with the same score so we need only bound the difference initialization. By the data processing inequality, we again have TV((C, "ee (((T))) <VVpp We can again apply the triangle inequality to get TV (ni, pt) S TV(p,, >d) TV(7", . S. For any 0 E O we have by the forward convergence of the OU process TV 7n)) (xpp_-T)yLLpppll>""" We can explicitly compute (T. as
KL(Pulh"*) dd+as a - d- --lg + awn T]]
1 la log(1 + +a]]
2
Thus, we obtain the following bound on TV((C++*) "eE(
)E((im) .7. (CF (G))))().))) of
TV((CT " )ee(f.rp (CT) (e))))))))))) S exp(-TT) log(1 a) Se.
D.3. Interweaving fransitions from other distributions
In this section, we extend our critical windows framework to the setting where certain steps of sampling procedure, instead of using the reverse Markov transition kerne from the original stochastic localization sampler, we an alternative distribution which is not necessarily related the original sampler. This includes many important applications of generative models, in which one seeks to combine the priors learned from data with some other algorithm. For exam- plc, one may want to combine the language model with problem generation oracle in in-context learning (Dong et al., 2024).
As (Xie et al., 2022) points out, the transition from the answer to one problem to the problem statement of another example in-context learning is determined by an alternative transition kernel (which they call Pprompt). Although the probability of transition from answer the problem statement of another example is extremely low under the natural data distribution, one still hopes that with sufficiently many samples, the model selects the correct 0 E θ if these lower probability transitions are overcome by the distributional difference for 0 E θ with 0 o 0" Similarly, under our critical windows framework, we can hope to capture the idea that we specialize to particular 0* given a sufficiently long context. In Section D.3.1, we first present a general framework for characterizing critical windows in this setting. Then, in section D.3.2, we consider the case of in-context learning by autoregressive language models and prove convergence.
1. GENERAL INTERWEAVING FRAMEWORK
We present this framework for the case where the index set I {0, 1,..., m) is discrete. Like before, assume have series of reverse Markov transition kernels P - (-I.), for k E I, but we also assume we have an alternative distribution Praal (+I-) that we use to sample for transitions k E A Ç I. For our sampling procedure, we sample Ym.. and for wilt (-((++1) for k E A. We denote
== m - 1, m - 2,..., O, we take Y4 ~ Piti (-1++) for & E I - A and Y2 ~ the final distribution all..
Now, we also need to adjust our definitions of p to this particular sampling procedure. We define S.alt for S C θ to the distribution over outputs when we instead use the kernels P (-1Y&+1. instead of P (- (Y ). To relate paalt to o for 0 E O we need to assume transitions from alt do not affect the posterior distribution over werpp" (01Y.)
Assumption D.10. For all y E IR and x E supp(PPritl (-1g)), we have for all 8, A' E e, the equality P P 9) = Pit1-+ (xly,,
21

---

Blink of an eye: a simple theory for feature localization in generative models
Adopting definitions from Section let
T alte E E (P Siitsaalt Stargaat) S [(}
T aupper,ale E E TV Stoggaalt @ S 2 21-
Ptr
The main challenge the below corollary is simply show that the final distribution p po,ult can written as a mixture of plaat with the same mixing weights as before.
Corollary D.11. Under Assumption D.10, for O, if 2 alt and T S T Tupper.ate(e), then
T p srrrealtt) S 1 1 222xxxxx (1, Loco-smo Stare We
"EoeSrare Wg ) /2 c.
Sma
12)
Proof. We need only show that p g A Ym, Ym.1...., Yo are the same under both density functions. We need only consider transitions for k E -I-A A, because for E A, the transitions are both given by the alternative distribution. For the transitions not given by alt, note that we are using the original model, so
Bee wop&.alt It suffices to shows that the probability of generating a path
g (Yx-1 Eoso wep ( (Yx-1 Yx,
Eoco war)(Y)
Furthermore, for the mixture model, this probability is
mix Eoso waph.alt (Ymn Ym-1,..., YK+1, Eeco allp (Ym, Ym-1..... Y YK+1,YA
The distinction between Equation D and Equation D.3.1 is that in the former we are using the likelihood of pf instead of phaat Thus it suffices to show that p" pfialt We explicitly write the probability.
m pº (Y V........ Ye,,, = II i=++1
P (Y--
x
II
p (Y,,,,
where the proportionality follows from the fact that we can ignore the probability of the transitions produced by alt under Assumption D. 10. By definition, this is proportional to pa (Ym: Ym-.,.... Ye+1, Y2) up to a normalization constant independent of
D.3.2. IN-CONTEXT LEARNING
Now, we will specialize our framework to the case of in-context learning. As in (Xie et al., 2022), we assume that the language model is given inputs of the form [x1, 3/1 o, Ø2 //2 o,..., Tt: Mr, O, Mr+i], where æ is the input, y1 is the output, and o is a delimiter token that separate different in-context samples from each other. We assume that the transitions y: 0 Ti+1 sampled by some alternative probability distribution Pother (-I-). We require that Pother selects the 2; i.i.d.
Assumption D.12. The distribution of Pothe((2t411) 91. O, 22, y2, It, Vt. o) Potter (x1).
Then we assume that the transitions x; + ys are generated by some 0- O, which does not depend on any of the previous tokens before the delimiter.
Assumption D.13. (Well-specification) There exists some 0" + such that yi is generated from Vi p- (-1[21, V1, 01...., 0:-1, 222]),00 2;,1, 0*).
Assumption D.14. For all 0 E O, we have p- (-1[21, y1, 01, 0,,,
mr).0) ar,) p (-121, 0).
We also assume statistical separation of 6* from θ {0*) in terms of Hellinger distance. Assumption D.15. Let p for S C O be the distribution of (x1 V1, o) where 11 Potner(-lo) and y1 ~ (p(1-)))) There (0") 0-10"' ,(Ty))) 2 6.
exists 0 such that (p (m.))
22

---

Blink of an eye: a simple theory for feature localization in generative models
Example D.16. Let T 2 In 6//
Under Assumptions D.10, D.12, D.13, D.14, and D.15, we have
IT+1, vT+1 [#1..... -T Ur++1 S e/wg-.
Proof. It suffices upper bound 1 HP (P (B-())) by O(c). First observe that the distribution pSsalt for S C e factors along the delimiters by factor independent of S using Assumptions D.12 and D.14, so we have
3T
pla (11 V1. 01, 22, y2, 02. ....... 2t, Yt. or) & III
S.alt (ot, Br, o.).
i=1
Using the tensorization property of Hellinger distance and our definition of T, 8, we have
1 - H (P
1
pfat
PSt (
x - 112 (Pg P23 p (0) 1"
Sc.
2
D.4. All-or-nothing
D.4.1. ALL-OR-NOTIINNG PHENOMENON
Here we elucidate a formal connection between the critical windows phenomenon in in-context learning and the all-or- nothing phenomenon. To begin, we first define the notions of strong and weak detection:
Definition D.17. Let (N,) be an increasing sequence of integers. Given sequences of distributions (p.), (q.) over RN,, sequence of test statistics (A, IRN. IR) with threshold (Tx) achieves:
strong detection if lim sup, p.. (Ao(2) < Ta] Pr Prong (Ao(2) 2 Tz]} 0.
weak detection if sup, (Przmp, (A.(z) T.] + Prz (r. rz]} <
By the operational characterization of TV distance, strong detection is cnfformation-theorriiii possible if and and weak detection is (information-theorreiclllll possible if and only if
only if lim inf, +000 TV(p..
lim inf TV(p. 9e) 0.
Now we consider the following Bayesian inference problem, given by joint distribution т over (0, E RR x R". Nature samples unknown signal 0 E R. from TTg Given sample size N, we receive observations { drawn i..d. from zajoi the goal is to infer 0 from these observations. Let (N) denote the distribution over {2, J.N the mixture of product measures parametrized by B.
Definition D.18. Let (T,) be a sequence of inference tasks over TR R. and mullll) be sequence of distributions over Rm, (m.) exhibits an all-or-nothing phase transition at threshold (N.) with respect to null models (ulll, if:
For any < 1: weak detection between (B()))) and (( ) is infornation-theorcticlll theoretically impossible
For any 1: strong detection between the planted model (अr (8N.)) and the null model ((rr joo8N..) is information- theoretically possible
All-or-nothing phase transitions have been established for a number of natural inference tasks like sparse linear re- gression (Reeves et al., 2019; Gamarnik & Zadik, 2019), sparse PCA (Niles Weed & Zadik, 2020), generalized linear models (Barbier et al., 2020), group testing (Truong et al., 2021; Coja-Oghlan et al., 2022), linear and phase retrieval models (Scarlett& & Cevher, 2016; Truong & Scarlett, 2020), planted subgraphs (Mossel et al., 2023), and planted Gaussian perceptron (Niles-Wear & Zadik, 2023). Here is an example for sparse linear regression:
Theorem D.19 (Reeves et al., 2019)). Let n, be the distribution over TR TRm, for n, = s and m, = 8 where the marginal over 0 is given by the uniform distribution over ky--parse vectors in {0,1}* 1},, and the conditional distribution π 22 given by sampling x. (0, Id,), taking y (0,x) + for N, o2), and outputting observation (x. y) The null model null is given by sampling
(0, Ide) and outputting y (O k. + +?))
If o² < K2 S 0.499, then (T,) exhibits an all-or-nothing phase transition at threshold (N") with respect to null models (ITBull) for No & 2k. log(s/k.)
log(1+k,/o2)
23

---

Blink of an eye: a simple theory for feature localization in generative models
Having defined the all-or-nohhing phenomenon, we rigorously instantiate it as a critical window for in-context learning. We first define a mixture model P(N) for sequence lengths N onto which we will identify a critical window.
Definition D.20. To any inference task т,, null model =r"ll, and sequence length N, we can associate the following in-context learning task. Let θ signal LI (NULL) where signnll _ supp(<o)) Given 0 E supp(ro), let p(x) denote the distribution over sequences (Z1,..., AN, ?, where 2N are i..d. samples from P=10. Let P(N) denote the distribution over observations (z1,..., 2N,? ?, NULL) where ,1,..., FN are i.i.d. samples from n n We then take p p - B fo 1o P(N)
Under this model of data, we have the following theorem expressing the all-ornnothing phase transition in terms of before: TC
Theorem D.21. Suppose (T,) is a sequence of inference tasks that exhibits an all-or-nothing phase transition at threshold (NÇ) with respect null models (xg). Given N₂ 2 N2 let ((N)) denote the sequence of in-context learning tasks.
For any constant 0 6 < 1, there exist constants ö, 5 such that for all 8 next-token prediction for (P(N)) )oes, exhibits a critical window over [Ns + - (1+ 6) .N. + - -1- N in which we transition from sampling a distribution Oclo-c TV Shefore θ s:signal to sampling from a distribution -cco-l TV Satter = Ox
In other words, we have Thefore - N + 2 (1- N) and Tattrr NN+ + 5)
The proof of Theorem D.21 is essentially immediate from Theorem 3.1 and the definition of the all-or-nothing phase transition:
Proof. Let us first apply Theorem 3.1 to Siit Starg s;signal- By the definition of De the parameter W therein is 1. Furthermore, we trivially have that (e) 0. Finally, because strong detection is possible provided there are N 2 BN; in-context examples for B >> 1, there exists 51 depending only on e for which TV Suiz .Pt 9, SSuxx) 2 1 e for f = .N + (1 + Nj).NN By Theorem 3.1 we conclude that
S e. Next, let us apply Theorem 3.1 to Smiit θ Dsgiinall and Starg = O,, The parameter W therein is now O. Furthermore, we trivially have that Tend (e) N₂ + 2. Finally, because weak detection is impossible provided there are N < 8N,"" in-context examples for .Pr Stars) I c for t = N, + (1 -1-60 - By Theorem 3.1 S e. Taking 8 max( (1, 52) concludes the proof.
1, there exists 52 depending only on for which TV Siii
we conclude that TV
E. Details from Section 5
Definition E.1. An e-mixture tree is tuple (T, (--(--))),, O, Joce, Subset, NoiseAmount). The tree T (V,
is associated with a function Subset: V 2 (0), which maps vertices to sub-mixtures. We require Subset satisfies the following two properties: (1) Subset(root) = θ: (2) If u is a parent of v, Subset(w) C Subset(u). We consider a NoiseAmount: V + RR20 which characterizes the noise levels that result in the aggregations of mixture components described by vertices in the mixture tree. Thus we require that Aoiutt satisfy three properties: (I) For distinct 0, f θ with leaf nodes w, such that 8, E Subset(w), O5 Subset(r), if is the lowest common ancestor of w, v, then we require (P (PrssseAmmuun( > Priisee Ammunt(())) ..)) S c (2) For u E V, we have statistical separation between Subset(u) and e Subbserr(())
O Subset(u) in terms of TV, TV Subect(u) (NNoscAAmonntu (u) NoiiscAmuun(()) 2 1 e²; and (3) If v E V is parent of we have Amoun(n)) < Noise/mount(p)) Property establishes bounds on TS and. and properties 2 and 3 establishes bounds on TS.
Corollary E.2. Consider an e-mixture tree. For 0; E O consider the path u1 U2, Uz Ug...., unr E V where un is the leaf node with 0g E Subset(ur) and ugr is the root. There is a sequence of times Ti < Ta < ... < Th with TVppli.Teppuubrttn pSubeer(ue)) Sw c.
Proof. For E E [H'], we let T2 NoiseAmount( uz). We apply Theorem 3.1 with Smit { and Stare Subset(ue). We Stare G-SweC) 2 1 e² by Condition in Definition E. By Lemma 15 of (Li & Chen, 2024), we know (i) ). This is & e for all j E Smit by Condition 3 on Noiut and the data
know TV .Pr.
Stare,
TV PTe
(i) Smmtt < maxjeSia
TV
processing inequality.
24

---

Blink of an eye: a simple theory for feature localization in generative models
Corollary E.3. Consider an e-mixture tree (T, (P(P-(( (..)).)),,I (pe Jeea, Subset, NoiseAmount). Suppose we have an- other distribution such that TV (pp -q)) S for all 0 G. Then we have Vo-mixture tree given by (.AP-(1-)), 0, {a" Jeee, Subset, NoiseAmount).
Proof. We need only check the first and second properties of NoiseAmount with parameter + v5. To do this, it suffices to
9, show TV TV(Gaasee (Naise/mmumm( Amiisem))))(()) I e e-- and TV (4NaiseAmsum(()) Noiieemmmumt( (m)) > 1 1-e2 By the data processing
Subset(u)
@ -Subse((()
inequality, we just need to show this at = 0, and we prove the stronger statement that for S C 6, TV(p g - 6/2.
This follows from Lemma 15 of (Li & Chen, 2024) and TV(p", a4" S 6/2 for all 0 E e.
Example E.4. Consider set of alphabets (A.)=-1 and define e {(a))=1 Vi E [d] A. and p óe. Let I= [0, 1,2, ..., d]. and for any permutation f,, da of [d], define forward process such that at t I, mask all id, ia 1.... ................................ da-r. This constructs a hierarchy where the values for i,, i2...... da are decided in that order.
Proof. We construct the following O-mixture tree as follows. We let the leaf nodes be the set θ We let two leaf nodes u, u have the same parent and only if they share the same values on the alphabet at i,, i2, ..... dd-1; we also define the parent as the union of all of its children. We now treat the parents we constructed as the roots, and let them have the same parent if and only if they share the same values the tuple i,, i2; ..... ia-2- We continue to do this until are left with one root node. We let Subset map each node the corresponding set and NoiseAmount map each node to its distance from a leaf node.
By the construction of T', it is clear that Subset satisfies the desired properties. For distinct 0;, 05 E O the lowest common ancestor of 0; 0g represents the largest & such that indices i,,..., ig are the same for 0,, 0j. Because 8. (NoiscAmount()) is just
02 the tuple of the values of 0r, 0g 21,...,ii ik. we know TV (p (PNaiseAmomm(1)) (NoisscAmmun(( (u) the values at index (ie) all 0 & Subset(u) does not share the same values at these indices by definition, so we also know
0. For any V representing
9- 9-Subset(u)) u) (NsssAAmuum(()) (1) PNoiseAmount(u))
Subset(u)
TV
F. Jailbreak Experiments
F.1. Reproducing critical windows for jailbreaks from existing papers
Existing work has already identified the presence of critical windows the domain of jailbreaks. Here we present critical windows for simplitied prefill jailbreak based on the prefill attack (Haize Labs, 2024b) and repeating token jailbreak (Nasr al., 2023) for LLAMA-3.1-8B- Instruct.. In the first figure, plot the probability of the model giving harmful response, computed using the StrongRe Gemma 7b auditor from (Souly et al., 2024), as function of the fraction of the phrase Sure, here is how to appended to the front of the model's generation. We can see that there is large jump in the attack success rate after only including a few tokens in the prefix. The second figure is reproduction of Figure 12 from (Nast et al., 2023). It shows that the probability of repeating the next token increases substantially as the first few tokens are included a few times in the context.
25

---

Blink of an eye: simple theory for feature localization in generative models
Prefill attack
Repeating token behavior
0.6
0.8
0.5
Probability token is repeated
Attack Success Rate
0.4
0.6
0.3
0.5
0.2
0.4
0.3
0.1
0.0
0.2
0.4
0.6
10c
10a
10a
Frac of prefix included
token repeated context
Figure Examples of critical windows in jailbreaks for LLAMA-3.14 1.MA--311-BB 1-0A--..-SB-Instrucct Left: Critical window prefill jailbreak (Haize Labs, 2024b). Right: Critical window for repeating token jailbreak (Nasr et al., 2023).
F.2. Experimental details from jailbreak
Now we apply our theory to develop a new jailbreak detection method, based on a likelihood ratio between an aligned and unaligned model. Intuitively, our theory states that when the unaligned component assigns high probability the text compared the entire model, the model is likely be jailbroken. We use LLAMA-3.1- 1.---..18B-Inttruue model jailbroken with LoRA to not refuse harmful prompts (grimjim, 2024) proxy for the unaligned model. We evaluate these different methods on a dataset of jailbreaks and benign prompts from (Bailey et al., 2024).
Dataset. We use the same dataset as (Bailey et al., 2024) but provide details here for completeness. The benign dataset consists of inputs from UltraChat (Ding et al., 2023), a large dialogue dataset, and Xstest (Röttger et al., 2024), which contains benign queries that are often incorrectly refused by language models. The benign queries are filtered to ensure that LLLMM--3.--BB-Imee
does not refuse any of them. The dataset of harmful prompts based off the Circuit Breakers dataset (Zou al., 2024). The datasets include the following jailbreaking methods from the extant literature: PAIR (Chao al., 2023), AutoDAN (Liu et al., 2024), Many-Shot Jailbreaking (MSJ) (Anil et al., 2024), Multi-Tum Attacks (Li et al., 2024; Haize Labs, 2024a), Prefill, GCG (Zou et al., 2023), and other Misc. attakes from (Wei et al., 2023). For each jailbreaking method, it is applied to a prompt from the Circuit Breaker dataset and evaluated to see if the generation from 1---AA-.1--BB-nsstee is helpful and harmful, as determined by the StrongReject jailbreaking classifier (Souly et al., 2024)).
Evaluation Metric. As is standard in the jailbreak detection literature (Bailey al., 2024), we report the recall at the false positive rate at 0.01.
Table 2 displays the recall and several other basclines. Crucially, the log likelihood ratio methods does obtain recall > 1 for 5 different categories of jailbreaks. While our methods do perform worse than existing methods, it is important to note that they still work and that their poor performance could be explained by the fact that we have to use a proxy for the unaligned mode of the model.
26

---

Blink of an eye: a simple theory for feature localization in generative models
Table Recall (FPR=0.01) for likelihood ratio threshold, perplexity threshold (Alon & Kamfonas, 2023), and ML.P-hased detector trained on activations (Bailey et al., 2024) for predicting different jailbreaks, prompt/gen denote the logprobs of the prompt and generation.
respectively.
AutoDAN
GCG
Multi-Turn
Misc
MSJ
Pair
Prefill
log prompt
unaligned
log P prompe aligned
aligned
0.000
0.000
0.028
0.000
0.063
0.000
0.077
log Pgen
unaligned
log Pgen
0.082
0.030
0.000
0.100
0.000
0.061
0.051
log
proomp
aligned
0.000
0.576
0.056
0.063
0.013
0.000
0.077
alyned
log
Pgen
0.205
0.150
0.570
0.200
0.006
0.015
0,416
MLP
1.00
0.956
0.873
0.663
1.00
0.833
1.00
G. Structured output experiments
We have LLAMA-3.1-8B- Instruct (default sampling parameters of temperature of 0.6 and top-p sampling of 0.9) respond following prompt, which asks it answer series of fill-in-the-blank questions in structured format. We prefill the model's generations with \n In)nn 1. ensure that the outputs comport that format. Figure plots the probability of obtaining the same answers as the original generation after truncating different amounts from the generation in the forward-reverse experiments, computed with 10, 000 generations from the prompt. Our theory predicts that jumps in the probability will occur at Teforre after which represent when the model has committed to particular answer to a question in the generation. To compute hefores afterr we look at when the generations diverge based on the first occurrence the identifying information. For example, the TC before of the first critical window is 1. The because the first answer has not appeared in the generation, and the T after the first critical window is The or The N, because that uniquely vindicated empirically, the jumps in probability, representing the TÔ
identifies the answer. These theoretical predictions model localization to particular set of answers, occur exactly at T before: after in Figure 8.
Structured Output Prompt
Complete the following by choosing only one option for each blank. The options are provided in parentheses, and your response must match the exact case and meaning of the chosen option. Respond with only the completed sentence, no explanations additional text.
1. The (Pirate/Ninja) jumped across the ship.
2. She adopted (Dog/Cat) from the shelter.
3. The (River/Bridge) sparkled under the sun.
4. A (Dragon/Knight) guarded the castle gates.
5. He ordered (Pizza/Sushi) for dinner.
27

---

Blink of an eye: a simple theory for feature localization in generative models
Critical
generation
PCAKP
Rax consittont olth orggmnl UPSAND
0.6
nanariig
Critical windows for structured data generation
Crttical windows for structured data generation
1.0
Toefore
1.0
Toefore lafter
P.CKK.)
biter
FCREP
FDR.TT
0.8
0.4
0.4
P.O.E
0.2
0.0
0.6
0.8
1.2
0.0
0.2
0.4
0.6
0.5
3.0
Fac gereretionn remaining
Frx generation rena
Figure Structured output plots LLAMA-3. -LAMA-3.1-BB-nstruu denotes that we are sampling from responses whose answer to first question was Pirate; P.C. denotes that we are sampling from responses whose answers the first two questions were Pirate and Cat, respectively. We can see that the critical windows directly correspond to our thcoretical values for before: TC
after
H. Chain of thought experiments
H.1. Experimental details
We select different math and reasoning benchmarks which performance is known improve with chain of thought reasoning (Lanham et al., 2023): ARC Challenge and Easy (Clark et al., 2018), AQua (Ling et al., 2017), LogiQA (Liu et al., 2020), MMLU (Hendrycks et al., 2021a), and TruthfulQA (Lin et al., 2022) multiple-choice benchmarks and the MATH benchmark from (Hendrycks et al., 2021b).
We describe the prompts we used in our evaluation of different datasets. The system prompt for all datasets is Produce correct solution to the following /TASK/ question., where /TASK/ is the type of question of the dataset, i.e. science, math, logic. For each question, create user prompt by appending Think of the /TASK/ question thoroughly step by step. Please only respond with the answer after reasoning thoroughly. in front of the question. Once the model completes its generation (max generation length set to 2048 and default sampling parameters), we append the user prompt Given all of the above, what's s the single, most likely answer? Your answer should have the format. "The answer is ANSWER", where ANSWER is your answer. for the multiple choice benchmarks and Given all of the above, what's the single, most likely answer? Simplify it completely. Your answer should have the format "The answer is $ANSWER$", where ANSWER is your answer in LaTeX. Note that when ask the model for the final answer, set the temperature O.
For 400 questions from each dataset, we prompt the language model with the query, cut off the generation at different percentages of completion (in tokens), and regenerate it 100 times. We compute the probability that the final answer is the same as the original answer, using a direct text comparison for the multiple choice benchmarks and the prm800k grader for MATH (Lightman al., 2023). Across three different models [LAMMA-3 -.---..---I--nttee- IB-Inrtu Phi-3-7B 1---BB-IIsttret and Qwen-
28

---

Blink of an eye: simple theory for feature localization in generative models
Table Differences between Accuracy (Ace) without versus with critical windows and frequency of critical windows (CW) when the
original generation is wrong versus right for LLAMA-3, Instrucct 1en-33-B-Instuue and Qwen Qwen-2.5-77 2.5-7B-Instruct
Dataset
-LAMA-3.1-BB-nnsttt
Phi-3-7B-Instruct
Qwen-2.5-77
2.-n-225-B-Inttruue
AAcc
ACW
AAcc
ACW
AAce
ACW
AQUA-RAT
0.42
0.20
0.36
0.16
0.03
0.01
ARC Challenge
0.53
0.22
0.50
0.24
0.38
0.11
ARC Easy
0.73
0.26
0.28
0.13
0.40
0.07
LogiQA
0.15
0.11
0.21
0.19
0.23
0.11
MATH
0,41
0.33
0.36
0.33
0.46
0,29
MMLU
0.45
0.24
0.38
0.21
0.26
0.11
TruthfulQA
0.44
0.20
0.42
0.23
0.55
0.23
2.5-7B-Instrruct, we find that conditioned on the occurrence of critical windows, the model generation is significantly less accurate compared to generations without critical windows.
Oriical
0.8
0.2
ABC Dhabange
Matro
Lapttar
Meto
Hinu
Debset Suace
Nois At2
Challngr
Delad Sexce
Figure Left-to-right: LLAMA-3.1-81) 1-8B- --Sttrrrru,,,, Mhi-3-7B-Instrun and Qwen-2.5-7B -I--trutt barplots of original generation accuracy for generations with (2 0.5 jump in probability) and without critical windows (<c 0.5).
29

---

Blink of an eye: simple theory for feature localization in generative models
1.-
02
0.2
0.0
00
0.3
MUALU
1ATH
ATC Easy
1.0
10.
2.-
U12-
Pac geniratinn samaiing
Dfc gENAEE Hamaunng
Tra gererEECE ESTATII
Figure 10. Probability that the answer is the same function of the percentage of the generation remaining for different math and reasoning benchmarks for LLAMA-3.1-88 Instrucct. Highlighted are generations with 50% increase and no 30% decrease in the probability step.
ABC Chall enge 1WTO-1.
Ter 1WV--1.
AQUAANT 1812-10
LOCRA JNTD-1L
3amp
Figure 11. Probability that the answer is the same as function of the percentage of the generation remaining for different math and reasoning benchmarks for Phi-3-7B- -n-sruuut.. Highlighted are generations with a 50% increase and 30% decrcase the probability over one step.
30

---

Blink of an eye: a simple theory for feature localization in generative models
24
104111 107H
J40. Eway
Figure 12. Probability that the answer is the same function of the percentage of the generation remaining for different math and reasoning benchmarks for 2aenn-2.5-B--Ittruree Highlighted are generations with a 50% increase and no 30% decrease in the probability over step.
H.2. Examples of critical windows
We provide few more examples of critical windows (bolded).
MATH example (LLAMA-3.1-BB-Instte)
Q: Let and be the solutions of the equation 2x 10z + 5 0. What is the value of (2a 3) 6)2
A: To find the value of (2a (4b 6), we need to first find the values of a and b by solving the equa- tion 2x 10r
We can use the quadratic formula solve for a and x bb vb 4nc 10, and
2a
In this case,
2,
Plugging these values into the quadratic formula, ger:
10+ I 10=2vI5 So, the solutions are a 10+2VI5 x -- 10) (- -1))5_ x 100 I
2(2)
and b 10-2v15 Now that we have the values of and b, we can substitute them into the expression
(2a (4h 6): (2a 34b 6) = (2 10+22115) (4( 10 2/15 6) Simplifying the expression,
4
we get: (2a (4b - 6) = (10-2/15 3)( 10-2v11 6) (2a 34b 6) (10+2 10 -00-2/1--
(2a (4b 6) (4+2/13) 4-2V15) Now, we can simplify the expression further by multiplying the Using the difference of squares
numerators and denominators: (2a- 3)(4b 6)
formula (a (a - b) a² b², we can simplify the numerator: (2a 3)(4b 6) 42 42-(2v15)²
16-60 Simplifying the expression further, get: (2a (4b 6) (2a (4b 6) -22 Therefore, the value of (2a 34b 6) is -2.
(2a (4b
Correct answer: -22
31

---

Blink of an eye: a simple theory for feature localization in generative models
LogiQA example (Phi-3-7B -n-struttct
The person in charge the relevant department said that pilots will be carried out some areas this year to prepare for the system and policy preparations for the comprehensive clean-up of "small property houses", and require all localities to confirm the registration and certification of rural collective land. All small-property houses are not subject to registration and are not protected by law. Therefore, this newly built house in Hexi Village is not protected by law. Which of the following is true will most weaken the above argument. (A) This newly built house in Hexi Village has been tacitly approved by relevant departments. (B) This newly built house in Hexi Village is a small property house. (C) This newly built house in Hexi Vil- lage built on rural collective land. (D) Some of the newly built houses Hexi Village are not small property houses.
Q:
A: The statement implies that small property houses are not protected by law and can be demolished if necessary in the preparations for the system and policy preparations for the comprehensive clean-up. We have find the option that weakens this argument the most. Option A suggests that the department has tacitly approved the newly built house. This weakens the argument, but it can still be argued that the house is subject to demolition if the comprehensive clean-up policy is implemented. Option B states that this newly built house in Hexi Village is small property house. This directly contradicts the argument that all small-property houses are not protected by law and not subject to registration and certification. Therefore, this house would not be demolished and would be protected by law. Option C says that the house built on rural collective land. This doesn't 't disprove the argument might still be srapl-property house without legal protection. Option D makes only some of the newly built houses in Hexi Village not small property houses. This doesn't 't address the argument that all srool-property houses not protected by law and not subject registration, but only carries exceptional circumstance. Thus, the option that most weakens the argument is (B) as it directly contradicts the statement by implying that the specific house is protected by law. The answer is B.
Correct answer:
MATH example (Qwen-2.5-7BBInnttree))
Q: Carson flips over the cards a standard 52-card deck one at a time. What is the probability that he flips over the ace of spades before any face card (jack, queen or king)?
A: First, we need to identify the relevant cards for our problem. The ace of spades (A and the face cards (J, Q.
K) the ones we're re're focusing on. There 4 face cards in standard deck of 52 cards (12 face cards total, but only consider 4 for this problem). 2. There are 52 cards in total, and 5 of these (A face card. Therefore, there are 47 other cards that are
Q4 KA JC, QD, KV,
JO, QC, K4, J Q4 ) are either the ace of spades neither the ace of spades nor face cards. 3. The problem can be simplified by considering the probability of the ace of spades coming before any face card a sequence of drawing these cards (the ace of spades and the 8 face cards). We can think the problem as sequence of drawing 9 cards, where we are interested in the position of the ace of spades relative the face cards. The ace of spades can come in any of the 9 positions, and we want find the probability that comes before any of the face cards. 5. For the ace of spades to come before any face card, it must be in of the first 5 positions (since there are 4 face cards, and the ace of spades must come before any of them). There are 9 positions in total, and of these are favorable. 6. Therefore, the probability that the ace of spades comes before any face card is the ratio of the number of favorable outcomes to the total number of outcomes, which is y. So, the probability that Carson flips over the ace of spades before any face
5
card is
Correct answer:
32

---

Blink of an eye: a simple theory for feature localization in generative models
H.3. Varied temperature sampling
In our above experiments with critical windows, we have been generating model responses at their default sampling temperatures:
1ELLL--.--8B-Isttii 0.6, Phi-3-7B-Instrucee 1.0, and Qwen-2.5-7B Insttruct:: 0.7. A natural question to ask is whether the sampling temperature affects the presence strength of critical windows across these models. At one extreme, sampling at infinite temperature exhibits no critical windows since each token is generated uniformly at random. At the other extreme, zero temperature sampling is deterministic, so there is only one mode of behavior and critical windows again cannot exist. At intermediate temperatures, decreasing sampling temperature reduces token variance, so the sharpness of transition from general mixture to target submixture can be heightened. At the same time, reduced sampling variance also reduces the likelihood of transitioning between modes as well, leaving the overall effect on eritical windows unclear.
First, we report the frequeney of critical windows across the different temperatures for the MATH dataset: 0.1, 0.4, 0.7, and 1.0. Then, we report the average size of critical window 400 samples from MATH, running the bereerd-reverse experiment on each sample 25 times and taking the max consecutive jump as the critical window size for that problem. From Figures 13 and 14, each model has its own relationship between temperature and critical window frequency and size, with no clear definite relationship across all models.
Cal Worthe Resquenry
0.38
0.00
0.38
0.28
322
06.00
Tenperouure (1 ama MOTH-
Tompenbuut (ATH
011.00
03335
4109
Terprratre (DONe-AAATi
Figure 13. Temperature versus critical window frequency for LLAMA-2.1-8B-In-tt Phi-3-7B-Instren and Dwen-2.5-7BB IB-strrtct..
33

---

Blink of an eye: simple theory for feature localization in generative models
Cficc roen 50m
1.0
2.44
0279
0.79
14
0315
Senperatore (1 eMS-MATT
Temperature (PN--ATIIO
222
3.59
82
Q4
36
10
Figure 14. Temperature versus mean critical window size for LLAMA-3. 3.1--BB-1-BBB- 1-8B Instruct, Phi-3-BB Instruct, and Qwen-2.5- IB-tnruuc...
34


Reviews:
[{"review_id": "RmXDWPTJyJ", "reviewer": "Reviewer_D7NV", "summary": "The paper investigates \"critical windows\" in generative models—brief intervals during the generation process in which features of the final output are determined. The authors introduce a general theoretical framework based on stochastic localization samplers, a class that includes both diffusion models and autoregressive models as special cases. Their data model assumes samples are drawn from a mixture distribution, with sub-mixtures corresponding to specific features. The core theoretical contribution involves analyzing forward-reverse experiments to identify critical windows: time intervals during which the inversion of a noised observation yields a distribution localized on a sub-mixture, corresponding to the emergence of a feature.\n\nThe authors instantiate their theory with examples such as Gaussian mixture models under diffusion and stylized settings modeling jailbreaks, math reasoning, and in-context learning in large language models. They then extend the theory to handle hierarchical mixture models, where modes are recursively nested. Finally, they perform experiments with large language models showing the presence of critical windows during their generation and that these windows are more likely to occur when the model outputs incorrect answers.\n\n## Update after rebuttal\nThe paper is technically sound, offers some unifying perspectives, and presents interesting experiments on LLMs.\nWhile I still find the predictive power of the framework in empirical settings somewhat limited, the authors have addressed most of my concerns. I have therefore raised my score to recommend acceptance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and thoughtful comments, especially with respect to our theory’s assumptions and the relationship between our experiments and theory. We are glad that you found that the generality of our theory and LLM reasoning experiments interesting.\n\n**Modeling assumptions**\n* *“The claim that their theory requires \"no distributional assumptions\" is overstated.. it relies on having data from a mixture model… Can the authors clarify how general the data distribution assumption and the theoretical results are?”*\n\nIn the final revision, we will rephrase “no distributional assumptions” to “very few distributional assumptions.” The mixture model assumption is **extremely mild**; any partition of the outputs from a generative model defines a mixture model, where the classes are given by the different partitions. For example, for a list of outputs of a language model, we can split them into labeled groups such as {correct answer, incorrect answer}, {safe answer, unsafe answer}, etc. The ability to attach these labels and partition into subpopulations is broadly applicable to the datasets, which we also make use of in our experiments.\n\nSee Remark 2.4 as well for a re-emphasis of the generality of the mixture model assumption.\n\n* *“It seems that the width of the critical windows varies significantly according to the considered model... the applicability of the theory crucially depends on the considered data structure... limit of validity of this modeling assumption should be better clarified.”*\n\nWe view one of our main contributions as offering a **unifying framework** to distill the phenomenon of critical windows to very general facts about the data distribution, so that their **presence/absence or width** depends only on computations with the data model and forward process. This is a major strength compared to extant literature, which only discusses critical windows for **particular data distributions**. In the final revision of this paper, we will clarify that our bounds and the narrowness of our critical windows are affected by the specifics of the data distribution.\n\n* *“there are mixture distributions and samplers for which their bounds may be vacuous.”*\n\nIn Sec. 4 and 5, we verify that our bounds are non-vacuous in many contexts. \n\n* *\"Is there an intuition about when to expect sharp critical windows?\"*\n\nIn Ex. 4.3, we mention an example providing general intuition when critical windows can be sharp, i.e. when a few tokens are very unlikely under one mode compared to the other. In general, we expect sharp critical windows when it only takes a few steps from the forward process to erase the differences between $S_{\\textrm{before}}$ and $S_{\\textrm{after}}$. This could happen if the data has a multi-scale hierarchical structure (Definition 5.1), where a feature is decided in a narrow intermediate band of the tree. \n\n**Experiments**\n* *“Is there some qualitative phenomenon we can predict from the theory (e.g., existence or not of critical windows, their width, etc.) that can be then verified in the experimental data?”*\n\nWe highlight **several predictions of our theory** verified by experiments: \n* We provide structured output experiments where our predictions for the location of critical windows match experiments (Fig. 5, App. G). \n\n* We predict that prefill jailbreaks yield narrow critical windows, because the probability that a model agrees to a harmful request in the first few tokens but refuses in the end is low (Ex. 4.3). This is consistent with (Haize Labs 2024b) and Fig. 4a. \n\nFor LLM reasoning experiments, our only claim is that critical windows coincide with reasoning mistakes (Table 1). In the final revision, we will make which aspects of our theory that the experiments verify more clear. \n\nWe note other works verifying theory with experiments: Li and Chen 2024 showed that positions of theoretical computations of critical windows matched up with experiments for Gaussian mixtures, and Biroli et al. 2024 showed a measure of the separation (size of principal component) between classes predicts real-life critical windows for diffusion. \n\n* *“In the experiments, the presence or absence of critical windows seems to depend strongly on the starting data. Can you elaborate more on that?”*\n\nWe agree that the specifics of the starting point could affect the location or presence of the critical window (Fig. 7). One explanation is that sometimes certain parts of the solution are more important to the final answer. For example, the bolded critical window in Figure 3 occurs at the point where the model finds the correct formula is crucial to solving the problem. In other instances, no particular part of the text could be crucial to the answer. We will clarify this in the final revision.\n\nThank you again for your time in reviewing the paper and providing much helpful feedback. If we have addressed your concerns about the paper, we hope you consider raising our score.", "claims_and_evidence": "The authors claim that their theory applies to a wide class of generative models (both diffusion and autoregressive), that it improves upon prior theoretical results, and that it avoids strong distributional assumptions. \nThe first part of their claims is well supported. However, I think that the claim that their theory requires \"no distributional assumptions\" is overstated. In fact, it relies on having data from a mixture model, and it is unclear how to apply it to more complex data structures. Moreover, as the authors acknowledge in Remark 3.2, there are mixture distributions and samplers for which their bounds may be vacuous. Therefore, the applicability of the theory crucially depends on the considered data structure.\n\nOn the empirical side, the reported experiments with LLMs provide evidence that critical windows can be identified in their generative process.", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria make sense.", "theoretical_claims": "I checked the correctness of the proof of the main theorem, which is sound and logically well-organized.", "experimental_designs_or_analyses": "The experimental designs are sound for the considered tasks. The observation of the existence of critical windows in many LLMs tasks is interesting on its own. However, I am not sure about the connection between the experiments and the proposed theory: is there some qualitative phenomenon we can predict from the theory (e.g., existence or not of critical windows, their width, etc.) that can be then verified in the experimental data?", "relation_to_broader_scientific_literature": "The paper improves previous theoretical results on critical windows for diffusion of mixtures of log-concave distributions [Li&Chen 2024].\nThe studied phenomena are connected to similar studies of diffusion models from a statistical physics perspective, which focus more on specific data models [Raya&Ambrogioni2023, Sclocchi et al. 2024, 2025; Biroli et al. 2024].\nThe paper connects these ideas of critical windows in diffusion models with recent observations in Jailbreaks and chain-of-thought in LLMs.", "essential_references_not_discussed": "The essential scientific literature is cited.", "other_strengths_and_weaknesses": "Strengths\n- The theoretical framework is general and draws connections between stochastic localization and different generative models, such as diffusion and autoregressive models.\n- The identification of critical windows in LLM tasks and how they correlate with accuracy is interesting.\n\nWeaknesses\n- The theory relies on a specific structure of the data distribution and its features, and the limit of validity of this modeling assumption should be better clarified.\n- The connection between experiments and theory is not very compelling.", "other_comments_or_suggestions": "- Example 4.2: $S_{after}$ should be $\\{\\mu_i\\}$", "questions_for_authors": "Can the authors clarify how general the data distribution assumption and the theoretical results are?\nIt seems that the width of the critical windows varies significantly according to the considered model. Is there an intuition about when to expect sharp critical windows?\nIn the experiments, the presence or absence of critical windows seems to depend strongly on the starting data. Can you elaborate more on that?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eAIGBXr33N", "reviewer": "Reviewer_n2Rr", "summary": "The paper presents a theory of “critical windows” – intervals in the generation process in which specific features of the generated data emerge – in both diffusion and autoregressive systems. Leveraging the framework of stochastic localization, the authors rigorously characterize when such windows appear. The theory is applied to several scenarios, including diffusion of Gaussian Mixture Models, jailbreaks in LLMs, a minimal model of problem-solving, and in-context learning. It also considers hierarchical distributions, where a hierarchy of critical windows separate different subpopulations. Finally, experiments demonstrate critical windows in LLM solving reasoning tasks.\n\n## Update after rebuttal\n\nI maintain my positive assessment of this work.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and thoughtful comments, especially with respect to the contribution section and the exposition of our theory. We were glad to hear that you found that our rigorous unifying framework was interesting and that our experiments were well-executed.\n\n# Contributions\n\n* *“Can the authors give a more standard list of contributions, briefly explaining the setting, the obtained insights, and only then the experimental results?\"*\n\nWe will revise the contributions in the final version, separating into clear “theoretical” and “empirical” sections for clarity. On the theoretical side, we will specify that, unlike existing frameworks, our theory **applies across all generative models and data distributions** represented by the stochastic localization framework, as accomplished by Theorem 3.1 and Definition 3.3. Moreover, the theoretical “bounds”, which we clarify as predictions for the location of critical windows, are more precise than that in Li and Chen 2024. \n\nWe also highlight applications of our general theory into specific but important contexts as contributions: for example, we can now **compute critical windows in many different contexts** (discrete diffusion, in-context learning, statistical inference), **unlike existing work (Section 4)**.  We will also specify a novel result for hierarchically structured data, where, if the learned sampler and true model are based on the same localization sampler and the learned sampler is good, then they have the same hierarchical structure (Corollary 5.3).\n\nWe will add a section that better explains our framework intuitively as well.\n\n* *“It briefly mentions “bounds””*\n\nBy this we mean a comparison between the computations of the locations for critical windows for Gaussian diffusion in Li and Chen 2024 versus this paper (Theorem 3.1). They were only able to control the total variation by epsilon times a factor polynomial with the dimension $d$. We were able to upper bound the total variation by epsilon times a constant, and our theorem **improves on their results by a factor that grows polynomially with $d$**. We will clarify this in our contributions. \n\n* *“is “Generality” really a contribution of the paper?”*\n\nBy generality, we mean that our framework applies to all stochastic localization samplers and models of data, not just the Gaussian diffusions and the toy models of data considered before in the literature. We view this ability of our **unifying framework** to explain critical windows across so many different contexts as a major contribution of our work. \n\n* *“On a side note, also the abstract is not fully informative of the paper's content.”*\n\nWe will synthesize the background in the abstract and better describe our contributions.\n\n# Other Comments or Suggestions\n\n* *“The plots in Figure 1 …  are hard to read”*\n\nIn the final revision, Figure 1 will only include three examples that will be explained: the Georgiev et al. 2023 critical window, the prefill attack from Haize Labs 2024, and Phi-4 critical tokens in Abdin et al. 2024. \n\n* *“That’s true only in the case of Gaussian diffusion.”*\n\nWhile the initial applications of stochastic localization (Eldan 2013; 2020) were Gaussian diffusion, extensions of stochastic localization by Montanari 2023 apply to a broader family of generative models, including discrete diffusion models (Example B.2 and Section 4.3. of Montanari 2023).\n\n* *“Is the formulation of autoregressive systems as a stochastic localization sampler a novel contribution of the work?\"*\n\nThis was first presented in Montanari 2023. In Section 2.1 and 2.2, we explicitly cite this work. and we will modify the text to cite it in Appendix B when we instantiate language models within this framework. \n\n* *“Can you please elaborate more on the complexity of hierarchies learned by diffusion vs autoregressive models, as speculated at the end of Sec. 5?”*\n\nWe view the dimension of a diffusion model as the dimension of the underlying space, $d$ in $\\mathbb{R}^d$, and the dimension of an autoregressive model as the length of its context, $T$ in $\\mathcal{A}^T$. We simply pointed out that the hierarchy depth was $O(\\ln d)$ for a continuous diffusion example and $\\Omega(T)$ for an autoregressive example. While they refer to different modalities, we wanted to highlight this vast difference between how hierarchy depth can vary with the dimension. In the final version, we will explain this more clearly. \n\n\nThank you again for your time in reviewing the paper and providing much helpful feedback. If we have addressed your concerns about the paper, we hope you consider raising our score. \n\nEldan, R. Thin shell implies spectral gap up to polylog via a stochastic localization scheme. *Geometric and Functional Analysis*, 23(2):532–569, 2013.\n\nEldan, R. Taming correlations through entropy-efficient measure decompositions with applications to mean-field approximation. *Probability Theory and Related Fields*, 176(3-4):737–755, 2020.", "claims_and_evidence": "The paper is primarily theoretical and rigorously supports its claims.", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I checked the validity of the main result (Theorem 3.1).", "experimental_designs_or_analyses": "Experiments are well-executed.", "relation_to_broader_scientific_literature": "The paper follows a recent rich literature on critical windows and phase transition in generative diffusion models. In particular, it extends the rigorous results of Li & Chen (2024) to general stochastic localization samplers, relaxing several technical assumptions and including autoregressive systems, and tightening the bounds in the case of Gaussian diffusion.", "essential_references_not_discussed": "None that I identified.", "other_strengths_and_weaknesses": "**Other strengths**\n- The paper provides an interesting and rigorous unifying theoretical framework for critical windows in both diffusion and autoregressive models.\n\n**Other weaknesses**\n- Section 1.1. “Our contributions” is quite unclear. In particular, I think it lacks a clear and comprehensive list of the contributions of the paper, especially the theoretical ones. It briefly mentions “bounds” (On what? Obtained how? In which framework?). Moreover, is “Generality” really a contribution of the paper? It also mixes theoretical and empirical contributions. Can the authors give a more standard list of contributions, briefly explaining the setting, the obtained insights, and only then the experimental results? The paper is rather dense in content, so I think it would greatly benefit from a clearer outline of contributions. On a side note, also the abstract is not fully informative of the paper's content.", "other_comments_or_suggestions": "- The plots in Figure 1, taken from related work, are not explained and are hard to read/understand, especially at that point of the introduction. Personally, I don’t see the necessity of including such a figure. I’d encourage the authors either to remove it or – in case they wish to keep it – to make it larger (use vectorized graphics) and explain the content.\n- Is the formulation of autoregressive systems as a stochastic localization sampler a novel contribution of the work? If so, I would suggest highlighting it more. Otherwise, the paper should cite previous work showing it. \n- L132 (right column): “The can be understood”?\n- L144 (right column): That’s true only in the case of Gaussian diffusion.\n- The running title is still the ICML template default and should be updated.", "questions_for_authors": "- Can you please elaborate more on the complexity of hierarchies learned by diffusion vs autoregressive models, as speculated at the end of Sec. 5? Don’t the two results refer to different data models/distributions?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2KDNhDhwA0", "reviewer": "Reviewer_HCM5", "summary": "The paper theoretically explains sudden behavioral shifts in generative models through critical windows, employing a forward-reverse experiment to study this phenomenon. It introduces Theorem 3.1, which bounds total variation distance to demonstrate that these windows signify transitions between sub-mixtures. The findings are substantiated with examples from diffusion and autoregressive processes.\n\n## update after rebuttal\nThe authors' response has addressed my concern, so I have raised the score from 3 to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their comments. We were glad to hear that you found that our experimental results for LLMs are convincing.\n\n* *“Since LLM performance is sensitive to evaluation metrics, a deeper discussion on the robustness of critical windows across different metrics is needed.”*\n\nIn appendix H.1, we include a discussion about the evaluation metrics we used to test our model, and in appendix H.4, we explore the effect of different temperatures. Note that we use standard methods like direct text comparison for multiple choice questions (Lanham et al. 2023) and existing math graders from the literature (Lightman et al. 2023). Given the primary focus on this paper is theoretical, and our experiments are commensurate with these well-cited manuscripts on LLM evaluation and performance, our evaluation metrics and discussion cover a broad range of datasets, models, and other empirical settings that demonstrate the robustness of the critical windows phenomenon.  \n\n* *“The experimental setup is not rigorously defined or directly validated against the theory, limiting its connection to Theorem 3.1. While TV distance may not be feasible for real distributions, simulations could provide a more direct validation of the theoretical results.”*\n\nAppendix G describes a **direct validation of our theoretical results for LLMs**, where we actually compute the TV distance to verify our bounds for a real-world model. \n\n\nThank you again for your comments. If we have addressed your concerns about the paper, we hope you consider raising our score. \n\n\nLanham et al. (2023). Measuring Faithfulness in Chain-of-Thought Reasoning. *arXiv preprint arXiv:2307.13702*. Retrieved from https://arxiv.org/abs/2307.13702\n\nLightman et al. (2023). Let's Verify Step by Step. *arXiv preprint arXiv:2305.20050*. Retrieved from https://arxiv.org/abs/2305.20050", "claims_and_evidence": "The central claim is supported by Theorem 3.1. Experimental results further confirm the presence of critical windows in generations.", "methods_and_evaluation_criteria": "The theoretical method uses stochastic localization samplers and mixture models and arrives at a conclusion in TV distance bounds. This is a sensible approach for studying feature localization in generative models.", "theoretical_claims": "I did not check the proofs in the Appendix.", "experimental_designs_or_analyses": "The results on LLMs clearly demonstrate abrupt changes in output probabilities.  \n\nHowever, several concerns remain:  \n- Since LLM performance is sensitive to evaluation metrics, a deeper discussion on the robustness of critical windows across different metrics is needed.  \n- The experimental setup is not rigorously defined or directly validated against the theory, limiting its connection to Theorem 3.1. While TV distance may not be feasible for real distributions, simulations could provide a more direct validation of the theoretical results.", "relation_to_broader_scientific_literature": "The paper builds on prior work on critical windows in diffusion models (e.g., Sclocchi et al., 2024; Li & Chen, 2024).", "essential_references_not_discussed": "I did not notice any missing key references.", "other_strengths_and_weaknesses": "See *Experimental Designs Or Analyses*.", "other_comments_or_suggestions": "See *Experimental Designs Or Analyses*.", "questions_for_authors": "See *Experimental Designs Or Analyses*.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZFTceIG2Yu", "reviewer": "Reviewer_TMA6", "summary": "The authors present a paper that explores critical windows in generative models.  Their paper is heavily theoretical and they propose an understanding that can be applied to a wide range of models.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their kind comments and strong recommendation. We will modify the abstract to say ''hacks’’ instead of ''jailbreaks.’’", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "I did not check the accuracy of the proofs", "experimental_designs_or_analyses": "NA", "relation_to_broader_scientific_literature": "I believe this paper will have broad appeal to the machine learning community", "essential_references_not_discussed": "no", "other_strengths_and_weaknesses": "First, I will admit to being somewhat biased in favour of strong theoretical contributions, but this paper stands out an exceptionally well written, instructional and informative example.  While the underlying theme of the paper is theoretical, the authors bring their contribution to a real issue in diffusion and LLMs.  I also appreciate the included clear examples adding context to the theoretical formulation.", "other_comments_or_suggestions": "undefined terms in the abstract (jailbreak), though this is defined quite early in the introduction", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "l7ueukwI3m", "reviewer": "Reviewer_aK3T", "summary": "This paper discusses the phenomenon of critical windows in generative models. It is an interesting topic, and the paper presents a general theory with minimal assumptions, enabling the explanation of abrupt shifts during the sampling phase across different modeling paradigms and data modalities. The writing is clear, and the definition of critical windows based on sub-mixtures, along with the discussion on hierarchical sampling, is engaging.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and thoughtful comments. We were glad to hear that you found our writing clear and engaging and our theory interesting.\n\n\n* *“If the reverse process is deterministic, such as an ODE, or includes additional conditions, such as text-to-image, does this theoretical framework still apply?”*\n\n\nIf the reverse process is deterministic, then there is no notion of a critical window under our framework. The initial position at the start of sampling completely characterizes the final image. For example, given a fixed piece of text and language model, truncating it anywhere in the model’s response and resampling at temperature $0$ would yield the same final text. This means that the probability it would yield the same answer as the original generation would be 0. We view this as a fruitful future direction of work to extend our framework to deterministic samplers. \n\n\n* *“Section 4 presents some case studies. Could you provide further experimental results to verify the accuracy of the computed critical windows from the theoretical analysis?”*\n\n\nWe would like to highlight that **many of the case studies in Section 4** are accompanied by  experiments either in the appendix or in the existing literature:\n* Li and Chen 2024 confirmed that critical windows for mixtures of Gaussians matched with experiments.\n\n* A critical window for the all-or-nothing phenomenon in sparse linear regression can be seen in Figure 2 of Reeves et al. 2019\n\n* The jailbreak critical windows are demonstrated in previous literature, e.g. Haize Labs 2024, and are reproduced in our Appendix F.1. \n\n* In the final revision, we will mention these experiments along with these examples or present some as well as a diagram for a discrete diffusion model that is a mixture of delta measures. \n\n\nThank you again for your time in reviewing the paper and providing much helpful feedback. If we have addressed your concerns about the paper, we hope you consider raising our score. \n\nHaize Labs. (2024). *A trivial jailbreak against LLaMA 3.* https://github.com/haizelabs/llama3-jailbreak\n\nLi, M., & Chen, S. (2024). Critical windows: Non-asymptotic theory for feature emergence in diffusion models. arXiv preprint arXiv:2403.01633.\n\nReeves, G., Xu, J. & Zadik, I. (2019). The All-or-Nothing Phenomenon in Sparse Linear Regression. *Proceedings of the Thirty-Second Conference on Learning Theory in Proceedings of Machine Learning Research* 99:2652-2663.", "claims_and_evidence": "1. If the reverse process is deterministic, such as an ODE, or includes additional conditions, such as text-to-image, does this theoretical framework still apply?", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I have checked the proof of Theorem 3.1 and found no additional issues.", "experimental_designs_or_analyses": "2. Section 4 presents some case studies. Could you provide further experimental results to verify the accuracy of the computed critical windows from the theoretical analysis?", "relation_to_broader_scientific_literature": "This paper proposes a unified and concise theoretical framework that explains the critical windows phenomenon observed in autoregressive and diffusion models in previous studies.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AJ3ESzzpbs", "reviewer": "Reviewer_MfZY", "summary": "This paper introduces a general framework for critical windows in stochastic localization. After a lengthy but valuable description of some key notions such as stochastic localization sampling and the \"forward-reverse experiment\", the authors prove their key result, which shows that there exist (possibly empty) \"critical windows\" during stochastic localization sampling in which the forward process has destroyed the information needed to distinguish a submixture $S_{init}$ from a larger submixture $S_{targ}$, but not the information needed to distinguish $S_{targ}$ from the remainder of the data distribution. They provide a simple toy example in the case of diffusion models and a number of examples drawn from autoregressive learning. They then sketch a general theory of how stochastic localization interacts with hierarchical semantic structure and briefly describe the results of some experiments on LLMs.\n\n## Update after rebuttal.\n\nI appreciate the authors' engagement with my review and their promise to add figures and examples to the camera-ready as intuition-building tools. I maintain my positive assessment of this paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and thoughtful comments. We were glad to hear that you thought that the theoretical results were interesting and that experiments were sound and illustrative of our main points.\n\n## Writing changes\n\n* *“While this paper is fairly well-written, it is dense with notation and short on figures”*\n\nIn the final revision, we will include a section titled “Intuition for critical windows” before the technical prelims that **describes our theory informally** and introduces the forward-reverse experiment, the definition of critical windows, the definition of $\\epsilon$-mixture trees, and our main Theorem 3.1 through a **simple vignette**. All of these definitions and theorems will be accompanied by **figures that visually explain them and text which concretely places them within our vignette.** For example, the definition of the forward-reverse experiment will be accompanied by a figure which shows how a “sweet spot” of noise leads to the specialization to a target sub-mixture; the definition of $\\epsilon$-mixture trees will be shown with a graph that shows the hierarchy of features in our vignette; we will expand and add more detail to Figure 2 of a critical window in this section.\n\nThe section will very loosely follow this structure: we will describe critical windows as the transition from sampling from a larger subset of features to a smaller subset of features. This motivates trying to understand when the generative model is sampling from a subset of features, and thus the forward-reverse experiment and our main Theorem 3.1. We will provide intuition into the location of the bounds for Theorem 3.1 and then explain how sequences of critical windows motivate understanding a hierarchy of feature specialization and thus the definition of $\\epsilon$-mixture trees.\n\n* *“For example, the authors could include a figure illustrating the forward-reverse experiment...”*\n\nYes, in the aforementioned new section we will illustrate the forward-reverse experiment with a very concrete example and figure.\n\n* *\"The definition of ϵ-mixture trees in Section 5 is also very abstract, and while I appreciate the benefits of this approach from the standpoint of generality, a few figures or additional examples would help readers parse the definitions better.\"*\n\n\nWe also plan to include an example and figure of an “$\\epsilon$-mixture tree” in the section “Intuition for critical windows” accompanying our text and the vignette.\n\n* *“include a more thorough discussion of their experiments in the main body… consider including figures, an expanded experiments section, and at least an abbreviated related work section in the main body”*\n\nIn addition to the new section, we will move many details from the examples and hierarchy sections to the appendix, add a short related works section in the main body, and thoroughly expand the experiments section. The expanded experiments section will include our structured output experiments, which show that our theory is predictive of critical windows for LLMs when outputs are hierarchically structured, and we will add to our LLM reasoning experiments details that were originally relegated to the appendix, e.g., statistics and visualizations of critical windows across datasets and models. The abbreviated related works section will cover the theory of critical windows for diffusion, the forward-reverse experiment, and stochastic localization.\n\n## Future directions\n* *“not clear a priori that non-empty critical windows should exist. Exploring why this is the case would be an interesting future direction.”*\n\nIn the Yellowstone and jailbreaking example, some actions from the LLM, i.e. browsing Yellowstone or acceding to a harmful user request, are much likelier under one mode of behavior than another and completely determine to which mode the generation belongs, resulting in a critical window as explained by Example 4.3. We agree with the reviewer that further exploring why critical windows exist in different settings is an interesting direction of future research. \n\nThank you again for your time in reviewing the paper and providing much helpful feedback. If we have addressed your concerns about the paper, we hope you consider raising our score.", "claims_and_evidence": "The main claim in this paper is to have constructed a general theory to explain \"critical windows\" in generative models, which the authors colloquially define as a small subset of steps in which important features of a model sample emerge. I believe the paper largely achieves this objective. The stochastic localization framework is general enough to include diffusion models and autoregressive models, which are the two main paradigms in generative modeling nowadays. Their main result is general and does not rely on strong assumptions or particularly heavy sledgehammer results. However, as the authors acknowledge in Remark 3.2, their theory does not rule out the possibility of critical windows being empty sets, which I believe to be the main gap in their story. Nonetheless, I think the results are sufficiently interesting to stand on their own, and look forward to future work exploring why critical windows are often non-empty for common classes of generative models.", "methods_and_evaluation_criteria": "The theoretical methods are appropriate for demonstrating the authors' key claims. In particular, the \"forward-reverse experiment\" is an appropriate tool for formalizing the process of destroying and then recovering information in stochastic localization sampling.", "theoretical_claims": "I have reviewed the proof outline for Theorem 3.1. I believe the strategy is correct, and I was unable to find any specific errors in the outline.", "experimental_designs_or_analyses": "The experiments included in the main body seem to be sound and adequately illustrate that critical windows can occur in LLMs. However, I would have liked the authors to include a more thorough discussion of their experiments in the main body -- as I will note below, it generally seems like the authors have packed too much content into the 8-page limit at the price of e.g. a very abbreviated experiments section.", "relation_to_broader_scientific_literature": "This paper generalizes results from Li and Chen (2024), which studies the existence of critical windows in diffusion models. It draws heavily on tools from the stochastic localization literature, which is anchored by a series of papers by Eldan and connected to diffusion models in a set of notes by Montanari (2023).", "essential_references_not_discussed": "While I am not well-versed in the literature on critical windows and stochastic localization, it seems to me that this paper adequately situates itself in its literature.", "other_strengths_and_weaknesses": "While this paper is fairly well-written, it is dense with notation and short on figures. A few figures to illustrate the key notions would greatly improve the readers' intuition for the results. For example, the authors could include a figure illustrating the forward-reverse experiment for a simple case like a mixture of Gaussians and depicting the model distribution $p_t^S$ for various subsets $S$ during the critical window predicted by Theorem 3.1. The definition of $\\epsilon$-mixture trees in Section 5 is also *very* abstract, and while I appreciate the benefits of this approach from the standpoint of generality, a few figures or additional examples would help readers parse the definitions better.\n\nTo me, the most interesting takeaway from this paper is that despite knowing nothing about the semantics of the data a priori, a stochastic localization sampler generates information in a way that respects the data's semantic hierarchy. e.g. a dog is an animal, so the support of the distribution over images of dogs is contained in the support of the distribution over images of animals -- and stochastic localization features critical windows in which the sampler has \"decided\" to generate an animal image but not yet decided that it will generate a dog image. It seems surprising to me that one can prove a general result of this form. However, I believe that Remark 3.2 reveals the primary gap in this theory as it stands -- it is not clear a priori that non-empty critical windows should exist. Exploring *why* this is the case would be an interesting future direction.", "other_comments_or_suggestions": "In general, it seems like the authors attempted to pack too much content into the 8-page limit, and have consequently neglected to include illustrative figures and a related work section in the main body of the text. They have also heavily compressed their experiments section. If this paper is accepted, I'd ask the authors to consider including figures, an expanded experiments section, and at least an abbreviated related work section in the main body, perhaps at the price of moving some of the examples in Section 4 to the appendix.", "questions_for_authors": "No further questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper studies the emergence of “critical windows” in generative models, namely locations in a generative process where some underlying features are determined. While this has been studied previously in diffusion models under strong distributional assumptions, the theory presented here refines existing bounds and applies more broadly to different classes of generative models and with very mild assumptions on the data distribution. The reviewers agree that this is a solid contribution and are all in favor of acceptance. Please include the proposed clarifications in the final version.
Decision: Accept (oral)