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.

Writing changes

• “While this paper is fairly well-written, it is dense with notation and short on figures”

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

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

• “For example, the authors could include a figure illustrating the forward-reverse experiment...”

Yes, in the aforementioned new section we will illustrate the forward-reverse experiment with a very concrete example and figure.

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

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

• “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”

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

Future directions

• “not clear a priori that non-empty critical windows should exist. Exploring why this is the case would be an interesting future direction.”

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

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

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.

• “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?”

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

• “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?”

We 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:

• Li and Chen 2024 confirmed that critical windows for mixtures of Gaussians matched with experiments.

• A critical window for the all-or-nothing phenomenon in sparse linear regression can be seen in Figure 2 of Reeves et al. 2019

• The jailbreak critical windows are demonstrated in previous literature, e.g. Haize Labs 2024, and are reproduced in our Appendix F.1.

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

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

Haize Labs. (2024). A trivial jailbreak against LLaMA 3. https://github.com/haizelabs/llama3-jailbreak

Li, M., & Chen, S. (2024). Critical windows: Non-asymptotic theory for feature emergence in diffusion models. arXiv preprint arXiv:2403.01633.

Reeves, 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.

We thank the reviewer for their kind comments and strong recommendation. We will modify the abstract to say ''hacks’’ instead of ''jailbreaks.’’

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.

• “Since LLM performance is sensitive to evaluation metrics, a deeper discussion on the robustness of critical windows across different metrics is needed.”

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

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

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

Thank you again for your comments. If we have addressed your concerns about the paper, we hope you consider raising our score.

Lanham et al. (2023). Measuring Faithfulness in Chain-of-Thought Reasoning. arXiv preprint arXiv:2307.13702. Retrieved from https://arxiv.org/abs/2307.13702

Lightman et al. (2023). Let's Verify Step by Step. arXiv preprint arXiv:2305.20050. Retrieved from https://arxiv.org/abs/2305.20050

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.

Contributions

• “Can the authors give a more standard list of contributions, briefly explaining the setting, the obtained insights, and only then the experimental results?"

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

We 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).

We will add a section that better explains our framework intuitively as well.

• “It briefly mentions “bounds””

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

• “is “Generality” really a contribution of the paper?”

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

• “On a side note, also the abstract is not fully informative of the paper's content.”

We will synthesize the background in the abstract and better describe our contributions.

Other Comments or Suggestions

• “The plots in Figure 1 … are hard to read”

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

• “That’s true only in the case of Gaussian diffusion.”

While 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).

• “Is the formulation of autoregressive systems as a stochastic localization sampler a novel contribution of the work?"

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

• “Can you please elaborate more on the complexity of hierarchies learned by diffusion vs autoregressive models, as speculated at the end of Sec. 5?”

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

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

Eldan, R. Thin shell implies spectral gap up to polylog via a stochastic localization scheme. Geometric and Functional Analysis, 23(2):532–569, 2013.

Eldan, 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.

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.

Modeling assumptions

• “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?”

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

See Remark 2.4 as well for a re-emphasis of the generality of the mixture model assumption.

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

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

• “there are mixture distributions and samplers for which their bounds may be vacuous.”

In Sec. 4 and 5, we verify that our bounds are non-vacuous in many contexts.

• "Is there an intuition about when to expect sharp critical windows?"

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

Experiments

• “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?”

We highlight several predictions of our theory verified by experiments:

• We provide structured output experiments where our predictions for the location of critical windows match experiments (Fig. 5, App. G).

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

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

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

• “In the experiments, the presence or absence of critical windows seems to depend strongly on the starting data. Can you elaborate more on that?”

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

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