Blink of an eye: a simple theory for feature localization in generative models

TL;DR

This paper introduces a simple, unifying theory explaining sudden shifts in behavior of generative models during output, applicable to both language and diffusion models, and validated through empirical analysis.

Contribution

The work develops a theory that explains feature localization in generative models without relying on strong distributional assumptions, applicable to multiple model types.

Findings

Critical windows often coincide with failures in problem solving.

The theory applies to autoregressive and diffusion models.

Empirical validation confirms the theory's predictions.

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 samplers. 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 no distributional assumptions; (3) quantitatively improves previous bounds even when specialized to diffusions; and (4) requires basic tools and no stochastic calculus or statistical-physics-based machinery. We also identify an intriguing connection to the all-or-nothing phenomenon from statistical inference. 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.