Summary

I constructed a neural network classifier whose inputs were integers , , and their prime factorizations. Using three layers: , the classifier learned in epochs to predict with accuracy whether the quaternion algebra defined over the rational numbers was a division algebra or a matrix algebra.

The underlying mathematics of quaternion algebras over the rationals makes it natural to wonder if the neural network was somehow capturing the Hilbert symbol of at various primes. I investigated this using PyTorch’s hooks into the activation vectors, balanced linear (logistic) regression as well as a two-layer MLP, by checking if the two activation vectors could predict the Hilbert symbol at the primes and . I used a fresh sample of 2000 - pairs, and used a 80/20 train/validation split for both the logistic regression and the MLP probe.

Both the linear and the non-linear probes performed much worse than the majority class baseline (simply guessing the symbol to be all the time). This suggests that the network is learning by some other mechanism.

Motivation

Driven by curiosity about neural networks (and more broadly LLMs) and wanting some hands-on learning, I decided to code a three layer neural network that would take in non-zero integers and and their prime factorizations, and predict if the quaternion algebra generated over the rational numbers is a division algebra or is isomorphic to the matrix algebra with rational coefficients. The choice of this problem arose from my own background of research into division algebras and allied areas.

Here is a brief introduction to the mathematics of these objects: With and as above, the quaternion algebra over is the four-dimensional algebra over consisting of all -linear combinations , where , , and are formal symbols satisfying the multiplication relations: , , and . The last relation introduces noncommutativity in the multiplication in .

There are deep theorems about these objects. The first is that they fall into two categories: given a quaternion algebra , either it is a division algebra (all nonzero elements are multiplicatively invertible, so one can do division), or, up to a disguise, it is just the set of matrices with rational coefficients. (In the second case is said to be split.) The second theorem of interest is the local-global principle: the algebra is split if and only if it is split when you extend scalars to as well as to all the -adic completions . The third theorem is that the Hilbert symbol, denoted , which is formulated in terms of the prime factorizations of and and is computed based on some local happenings in the finite fields , is if is a division algebra over and otherwise. Undergirding the Hilbert symbol result is a theorem that a quaternion algebra over is split if and only if is a norm from to .

Given the second and third results above, it is natural to wonder if the neural network is using its knowledge of the prime factorizations of and to somehow capture the Hilbert symbol at various primes. It is this that motivated me to run some quick experiments to see if something like this is going on.

This was a learning-by-doing project for me, and I have to thank ChatGPT's and Gemini’s LLMs for guiding me and helping me with vibe-coding (although, the kludginess of the final code is entirely attributable to me).

The Network, and the Probes

Network

I constructed the neural net with three layers: a (input layer) matrix followed by ReLU, a followed by ReLU, and a followed by sigmoid representing the output. The input for training consisted of nonzero pairs of integers and randomly generated, tested to be within in magnitude and having no more than distinct prime factors (for ease of processing), along with their prime factorization (both the primes and their exponents). There was an training/validation split. The training input was constructed so that the numbers of division algebras and matrix algebras were roughly balanced. The output was the class of : whether it was a division algebra or a matrix algebra. In epochs with a learning rate of , it had learned to predict with a validation accuracy of and a training accuracy of (validation loss of and training loss of ).

The Probes

I used a fresh dataset of elements and ran it through the model trained as above. I extracted the activations after layer 1 (the layer) and after layer 2 (the layer). I precalculated the Hilbert symbols at the five primes and for each of these algebras, and I checked for mathematical correlation between the activation and the Hilbert symbols at these primes. I used an train/validation split.

I used both a linear probe and a non-linear one. For the linear probe I used logistics regression (with class balancing, as the Hilbert symbol would have as its dominant value). For the non-linear one, I trained a separate two-layer Neural network with nodes in the hidden layer followed by ReLU and a learning rate of .

The Results

Discussions, Interpretability

The chief result is negative: there does not appear to be a mathematical correlation between the activation vectors and the local Hilbert symbols at the first five primes. We checked for both linear and non-linear relations, and in all cases, we found that our probes’ accuracies were below majority baseline numbers (which in our case meant simply guessing the Hilbert symbol to be all the time). Our results are summarized in the table above. Although not a mathematical proof, our experiment suggests that the learning in the neural network is not happening by learning the local Hilbert symbols.

This would lead to probing for other learning keys. The network seems to use the prime factorizations of and , because some brief experiments made earlier where the network was redesigned to take in only and (without the prime factorizations) gave very poor accuracies. On the other hand, any phenomenon it is probing over the -adic numbers would ultimately be correlated with the Hilbert symbol, so we should have seen that correlation already. Thus, it appears to be using the primes but does not appear to be going to the fields .

Here is an idea: the equation for to be a norm from to reduces to the diophantine equation , where we may suitably factor the coefficients and assume, abusing notation, that , , and are pairwise coprime. Legendre has a theorem that says exactly when there are integers , and that satisfy this. It involves checking if is a quadratic residue mod , is a quadratic residue mod and is a quadratic residue mod . But for this, due to the multiplicativity of the Legendre symbol, it is sufficient to check if these same integers are squares mod the various prime factors of and — AHA! So here is a route that needs the prime factorization but does not use the fields at all! Instead, the network may be going directly into the arithmetic of the finite fields ! And to put this in perspective, even the Hilbert symbols ultimately depend on arithmetic in , so maybe it is what is going on in that the network cares about. If true, that would be awesome!

Future Work

One pathway would be to train the model over 100,000 or more epochs and hope for grokking behavior. If this happens, it is possible that the sudden boost in accuracy is precisely because the network has learned the Hilbert symbols at various primes, and the same experiment repeated then might confirm that.

My gurus (ChatGPT's and Gemini's LLMs) tell me there is a technique called Integrated Gradients that will give some handle on how much the various inputs affect the output. At some point I need to explore this technique to try to get some handle on what is going on.

But more promising and definitely more exciting is to explore Legendre’s solution to a rational point on a ternary quadric (like the norm equation). The grokking idea is a matter of letting the computer run (warming the planet in the process, sadly), but studying if the network is learning the inputs to Legendre’s test is mathematically richer. I will be doing this next.