Learning the greatest common divisor: explaining transformer predictions

TL;DR

This paper investigates how small transformers learn to compute the GCD of two integers, revealing that their predictions are based on learned divisibility patterns and how training data distribution affects both performance and interpretability.

Contribution

It provides a detailed characterization of transformer behavior in GCD computation and shows how training data distribution influences both accuracy and explainability.

Findings

Transformers predict the largest divisor from a learned list of products.

Training with uniform operands limits GCD learning to small values.

Log-uniform training distributions significantly improve GCD prediction accuracy.

Abstract

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.