Research map / Architectures and efficiency / Attention and transformers
Learning the greatest common divisor: explaining transformer predictions
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains small 4-layer transformers to compute GCD of integers, shows predictions are fully characterized by simple rules (sieve-like clustering by divisors of the base), and finds log-uniform operand and outcome distributions improve the number of correct GCD up to 91.
Problem
Whether small transformers can learn exact greatest common divisor computation, and how their predictions can be explained and improved by training distribution choices.
Approach
Model predictions can be fully characterized: the transformer clusters input pairs by GCD class and predicts the largest learned divisor (products of base divisors and later grokked small primes) dividing both inputs. Log-uniform training distributions of operands and outcomes accelerate learning and raise the number of correctly learned GCD.
Key results
- Base 2401 with log-uniform operands reaches 99.1% accuracy and 73 correct GCD under 100, versus 61-97% and 2-38 GCD with uniform operands.
- Log-uniform operands and log-uniform outcomes give models with base above 1000 predicting 87 to 91 GCD under 100, e.g. base 1000 with 91 correct GCD.
- Uniform outcomes lead to unstable, non-deterministic predictions, though up to 95 GCD are eventually predicted for base 1000.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use log-uniform operands instead of uniform operands | Transformer GCD training, base 10 · Stratified GCD ≤100 | correct GCD count: 13 -> 48 | Helped |
| Use log-uniform outcomes instead of natural outcomes | Transformer GCD training, base 1000 · Stratified GCD ≤100 | correct GCD count: 71 -> 91 | Helped |
| Use uniform outcomes instead of unbalanced outcomes | Transformer GCD training · GCD prediction explainability | deterministic predictions | Hurt |
| Use inverse square root outcomes instead of natural inverse square outcomes | Transformer GCD training with log-uniform operands · Stratified GCD ≤100 | correct GCD count: For 6 bases, 92 GCD under 100 are predicted correctly. | Helped |
| Use log-uniform outcomes sampled up to 200 instead of up to 100 | Transformer GCD training with log-uniform operands · Stratified GCD ≤100 | correct GCD count: 90 -> 85 | Hurt |
| Add 5% uniformly sampled GCD to natural training outcomes | Transformer GCD training, base 30 · Stratified GCD ≤100 | epochs to learn 27 GCD: 660 -> 175 | Helped |
| Use uniformly distributed GCD outcomes instead of unbalanced outcomes | Transformer GCD training, base 10 · GCD prediction explainability | prediction stability across epochs | Hurt |
| Reduce transformer size from the 4-layer baseline to 1 layer and 32 dimensions (ablation) | Transformer GCD, base 30 · Natural test set | accuracy (%): 94.7 -> 93.3 | No clear effect |
| Scale from the 4-layer baseline to 24 layers and 1024 dimensions (ablation) | Transformer GCD, base 30 · Natural test set | accuracy (%): 94.7 -> 93.4 | No clear effect |
| Use batch size 64 instead of 256 (ablation) | Transformer GCD training with log-uniform operands, base 1000 · Stratified GCD ≤100 | correct GCD count: 71 -> 70 | No clear effect |
| Use base 420 instead of base 11 | Transformer GCD with uniform operands · Natural test set | accuracy (%): 61.8 -> 96.8 | Helped |
| Train on uniform rather than log-uniform operands | Transformer GCD, base 2401 · Stratified GCD ≤100 | correct GCD count: 14 -> 73 | Hurt |
Surprising findings
- High accuracy (95%) can be obtained by models that predict only a handful of GCD, so accuracy on random test sets is a misleading metric.
- Uniform outcome distributions break explainability: predictions per GCD become split and vary from epoch to epoch.
Limitations
- Only GCD of positive integers up to 10^6 studied; rational arithmetic tasks (addition, multiplication, simplification) are not learned.
- Explainability breaks down with uniform outcome distributions and during grokking of new primes.
- Experiments run on a single V100 GPU with limited seeds (best of 3 or 6).
Experimental setup
Datasets: synthetic GCD pairs sampled uniformly from 1 to 10^6 log-uniform operand samples stratified GCD test set
Benchmarks: natural test set (uniform pairs 1 to 10^6) stratified test set (GCD 1 to 100)
Metrics: accuracy number of correct GCD under 100
Baselines: LSTM GRU natural GCD distribution training
Builds on: Adam grokking
Closest papers by meaning
- Teaching Arithmetic to Small Transformers (ICLR 2024 · similarity 0.82)
- Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks (NeurIPS 2024 · similarity 0.81)
- TRANSFORMER EXPLAINER: Interactive Learning of Text-Generative Models (AAAI 2025 · similarity 0.81)
- (How) Can Transformers Predict Pseudo-Random Numbers? (ICML 2025 · similarity 0.81)
- Understanding Transformers via N-Gram Statistics (NeurIPS 2024 · similarity 0.81)
- Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count (ICLR 2025 · similarity 0.80)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.