Can Neural Networks Learn Small Algebraic Worlds? An Investigation Into the Group-theoretic Structures Learned By Narrow Models Trained To Predict Group Operations

TL;DR

This work asks whether narrow neural nets trained to predict a group operation on a finite set can learn and reveal abstract group-theoretic structures such as commutativity, the identity, and subgroup structure. It deploys three detection modalities—training dynamics, generalization to unseen elements, and internal representations—across cyclic, symmetric, and dihedral groups using MLPs and decoder-only transformers. The findings show partial evidence for commutativity in some settings and clear access to subgroup structure via linear probes, while the identity element remains elusive to extraction with the proposed methods. Overall, the results support the view that small networks can encode rich mathematical world models, suggesting avenues for AI-assisted mathematical discovery while underscoring challenges in interpreting such models without targeted probes or broader hyperparameter exploration.

Abstract

While a real-world research program in mathematics may be guided by a motivating question, the process of mathematical discovery is typically open-ended. Ideally, exploration needed to answer the original question will reveal new structures, patterns, and insights that are valuable in their own right. This contrasts with the exam-style paradigm in which the machine learning community typically applies AI to math. To maximize progress in mathematics using AI, we will need to go beyond simple question answering. With this in mind, we explore the extent to which narrow models trained to solve a fixed mathematical task learn broader mathematical structure that can be extracted by a researcher or other AI system. As a basic test case for this, we use the task of training a neural network to predict a group operation (for example, performing modular arithmetic or composition of permutations). We describe a suite of tests designed to assess whether the model captures significant group-theoretic notions such as the identity element, commutativity, or subgroups. Through extensive experimentation we find evidence that models learn representations capable of capturing abstract algebraic properties. For example, we find hints that models capture the commutativity of modular arithmetic. We are also able to train linear classifiers that reliably distinguish between elements of certain subgroups (even though no labels for these subgroups are included in the data). On the other hand, we are unable to extract notions such as the concept of the identity element. Together, our results suggest that in some cases the representations of even small neural networks can be used to distill interesting abstract structure from new mathematical objects.

Figures (3)

• Figure 1: Plots of symmetric consistency vs. equal value consistency (Left) and OOD symmetric consistency vs. (Right) for five transformers trained to perform modular arithmetic in $C_{100}$. Shaded regions correspond to $95\%$ confidence intervals.
• Figure 2: A plot of the symmetric representational similarity (solid lines) and equal value similarity (dashed lines) at different blocks of a decoder only transformer trained on the group $C_{100}$. Shading corresponds to 95% confidence intervals over 5 random initializations.
• Figure 3: Linear probe performance on a (Left) transformer trained to predict the binary of $S_5$ and a (Right) MLP trained to predict the binary operation on $D_{30}$. Solid lines are probe accuracy when labels correspond to the alternating subgroup and rotation subgroup respectively, while the dashed line corresponds to a probe trained on a random labeling of the group. In the right plot, shading indicates 95% confidence intervals over three random initializations (transformer performance on $S_5$ varied too much between random initializations to be useful in the case of the transformer). In the case of the MLP, probing was performed after each ReLU layer, in the transformer it was performed after each attention layer.