Accepted at ECML PKDD 2026 · Paper: “Grokking Is a Geometric Phase Transition: Archetypal Analysis Reveals Causal Structure in Transformer Representations”
When Transformers grokk modular arithmetic, their internal activations reorganize into a simplex whose vertices are causally meaningful algorithmic states — one per output class. We show this using Archetypal Analysis (AA) and validate it through targeted causal interventions.
Grokking — the sudden transition from memorization to generalization in overparameterized models — remains poorly understood at the representation level. We apply Archetypal Analysis (AA) to decompose Transformer residual-stream activations into convex mixtures of extreme points ("archetypes") and track how this decomposition evolves across training. Our key findings:
- Co-emergence: Archetype purity jumps sharply at the grokking transition, co-emerging with test accuracy and causal efficacy.
- Causal control: Swapping AA mixture weights between inputs redirects model output with 98–100% accuracy (vs. ~6% for k-means and ~1% for random baselines).
- Label-free discovery: Behavior-based merging of overcomplete archetypes recovers exactly P output classes with 100% label alignment — using only the model's output head.
- Group-theoretic structure: Fixed archetype permutations implement the cyclic group action y → y+d (mod P) at 80–85% success for k=P, and at 100% using the behavior-merged k=2P→P representation (both addition and multiplication, all seeds). For addition the relevant group is ℤ/Pℤ (linear label order); for multiplication it is (ℤ/Pℤ)* (discrete-logarithm order, label 0 excluded).
These results hold across modular addition, multiplication, and parity, with a random-lookup control showing no structure.
Prior work on grokking has characterized the what (Fourier modes emerge, circuits simplify) but not the how at the representation geometry level. Standard tools face limitations:
- PCA reconstructs well but produces unconstrained coordinates without causal semantics
- k-means assigns hard clusters that fail as causal factors (~6% swap success)
- Neuron-level analysis is confounded by superposition
AA is uniquely suited because it produces simplex coordinates (convex weights summing to 1) over extreme points (archetypes on the convex hull). This gives a natural "coordinate system" for the representation space that turns out to be causally manipulable.
| Parameter | Addition | Multiplication | Parity | Random (control) |
|---|---|---|---|---|
| Operation | (a+b) mod 113 | (a×b) mod 113 | XOR(16-bit) | random lookup |
| P (classes) | 113 | 113 | 2 | 113 |
| seq_len | 2 | 2 | 16 | 2 |
| weight_decay | 1.0 | 1.0 | 0.1 | 1.0 |
| Grokking step | ~3,500 | ~3,500 | ~5,000 | N/A |
Architecture: single-layer Transformer (pre-norm, d_model=128, 4 heads). Trained with AdamW on 50/50 train/test split. 5 seeds per task, 7–8 checkpoints each.
We extract activations at hook_resid_post (final residual stream, before LayerNorm/Unembed) and fit AA:
X ≈ W × Z
where:
Z: [k, d_model] — archetypes (extreme points of the convex hull)
W: [n_samples, k] — simplex weights (convex coefficients, rows sum to 1)
We inject reconstructed hidden states at the residual stream:
Input → Embed → [Attention] → [MLP] → RESID_POST → LayerNorm → Unembed → Logits
↑
h = W_target × Z
Four intervention types:
- Weight swap: Replace source's W with target's W → reconstruct h
- Single-archetype edit: Interpolate weight mass toward a target archetype (α ∈ [0,1])
- Z-vector swap: Swap one archetype vector while keeping W fixed
- Archetype permutation: Globally reorder archetypes to implement y → y+d
For overcomplete AA (k >> P), we merge archetypes without ground-truth labels:
- Compute behavior signatures: S_j = softmax(LayerNorm(Z_j) × W_unembed)
- Cluster archetypes by cosine similarity of S_j
- Threshold τ determines merge granularity (robust over wide plateau)
All results report mean ± std over 5 random seeds.
Purity, test accuracy, and causal efficacy co-emerge at the grokking transition:
| Step | Purity | Test Acc | Phase |
|---|---|---|---|
| 1000 | 0.196 | 0.2% | Memorization |
| 2000 | 0.207 | 0.3% | Memorization |
| 3000 | 0.276 | 14.6% | Grokking starts |
| 4000 | 0.795 | 100.0% | Post-grokking |
| 7000 | 0.869 | 100.0% | Converged |
(Addition, k=113, mean over 5 seeds)
AA purity scales monotonically with k and saturates near k=P. Merged overcomplete AA matches k=P performance.
| Method (k=113; parity k=2) | Addition | Multiplication | Parity |
|---|---|---|---|
| AA (ours) | 0.87 | 0.86 | 0.89 |
| PCA | ≤0.01 | ≤0.01 | — |
| k-means | ≤0.01 | ≤0.01 | 0.66 |
| NMF | ≈0.17 | ≈0.17 | — |
PCA and k-means stay near the uniform baseline for all k, NMF is flat at ≈0.17 (no peak at P).
AA weight swaps redirect model output with near-perfect accuracy:
| k | Addition | Multiplication | Parity |
|---|---|---|---|
| 16 | 37.0% | 42.0% | 100.0% |
| 64 | 89.6% | 89.4% | 100.0% |
| 113 | 99.0% | 97.8% | — |
| 128 | 100.0% | 99.8% | — |
Parity is swept up to k=64. Raw k=P=2 yields 60±23% due to a binary orientation ambiguity, which behavior merging resolves to 100% (see below).
Baselines at k=113:
| Control | Addition | Multiplication |
|---|---|---|
| Random simplex | 0.6% | 1.4% |
| k-means | 5.8% | 2.0% |
| PCA | 100.0% | 100.0% |
| Chance (1/P) | 0.88% | 0.88% |
PCA achieves 100% via linear reconstruction but lacks interpretable coordinates. k-means and random controls are at chance, confirming AA's causal specificity beyond mere reconstruction quality.
Behavior-based merging of overcomplete archetypes recovers exactly P groups:
| Addition | Multiplication | Parity | |
|---|---|---|---|
| Overcomplete k | 226 | 226 | 64 |
| Groups found | 113 | 113 | 2 |
| Aligned accuracy | 100.0% | 100.0% | 100.0% |
| Merged purity | 0.921 | 0.927 | 0.984 |
| Tau plateau | [0.30, 0.95] | [0.30, 0.95] | [0.10, 0.99] |
All 5 seeds recover exactly P groups with perfect alignment. Wide tau plateaus indicate robustness.
Archetype permutations implement the cyclic group action y → (y+d) mod P:
| Addition | Multiplication | |
|---|---|---|
| Permutation success (k=P direct) | 80–85% | 77–88% |
| Permutation success (k=2P→P merged) | 100.0% | 100.0% |
| Identity control | 99.1% | 99.4% |
| Random-perm control | 1.1% | 0.9% |
The behavior-merged representation achieves perfect permutation success across all 5 seeds and all tested offsets d. For multiplication, the circular structure is only visible under discrete-logarithm ordering of labels in (ℤ/Pℤ)* (label 0 excluded).
Single-archetype edits (k=113, α=1.0): Moving weight mass to a target archetype achieves 92.5% success (addition) and 88.6% (multiplication).
Z-vector swaps (k=113): Swapping archetype vectors achieves 100% success on samples with activation ≥ 0.30, confirming individual archetypes causally control specific output classes.
Fourier alignment: 88.5% of archetypes show strong Fourier mode correlation (>0.5), connecting to prior work on grokking via Fourier features.
Orthogonality: Add/mul archetypes are approximately orthogonal (mean cosine similarity ≈ -0.006). Parity archetypes cluster into 2 groups (mean similarity ≈ 0.595).
Cross-seed stability: Individual archetype vectors are NOT preserved across seeds (mean pairwise cosine ≈0.20), but the functional organization is: behavior merging recovers exactly P groups with 100% label alignment in every seed.
| Metric | Addition | Multiplication | Parity |
|---|---|---|---|
| Groups found (label-free) | 113/113 | 113/113 | 2/2 |
| Merge aligned accuracy | 100.0% | 100.0% | ~100.0% |
| Causal swap (k=P) | 99.0% | 97.8% | 100.0% (merged) |
| Single-archetype edit | 92.5% | 88.6% | — |
| Z-swap (active samples) | 100.0% | 100.0% | — |
| Permutation k=P direct | 80–85% | 77–88% | — |
| Permutation k=2P→P merged | 100.0% | 100.0% | — |
| Post-grokking purity | 0.87 | 0.86 | 0.89 |
This work connects several research threads:
-
Grokking: Power et al. (2022) discovered delayed generalization. Kumar et al. (2024) interpret it as a lazy-to-rich transition of the learning dynamics, and Nanda et al. (2023) identify Fourier modes in the learned weights. We contribute a complementary geometric lens — tracking how activations become well-approximated by a simplex whose vertices correspond to output classes.
-
Mechanistic Interpretability: The Transformer Circuits framework (Olah et al., 2021) and superposition results (Elhage et al., 2022) motivate structured decompositions of the residual stream. AA provides such a decomposition with built-in causal testability.
-
Concept-Based Explanations: TCAV (Kim et al., 2018) and ProtoPNet (Chen et al., 2019) explain predictions via concepts/prototypes. Our archetypes serve a similar role but with simplex-constrained coordinates and causal validation.
-
Causal Abstraction: Geiger et al. (2023) argue that useful internal representations should support counterfactual manipulation. Our weight swaps, Z-vector swaps, and archetype permutations directly test this criterion. Zhang & Nanda (2023) provide best practices for activation patching that inform our evaluation methodology.
-
Archetypal Analysis: Originally proposed by Cutler & Breiman (1994), with deep variants by Keller et al. (2019). We apply AA to Transformer internals and show that its unique properties (simplex weights, convex-hull extremes) translate to causal controllability — unlike PCA (unconstrained) or k-means (hard assignment).
Main figures (in paper_figures/):
- Fig. 2: Triple overlay — test accuracy, AA purity, and causal efficacy co-emerging at grokking
- Fig. 3: Method comparison — AA purity vs. k (with baselines and merged points)
- Fig. 4: Causal control — k-sweep accuracy with baseline controls
- Fig. 5: Label-free discovery — tau sweep showing group count convergence
- Fig. 6: Group actions — confusion matrices and permutation success
Appendix figures:
- Fig. A4: Cross-seed archetype stability
- Fig. A8: Training curves (loss + accuracy per task)
- Fig. A9: DFT-projected archetype circles for addition (linear label order, ℤ/Pℤ) and multiplication (discrete-log order, (ℤ/Pℤ)*, label 0 excluded)
- Python 3.9+
- NVIDIA GPU (RTX 3090 or similar recommended)
Optional power management for RTX 3090 (~250W instead of 350W, <10% performance loss):
sudo nvidia-smi -pm 1
sudo nvidia-smi -i 0 -lgc 0,1600git clone https://github.com/AppliedMachineLearning-Lab/grokking-aa.git
cd grokking-aa
conda create -n grokking-aa python=3.10 -y && conda activate grokking-aa
pip install -r requirements.txt# Run everything (phases 1-8, all operations):
./run_all_experiments.sh
# Or specific phases:
./run_all_experiments.sh 3 # Start from phase 3
./run_all_experiments.sh 3 5 # Run phases 3 through 5 onlyEach script supports --help for its individual options.
Output locations:
- Per-run results:
outputs/{operation}/run_{seed}/ - Aggregated metrics:
outputs/{operation}/aggregated/ - Paper figures:
paper_figures/
src/
├── config.py # Hyperparameters (ModelConfig, TrainConfig, AAConfig)
├── model.py # GrokkingTransformer with activation hooks
├── data.py # ModularDataset + sequence tasks (parity, random)
├── causal_utils.py # Shared causal intervention utilities
├── analysis/
│ ├── archetypes.py # Archetypal Analysis (gradient descent on simplex)
│ ├── baselines.py # PCA, k-means, NMF baselines
│ ├── behavior_merge.py # Label-free behavior-based archetype merging
│ └── geometry.py # Purity, Hungarian alignment, diversity metrics
├── plotting/
│ └── plots.py # Publication-quality figure generation
└── utils/
├── seeding.py # Reproducibility utilities
├── metrics_logger.py # CSV logging
└── aggregation.py # Multi-run metric aggregation
scripts/
├── 01_train_grokking.py # Phase 1: Train models until grokking
├── 02_run_aa.py # Phase 2: AA k-sweep + baselines
├── 03_behavior_merge.py # Phase 3: Behavior-based merging (label-free)
├── 04_causal_test.py # Phase 4: Causal interventions (swap/edit/z_swap/permutation)
├── 05_temporal_analysis.py # Phase 5: Temporal metric tracking
├── 06_visualize_archetypes.py # Phase 6: Archetype visualization
├── 07_aa_stability.py # Phase 7: Cross-seed stability
└── 08_paper_figures.py # Phase 8: Generate all paper figures
run_all_experiments.sh # Master pipeline orchestrator
tests/ # Comprehensive test suite
paper_figures/ # Generated figures (PDF + PNG, 300 DPI)
outputs/ # Experiment results (created by the pipeline)
@inproceedings{sparrenberg2026grokking,
title = {Grokking Is a Geometric Phase Transition: Archetypal Analysis
Reveals Causal Structure in Transformer Representations},
author = {Sparrenberg, Lorenz and Deu{\ss}er, Tobias and Sifa, Rafet},
booktitle = {Machine Learning and Knowledge Discovery in Databases (ECML PKDD)},
year = {2026}
}pytest tests/ -v # Run all tests
pytest tests/ --cov=src --cov-report=term-missing # With coverage
black src/ scripts/ tests/ # Format code
ruff check src/ scripts/ tests/ --fix # Lint and auto-fix