Open-sourcing math AI models, and why it matters
Mathematics is, at heart, a human endeavor, it is created (or is it discovered?) by human mathematicians, for the benefit of other human mathematicians. As new technologies develop, and new tools appear, the role of the mathematician evolves, tasks get automated, complex problems become trivial. Yet, mathematics remains a human endeavor.
Advances in artificial intelligence opens new perspectives for computer-assisted mathematics. This is reflected in the many benchmarks the AI community maintains to measure progress. However, the litmus test of these methods is their adoption by the mathematical community. This can prove a daunting task. The many architectural variants, a research corpus that grows exponentially, the necessity of adapting techniques developed for image or natural language processing to mathematics, and the technicalities of running such models on GPU, create a very steep learning curve for the interested mathematician. And it comes as no surprise that the question we hear most often is: where do I start?
At Axiom, we believe that AI-assisted mathematical discovery will only happen at scale when a large number of mathematicians have access to open-source models that they can use without extensive knowledge of current AI research. Axplorer (pronounced explorer), a generative AI tool for creating interesting mathematical objects, is a first step in this direction.
All happy families are alike
Suppose you want to discover rare mathematical objects, special members of a discrete class, elements of a family that maximize some quantity (like the number of edges in a graph) without violating a specific constraint. You know how to evaluate candidates, but have little guidance, and little theory, on how to generate good ones. Enumeration, generating a lot of random solutions and retaining the best, appears to be your only option.
After generating many candidates, you might notice that the best ones are subtly related. These common patterns are elusive, and hard to characterize. Sometimes they only exist at a statistical level, as higher than chance level correlations. In other words, the distribution of good solutions is different from the distribution of random candidate solutions.
An AI model, provided with a training set of such good candidate solutions, can learn their distribution (this is the point of gradient-based learning), and then be used to generate more good solutions.
This is the basic idea behind PatternBoost, a 2024 paper by François Charton, Jordan Ellenberg, Adam Wagner and Geordie Williamson. It follows a simple recipe: start from generated good candidates, use them to train a model, use the model to generate more good candidates, keep the best, and use them to train the model. It was shown to be able to provide competitive solutions to hard combinatorial problems.
Axplorer is an efficient and easy to use implementation of PatternBoost. It can be easily adapted to the problem of your choice, and can run on a laptop with a GPU (e.g. a MacBook pro).
The Apprentice (AI model) and the Sculptor (local search)
Axplorer leverages the training-generation paradigm proposed by PatternBoost: a self improvement cycle.
Generate initial data (Done only once): using a problem specific algorithm, generate random valid examples, score them, and keep the best. They will serve at the initial training set of the generative model.
Train the model. Encode the candidate solutions as sequences of words in some finite language (tokenization), and use them to train a decoder-only transformer, predicting the next token in the sequence. The model will learn the patterns that define good candidates.
Sample new candidates. Use the trained model to generate new candidate solutions, by predicting the next words of a random seed. This creates a new set of candidates.
Local search. The new candidates generated by the model do not always respect the constraints of the problem, or are locally optimal. Use a local search algorithm to fix and improve them.
Sample selection. Score all new candidates, combine them with the existing data and retain the best, as the new training data. Go back to step 2.
In PatternBoost, the AI (training and sampling) and local search parts were clearly separated. They used different languages which caused inefficiencies. Standard natural language methods, like Byte-Pair-Encoing tokenization, proved suboptimal. As a result, the proof of concept model would require extensive computation, and weeks of training, and would fail on harder problems.
For Axplorer, we wanted a toolkit that could be used by mathematicians with very little AI experience, and run in a reasonable amount of time (hours) without expensive compute resources. We dismantled the engine to its minimal components, optimized it for speed, and experimented with alternative configurations that allow for stronger results.
We are happy to report that one does not need a supercomputer to do state-of-the-art mathematical discovery. All you need is a smarter approach to learning.
Optimizing for Efficiency
Consider the Turán 4-cycles problem. Imagine you are drawing a graph and you want to connect as many nodes as possible with edges, but you are forbidden from creating any squares, a loop of four nodes where A is connected to B, B to C, C to D and D to A. How many edges can you add before you are forced to create a square?
For small graphs, we have detailed answers. But as the number of nodes $N$ grows, the complexity explodes. The number of possible graphs when $N=25$ is already higher than the number of atoms in the universe. Yet somewhere in the Hypercube of possibilities, there are a handful of examples that maximizes the edges without breaking the rule.
Original PatternBoost found the best known solution for $N=33$ nodes, but it was an exhausting effort. It had to generate 116.5 million candidates to find the needle. Axplorer found the optimal solution in just 2.5 hours on a single L4 GPU. It didn’t need 116 million attempts; it needed only 2.6 million.
How did we achieve a nearly 100x efficiency gain?
Don’t describe the void
There is a rule in language models: short, informative sentences are usually easier to learn than long verbose ones. This is because the attention mechanism that the model uses to figure the relations between words in the sentence, is quadratic in the length of the sentence.
This general principle can be applied to graphs. How do you represent a graph, say friends on social media? You could have a list of all pairs of members, and their friendship status, or just a list of the pairs of friends.
PatternBoost used a dense encoding: graphs were represented by their adjacency matrix, a list of all possible relations between nodes, most of them zeroes. For Axplorer, we switched to sparse encoding: graphs are represented by their edges only.
This brings huge benefits: on a graph of $N=30$ nodes, PatternBoost struggles for 30 epochs (generating 15 millions candidates) before finding a sub-optimal solution (81 edges). With sparse encoding, the optimal solution (85 edges) is found in 4 epochs (2 million candidates).
The temperature trap
Exploration vs exploitation is an old paradigm in optimization. Given a fixed amount of resources, and a quantity to be maximized, how should we balance between exploiting known but possibly subpar solutions, and exploring the space in the hope of finding better solutions? The typical answer is that one should start mostly exploring, and increase the exploitation budget as better solutions are found.
Such a dilemma exists in language models, via a parameter named temperature, which controls how solutions are generated from the word distribution learned by the model. A low temperature will generate the best solutions according to the probability learned by the model, and this will happen at the expense of diversity: with very low temperatures, successive calls to the models tend to generate duplicates. This is exploitation. With a high temperature, the learned distribution takes second role to diversity, and the model will generate a broad set of candidates, usually less efficient. This is exploration.
In Axplorer, we notice that a low temperature (e.g. $T=0.6$) is required at an early stage. This forces the model to respect the rules of the problem, and focus on the best parts of search space. After some time, however, this tends to use a drop in the diversity of the candidate solutions, that can be measured by the apparition of duplicates in the generated data. At this point, we increase the temperature, to keep a healthy pool of generated examples.
This empirical result is intriguing, as it is a departure from usual exploration-exploitation wisdom. In Axplorer, we exploit first, and then explore.
Is local search needed?
Perhaps the most critical insight of Axplorer is understanding the role of the language model. Watching Axplorer perform, one may conclude that the AI-model does most of the work and that the local search is relegated to a minor, fine-tuning, role. Could we, then, do without local search?
The answer is a definitive no. Even late during training, most of the solutions produced by the generative model are not valid solutions to the problem: they violate imposed constraints (e.g. the absence of four-cycles). Without local search, very few solutions are discovered, and the best solutions are never found.
The true power of Axplorer is the collaboration between a local search algorithm, that will make the actual discovery be provided with a close enough starting point, and a generative model that explores possible starting points, essentially making the local search global.
Getting started
The true power of Axplorer isn’t just in its efficiency; it’s in its accessibility. We designed a pipeline that is plug-and-play.
For the one time setup, you firstly need to create an environment:
conda env create -f environment.yml conda activate env_axplorer
You can launch a full training run from the command line without typing a single line of code. To hunt for square-free graph of size 30:
python main.py –env_name square –exp_name square_experiment –N 30
But don’t stop there! We think the real value of this tool is not to solve the problems we have set, but to solve yours. We have provided a tutorial notebook within our repository that shows you how to add a new problem. You simply need to code the following
The object: What does your object look like? A graph, a list of numbers, a grid?
The rules: What makes the object invalid? For example, a cycle of length 4.
The local search: How do you want to fix a broken object?
We provide two templates to show you how to implement the code for two other problems. We want you to focus on mathematical logic; Axplorer handles the tokenization, the neural network training and the search strategy.
The future is open
We applied Axplorer to three classical extremal combinatorics problems. These are problems where the search space is so vast that brute force is impossible, and standard optimization algorithms often get stuck in local optima.
Turán 4-cycles: At most, how many edges can a graph on $N$ vertices have, if it doesn’t contain a cycle of length 4? We recovered best-known solutions for $N \leq 40$ in less than a day.”
No 5 points on a sphere: In a grid $N \times N \times N$, how many vertices can we choose such that no 5 points lie on a sphere? Here, we matched the best known solutions for $7 \leq N \leq 11$.
Isosceles-free sets: How many points can we choose from a $N \times N$ grid such that no 3 points form an isosceles triangle? On this problem, we achieved the best known results for $N=64$ and competitive results for $N=100$ with a fraction of the compute used by other methods.
And the cost? To solve the Turán 4-cycles problem for $N=33$ that once took weeks, we spent 2.5 hours and 3 dollars in cloud computing.
But the real victory is what this means for the future. You don’t need a research grant and a supercomputer to push the boundaries of what is known. We open-sourced the Axplorer codebase to put this engine into the hands of students and researchers. We give every mathematician the power to construct their own counterexamples, right on their laptop.
Happy hunting!