EG-trained modular arithmetic models learn novel representations
Abstract
Small transformers trained on modular arithmetic offer a setting in which we can reverse-engineer what a network has learned, and compare its internal representations against the known task structure. Trained with gradient descent, models embed numbers as points on a few circles in embedding space, with the numbers stepping around each circle at a different frequency, and implement specific forms (e.g. Clock'',Pizza'') of an abstract algorithm called the approximate Chinese Remainder Theorem (aCRT). Whether this solution reflects the task or the optimiser is unknown, as every study to date has used gradient descent. Here we train networks with Exponentiated Gradient (EG), a multiplicative update rule that obeys Dale's Law and produces log-normal weights. Applying the gradient and logit-based measures that distinguish Clock from Pizza, we find EG networks differ from both. We find that EG discovers the same aCRT algorithm as gradient descent, but with a different, non-circular embedding geometry. Furthermore, EG networks utilise the same frequency-based subcircuits as gradient descent, but use far fewer neurons per subcircuit. Altogether, the geometry of the optimiser selects the representational geometry, the realisation of the solution, and the number of neurons it utilises but not the abstract algorithm it implements.