The Geometric Anatomy of Capability Acquisition in Transformers

TL;DR

This paper investigates how geometric changes in transformer representations relate to capability acquisition, revealing that geometric measures often precede behavioral improvements, especially in hard tasks.

Contribution

It introduces a detailed analysis of geometric measures across multiple transformer sizes and tasks, highlighting ankme as a reliable precursor to capability acquisition.

Findings

Representations first collapse to low-dimensional states then recover before behavior improves.

Linear probes show task-relevant information exists before the model acts on it.

Only ankme reliably predicts capability acquisition, especially in hard tasks.

Abstract

Neural networks gain capabilities during training, but the internal changes that precede capability acquisition are not well understood. In particular, the relationship between geometric change and behavioral change, and the effect of task difficulty and model scale on that relationship, is unclear. We track geometric measures and linear probes across six transformer sizes (405K--151M parameters), eight algorithmic tasks (144 task$\times$level$\times$model combinations), and three Pythia language models (160M--2.8B). Across all settings, representations first collapse to a low-dimensional state, then recover, and only then does behavioral performance improve. Linear probes show that the model's hidden states already contain task-relevant information before the model can act on it. The collapse floor is task-specific, the collapse propagates top-down through the network, and of the geometric measures tested, only \rankme reliably precedes capability acquisition for hard tasks. Whether this precursor is detectable depends on task difficulty relative to model capacity. For hard tasks, there is a clear gap: geometry changes first, behavior follows. For easy tasks, the model learns so quickly that both happen simultaneously and no precursor is detectable. On Pythia-2.8B, a logical deduction task that is genuinely hard for the model shows a precursor gap of ${\sim}$49K training steps, while easy benchmarks show none. This suggests that geometric patterns observed in small proxy models can persist at larger scale when the task remains difficult relative to model capacity.