The Geometric Anatomy of Capability Acquisition in Transformers
Organizations: San Jose, CA, USA
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 tasklevelmodel 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 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.
Figures & tables
| Task | Description | L1 | L2 | L3 |
| COPY | Copy tokens | 3 tokens | 5 tokens | 8 tokens |
| REV | Reverse tokens | 3 tokens | 5 tokens | 8 tokens |
| CMP | Compare two numbers | 1 digit | 2 digits | 3 digits |
| PAR | Parity of bits | 4 bits | 6 bits | 8 bits |
| ADD | Addition | 1+1 digit | 2+2 digits | 3+3 digits |
| MOD | Modular arithmetic |
| Size | Layers | Heads | Parameters | ||
| Nano | 2 | 128 | 4 | 512 | 405K |
| Micro | 4 | 192 | 6 | 768 | 1.8M |
| Small | 6 | 320 | 8 | 1,280 | 7.4M |
| Medium | 8 | 512 | 8 | 2,048 | 25.2M |
| Large | 12 | 768 | 12 | 3,072 | 85M |
| XLarge | 12 | 1,024 | 16 | 4,096 | 151M |
| Measure | Captures | Formulation | Cost |
| RankMe | Repr. dimensionality | , | 1s |
| Grad. eff. rank | Gradient concentration | of gradient eigenvalues | 85s |
| LLC | Loss landscape complexity | SGLD loss elevation | 11s |
| Hessian top- | Curvature | Stochastic Lanczos, (per task) | 1s |
| Grad. cov. rank | Gradient diversity | Effective rank of (per task) | 6s |
| Benchmark | Capability | Prompt format (example) | |
| Syntactic | Subject–verb agreement | “The keys to the cabinet is/are ” (logprob) | 36 |
| Semantic | Word analogies | “man : woman :: king : ___ ” (4-way choice) | 30 |
| Arithmetic | Few-shot addition/mult. | “3+7=10; 15+23=38; 42+8= ? ” (3-shot) | 200 |
| ICL | In-context learning | Novel label mappings (4-shot sent./categ.) | 16 |
| Factual | Factual recall | “The capital of France is ___ ” (next token) | 51 |
| Logical | Deductive reasoning | “All dogs are animals. Rex is a dog. ___ ” | 26 |
| Task | Nano | Micro | Small | Medium | Large | XLarge |
| (405K) | (1.8M) | (7.4M) | (25.2M) | (85M) | (151M) | |
| COPY | 700 | 700 | 400 | 600 | 500 | 400 |
| REV | 700 | 600 | 400 | 600 | 400 | 400 |
| CMP | 400 | 400 | 300 | 100 | 200 | 100 |
| PAR | 900 | 900 | 3,700 | 2,500 | 600 | 900 |
| SORT | 1,100 | 900 | 900 | 700 | 900 | 500 |
| Nano (405K) | Micro (1.8M) | Small (7.4M) | Medium (25.2M) | Large (85M) | ||||||
| Task | Behav | Probe | Behav | Probe | Behav | Probe | Behav | Probe | Behav | Probe |
| COPY | .00 | .81 | .39 | .94 | .00 | .76 | .47 | .95 | .26 | .96 |
| REV | .02 | .83 | .00 | .73 | .00 | .70 | .29 | .92 | .00 | .77 |
| CMP | .49 | .94 | .49 | .96 | .49 | .97 | .00 | .97 | .46 | .96 |
| PAR | .49 | .87 | .42 | .91 | .48 | .94 | .49 | .92 | .49 | .88 |
| SORT | .34 | .98 | .38 | .98 | .31 | .99 | .50 | .99 | .45 | .99 |
| Measure | Level | Hard-Task Precursor Rate | Cost/ckpt | Signal Quality |
| RankMe | Representation | 100% (8/8) | 1s | Clean collapse recovery |
| Hessian | Curvature | 100% (8/8) | 1s | Noisy (oscillates between ckpts) |
| Grad. cov. rank | Gradient | 100% (8/8) | 6s | Noisy |
| Grad. eff. rank | Gradient | 38% (3/8) | 85s | Transition too late for most tasks |
| LLC | Loss landscape | No discrete event | 11s | Peaks early, declines continuously |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Nano | Micro | Small | Medium | Large | XLarge | |
| Peak LR | ||||||
| Max steps | 100K | 100K | 100K | 200K | 200K | 200K |
| Warmup steps | 1,000 | 1,000 | 1,000 | 1,000 | 1,000 | 1,000 |
| Batch size | 64 | 64 | 64 | 64 | 64 | 64 |
| Weight decay | 0.1 | 0.1 | 0.1 | 0.1 | 0.1 | 0.1 |
| Grad clip | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 |
| Task | Level | Nano | Micro | Small | Medium | Large | XLarge |
| COPY | L1 | 800 | 600 | 400 | 500 | 400 | 300 |
| COPY | L2 | 700 | 700 | 400 | 600 | 500 | 400 |
| COPY | L3 | 700 | 700 | 500 | 600 | 500 | 400 |
| REV | L1 | 700 | 600 | 400 | 500 | 400 | 400 |
| REV | L2 | 700 | 600 | 400 | 600 | 400 | 400 |
| REV | L3 | 700 | 700 | 400 | 700 | 500 | 500 |
| Scale | LP 25% | LP 50% | LP 75% | LP 90% | Accuracy |
| (acc 0–10%) | (acc 0–15%) | (acc 30–65%) | (acc 50–90%) | ( 50%) | |
| Nano (405K) | 0% | 100% | 100% | 100% | 100% |
| Micro (1.8M) | 0% | 89% | 100% | 100% | 100% |
| Small (7.4M) | 0% | 33% | 100% | 100% | 100% |
| Medium (25.2M) | 0% | 33% | 89% | 100% | 89% |
| Large (85M) | 0% | 22% | 78% | 100% | 89% |
| Component | Ckpts | GPU-hrs |
| Algorithmic experiment | ||
| Training (5 sizes + 5 freeze) | – | 10 |
| Eval + task-specific RankMe (5 sizes) | 740 | 7 |
| Eval + task-specific RankMe (5 freeze) | 1,030 | 6 |
| Grad. eff. rank + LLC (nano: 206, large: 37) | 243 | 19 |
| Linear probing (nano–large) | 1,130 | 146 |
| Nano (405K) | Micro (1.8M) | Small (7.4M) | Medium (25.2M) | Large (85M) | XLarge (151M) | |
| Hard tasks (3) | 100% (3/3) | 100% (3/3) | 100% (3/3) | 100% (3/3) | 100% (3/3) | 100% (3/3) |
| Task | Baseline | Freeze-L1 | Delay | Freeze-L0 | Delay |
| COPY | 800 | 1,300 | 500 | 1,200 | 400 |
| REV | 700 | 1,500 | 800 | 1,100 | 400 |
| CMP | 300 | 900 | 600 | 600 | 300 |
| PAR | 1,000 | 2,500 | 1,500 | 300 | 700 |
| SORT | 1,800 | 2,000 | 200 | 1,700 | 100 |
| ADD | 19,500 | 33,500 | 14,000 | 11,500 | 8,000 |
Representation Geometry in Language Models
Representation geometry in language models studies how embeddings and hidden states organize concepts, computation, and behavior.
See all 417 papers