Transformers for Mathematics Tutorial

Simons Institute and SLMath Joint Workshop: AI for Mathematics and Theoretical Computer Science, April 8 2025

[Slides]

Contents

We provide Google Colab notebooks that have the required environment and files.

1. Bigram model:

   • A very simple language model based on counting consecutive tokens.

2. Transformer:

   • Implement and train a simple Transformer language model.

3. Addition:

   • Train a model for four-digit addition.

4. Triangle-free graphs:

   • Train a model to generate triangle-free graphs.

Note: you can safely skip notebook 1 and 2 if you are primarily interested in using existing transformer libraries for different applications.

(Optional) Local Setup

If you would like to run the notebooks locally instead of on Colab please follow these instructions.

Install uv. See installation instructions

Execute, in the terminal

uv sync
uv venv

These two commands download dependencies and create a virtual environment. The second command shows the instruction for activating the virtual environment. After activating the venv, check that Torch is available by

python3 -c "import torch"

To launch Jupyter notebook, run inside the virtual environment

jupyter notebook