Learning the greatest common divisor: explaining transformer predictions
This work provides insights into transformer interpretability for a specific mathematical task, but it is incremental as it focuses on a narrow, synthetic problem.
The study investigated how small transformers learn to compute the greatest common divisor (GCD) of two integers, finding that their predictions can be fully explained by learning a specific list of integers, with performance varying based on training distributions: uniform operands achieved up to 38 GCDs ≤100, log-uniform operands improved to 73, and log-uniform outcomes reached 91, though uniform GCD training disrupted explainability.
The predictions of small transformers, trained to calculate the greatest common divisor (GCD) of two positive integers, can be fully characterized by looking at model inputs and outputs. As training proceeds, the model learns a list $\mathcal D$ of integers, products of divisors of the base used to represent integers and small primes, and predicts the largest element of $\mathcal D$ that divides both inputs. Training distributions impact performance. Models trained from uniform operands only learn a handful of GCD (up to $38$ GCD $\leq100$). Log-uniform operands boost performance to $73$ GCD $\leq 100$, and a log-uniform distribution of outcomes (i.e. GCD) to $91$. However, training from uniform (balanced) GCD breaks explainability.