Probing Geometry of Next Token Prediction Using Cumulant Expansion of the Softmax Entropy
2510.04285v1
cs.CL, cond-mat.stat-mech, cs.LG, stat.ML
2025-10-08
Авторы:
Karthik Viswanathan, Sang Eon Park
Abstract
We introduce a cumulant-expansion framework for quantifying how large
language models (LLMs) internalize higher-order statistical structure during
next-token prediction. By treating the softmax entropy of each layer's logit
distribution as a perturbation around its "center" distribution, we derive
closed-form cumulant observables that isolate successively higher-order
correlations. Empirically, we track these cumulants in GPT-2 and Pythia models
on Pile-10K prompts. (i) Structured prompts exhibit a characteristic
rise-and-plateau profile across layers, whereas token-shuffled prompts remain
flat, revealing the dependence of the cumulant profile on meaningful context.
(ii) During training, all cumulants increase monotonically before saturating,
directly visualizing the model's progression from capturing variance to
learning skew, kurtosis, and higher-order statistical structures. (iii)
Mathematical prompts show distinct cumulant signatures compared to general
text, quantifying how models employ fundamentally different processing
mechanisms for mathematical versus linguistic content. Together, these results
establish cumulant analysis as a lightweight, mathematically grounded probe of
feature-learning dynamics in high-dimensional neural networks.