Representational Curvature: The Geometric Secret Behind LLM Uncertainty
Representational Curvature Modulates Behavioral Uncertainty in Large Language Models
The paper investigates how internal representational geometry in autoregressive LLMs influences model behavior, introducing "contextual curvature" as a predictor of next-token entropy. Across models like GPT-2 XL and Pythia-2.8B, the authors demonstrate that temporal straightening in middle layers directly modulates next-token predictive uncertainty.
TL;DR
Why are some tokens "easier" for an LLM to predict than others? This paper reveals that the answer lies in the geometry of the internal trajectory. In middle layers, LLMs learn to "straighten" the path of their representations. When this path remains straight, the model is confident; when it bends (high "contextual curvature"), next-token entropy spikes. By regularizing this curvature during training, we can actually "nudge" model confidence without hurting accuracy.
The Intuition: Predicting the Curve
In computational neuroscience, there is a principle called Temporal Straightening. Imagine watching a video of a ball moving: in raw pixel space, the frame-by-frame changes are highly nonlinear and messy. However, our brains transform these into a "straight" representation that makes it easy to predict where the ball will be next via simple linear extrapolation.
The authors hypothesize that LLMs do exactly the same thing. They transform discrete token inputs into a smooth, high-dimensional trajectory.
The Key Insight: If the trajectory is straight, the "next state" is a simple continuation. If the trajectory bends sharply, the future becomes unpredictable, leading to higher entropy (uncertainty).
Measuring the "Bend"
The authors define Contextual Curvature () as the angle between successive displacement vectors in the model's residual stream.
- Low Curvature: Moving in a straight line (Predictable).
- High Curvature: Sudden change in direction (Uncertain).
Evidence 1: The Middle-Layer Sweet Spot
Across different models (GPT-2, Pythia), a striking pattern emerges. In the very first and very last layers, curvature is high and messy. But in the middle layers, the model actively "untangles" the data, and curvature hits a minimum.
Critically, these middle layers are exactly where curvature most accurately predicts the output entropy. It's as if the middle of the network is the "engine room" for straightening the logic of the sequence before generating the final token.
Evidence 2: Causal Deviations (Perturbation Analysis)
Does curvature cause uncertainty, or is it just a byproduct? To find out, the authors performed "surgical" interventions on the residual stream.
- They found that if you perturb the representation in random directions, nothing happens to the entropy.
- However, if you perturb it specifically within the Planar-subspace (the 2D plane where the trajectory is currently "turning"), the entropy changes predictably.
This proves that the model's behavioral uncertainty is selectively tuned to the geometric path of the representation.
Evidence 3: Training the Geometry
Can we force a model to be more confident? The researchers added a "Curvature Penalty" to the loss function:
- Untangled Model: Penalized high curvature (forced straighter paths).
- Tangled Model: Penalized low curvature (forced "bendier" paths).
The result? The Untangled Model consistently showed lower token-level entropy (higher confidence) across datasets like LAMBADA, while maintaining the same validation loss as the baseline. This suggests that representational geometry is a "hidden knob" for controlling model calibration.
Critical Analysis & Takeaways
This work provides a beautiful bridge between geometric deep learning and behavioral uncertainty. It suggests that:
- Curvature is an Inductive Bias: Next-token prediction naturally drives models toward linear representational trajectories.
- Interpretability is Temporal: To understand LLMs, we shouldn't just look at "static" activations of one token, but rather how the "flow" evolves over the context window.
Limitations: The effect sizes for training regularization were modest. While "straightness" helps, it is likely only one of many geometric features (like manifold density or dimensionality) that govern complex language behavior.
Future Outlook: Could we use curvature as a real-time monitor for "hallucinations"? If the representational trajectory starts to bend wildly in the middle layers, it might be a signal that the model is about to generate low-confidence, potentially incorrect text.