|
|
--- |
|
|
title: Modular Addition Feature Learning |
|
|
emoji: "🔢" |
|
|
colorFrom: blue |
|
|
colorTo: yellow |
|
|
sdk: gradio |
|
|
sdk_version: "6.5.1" |
|
|
app_file: hf_app/app.py |
|
|
pinned: false |
|
|
--- |
|
|
|
|
|
# On the Mechanism and Dynamics of Modular Addition |
|
|
|
|
|
### Fourier Features, Lottery Ticket, and Grokking |
|
|
|
|
|
**Jianliang He, Leda Wang, Siyu Chen, Zhuoran Yang** |
|
|
*Department of Statistics and Data Science, Yale University* |
|
|
|
|
|
[[arXiv](https://arxiv.org/abs/2602.16849)] [[Blog](https://y-agent.github.io/posts/modular_addition_feature_learning/)] [[Interactive Demo](https://huggingface.co/spaces/y-agent/modular-addition-feature-learning)] |
|
|
|
|
|
--- |
|
|
|
|
|
## Overview |
|
|
|
|
|
This repository provides the code for studying how a two-layer neural network learns modular arithmetic $f(x,y) = (x+y) \bmod p$. We analyze three phenomena: |
|
|
|
|
|
1. **Fourier Feature Learning** — Each neuron independently discovers a cosine wave at a single frequency, collectively implementing a discrete Fourier transform that the network was never taught. |
|
|
2. **Lottery Ticket Dynamics** — Random initialization determines which frequency each neuron will specialize in: the frequency with the best initial phase alignment wins a winner-take-all competition. |
|
|
3. **Grokking** — Under partial data with weight decay, the network first memorizes, then suddenly generalizes through a three-stage process: memorization → sparsification → cleanup. |
|
|
|
|
|
An [**Interactive Demo**](https://huggingface.co/spaces/y-agent/modular-addition-feature-learning) on Hugging Face Spaces visualizes all results with 9 analysis tabs, interactive Plotly charts, and on-demand training for any odd $p \geq 3$. Pre-computed examples are included for $p = 15, 23, 29, 31$. |
|
|
|
|
|
### Launch Locally |
|
|
|
|
|
```bash |
|
|
pip install -r requirements.txt |
|
|
python hf_app/app.py |
|
|
# Opens at http://localhost:7860 |
|
|
``` |
|
|
|
|
|
### Deploy to Hugging Face Spaces |
|
|
|
|
|
We use the [Hugging Face Python API](https://huggingface.co/docs/huggingface_hub/) to upload to Spaces, since HF now requires [Xet storage](https://huggingface.co/docs/hub/xet) for binary files (PNGs, etc.) which standard `git push` does not handle. |
|
|
|
|
|
**First-time setup:** |
|
|
|
|
|
```bash |
|
|
pip install huggingface_hub hf_xet |
|
|
``` |
|
|
|
|
|
Log in (get a **write** token from https://huggingface.co/settings/tokens): |
|
|
|
|
|
```bash |
|
|
huggingface-cli login |
|
|
``` |
|
|
|
|
|
**Upload to the Space:** |
|
|
|
|
|
```bash |
|
|
python deploy_to_hf.py |
|
|
# or with a custom commit message: |
|
|
python deploy_to_hf.py --message "Update app" |
|
|
``` |
|
|
|
|
|
The deploy script prepends the required HuggingFace Space metadata (SDK config, app path, etc.) to `README.md` before uploading, so the GitHub README stays clean. |
|
|
|
|
|
**What gets uploaded:** Only the files the app needs — `hf_app/`, `precompute/`, `precomputed_results/`, `src/`, `requirements.txt`, `README.md`. Model checkpoints, notebooks, and figures are excluded. |
|
|
|
|
|
**On-demand training:** Users can generate results for new $p$ values directly from the app's "Generate" button. Streaming logs show real-time training progress. New results are auto-committed back to the Space repo so they persist across restarts. |
|
|
|
|
|
> **Tip:** For GPU-accelerated on-demand training, select a GPU runtime in your Space settings. |
|
|
|
|
|
## Pre-computation Pipeline |
|
|
|
|
|
The `precompute/` directory trains 5 model configurations per modulus and generates all plots + interactive JSON data. See [`precompute/README.md`](precompute/README.md) for full documentation. |
|
|
|
|
|
### Quick Start |
|
|
|
|
|
```bash |
|
|
# Full pipeline for a single modulus (train → plots → analytical → verify) |
|
|
bash precompute/run_pipeline.sh 23 |
|
|
|
|
|
# With custom d_mlp |
|
|
bash precompute/run_pipeline.sh 23 --d_mlp 128 |
|
|
|
|
|
# Delete checkpoints after generating plots (saves disk space) |
|
|
CLEANUP=1 bash precompute/run_pipeline.sh 23 |
|
|
|
|
|
# Batch: all odd p in [3, 99] |
|
|
bash precompute/run_all.sh |
|
|
|
|
|
# Or up to p=199 |
|
|
MAX_P=199 bash precompute/run_all.sh |
|
|
``` |
|
|
|
|
|
### Manual Steps |
|
|
|
|
|
```bash |
|
|
# Step 1: Train all 5 configurations |
|
|
python precompute/train_all.py --p 23 --output ./trained_models --resume |
|
|
|
|
|
# Step 2: Generate model-based plots (21 PNGs + 7 JSONs) |
|
|
python precompute/generate_plots.py --p 23 --input ./trained_models --output ./precomputed_results |
|
|
|
|
|
# Step 3: Generate analytical simulation plots (2 PNGs, no model needed) |
|
|
python precompute/generate_analytical.py --p 23 --output ./precomputed_results |
|
|
``` |
|
|
|
|
|
### Output |
|
|
|
|
|
Each modulus produces ~33 files in `precomputed_results/p_XXX/`: |
|
|
|
|
|
| Category | Files | Description | |
|
|
|----------|-------|-------------| |
|
|
| Overview (Tab 1) | 2 PNGs + 1 JSON | Loss, IPR, phase scatter | |
|
|
| Fourier Weights (Tab 2) | 3 PNGs + 1 JSON | DFT heatmaps, cosine fits, neuron spectra | |
|
|
| Phase Analysis (Tab 3) | 3 PNGs | Phase distribution, alignment, magnitudes | |
|
|
| Output Logits (Tab 4) | 1 PNG + 1 JSON | Logit heatmap, interactive explorer | |
|
|
| Lottery Mechanism (Tab 5) | 3 PNGs | Magnitude race, phase convergence, contour | |
|
|
| Grokking (Tab 6) | 5 PNGs + 3 JSONs | Loss/acc curves, memorization, weight evolution | |
|
|
| Gradient Dynamics (Tab 7) | 4 PNGs | Phase alignment + DFT for Quad and ReLU | |
|
|
| Decoupled Simulation (Tab 8) | 2 PNGs | Analytical ODE integration | |
|
|
| Metadata | 2 JSONs | Config + training log | |
|
|
|
|
|
> **Note:** Grokking results (Tab 6) require $p \geq 19$. Smaller values of $p$ have too few data points for a meaningful train/test split. |
|
|
|
|
|
## The 5 Training Configurations |
|
|
|
|
|
| Config | Activation | Optimizer | LR | Weight Decay | Data | Epochs | Used In | |
|
|
|--------|-----------|-----------|-----|-------------|------|--------|---------| |
|
|
| `standard` | ReLU | AdamW | 5e-5 | 0 | 100% | 5,000 | Tabs 1–4 | |
|
|
| `grokking` | ReLU | AdamW | 1e-4 | 2.0 | 75% | 50,000 | Tabs 1, 6 | |
|
|
| `quad_random` | Quad | AdamW | 5e-5 | 0 | 100% | 5,000 | Tab 5 | |
|
|
| `quad_single_freq` | Quad | SGD | 0.1 | 0 | 100% | 10,000 | Tab 7 | |
|
|
| `relu_single_freq` | ReLU | SGD | 0.01 | 0 | 100% | 10,000 | Tab 7 | |
|
|
|
|
|
## Running a Single Experiment |
|
|
|
|
|
For custom experiments outside the pre-computation pipeline: |
|
|
|
|
|
```bash |
|
|
cd src |
|
|
|
|
|
# Train with default config (p=97, d_mlp=1024, ReLU, 5000 epochs) |
|
|
python module_nn.py |
|
|
|
|
|
# Train with specific parameters |
|
|
python module_nn.py --p 23 --d_mlp 512 --num_epochs 5000 --lr 5e-5 |
|
|
|
|
|
# Dry run: see config without training |
|
|
python module_nn.py --dry_run --p 23 --d_mlp 512 |
|
|
``` |
|
|
|
|
|
## Notebooks |
|
|
|
|
|
Interactive analysis notebooks in `notebooks/`: |
|
|
|
|
|
| Notebook | Description | |
|
|
|----------|-------------| |
|
|
| `empirical_insight_standard.ipynb` | Fourier weight analysis, phase distributions, output logits | |
|
|
| `empirical_insight_grokk.ipynb` | Grokking stages, weight dynamics, IPR evolution | |
|
|
| `lottery_mechanism.ipynb` | Neuron specialization, frequency magnitude/phase tracking | |
|
|
| `interprete_gd_dynamics.ipynb` | Phase alignment under single-frequency initialization | |
|
|
| `decouple_dynamics_simulation.ipynb` | Analytical gradient flow simulation | |
|
|
|
|
|
## Setup |
|
|
|
|
|
### Requirements |
|
|
|
|
|
- Python 3.8+ |
|
|
- PyTorch 2.0+ |
|
|
- CUDA-capable GPU (recommended for $p > 50$; CPU works for small $p$) |
|
|
|
|
|
### Installation |
|
|
|
|
|
```bash |
|
|
git clone https://github.com/Y-Agent/modular-addition-feature-learning.git |
|
|
cd modular-addition-feature-learning |
|
|
pip install -r requirements.txt |
|
|
``` |
|
|
|
|
|
## Project Structure |
|
|
|
|
|
``` |
|
|
modular-addition-feature-learning/ |
|
|
├── src/ # Core source code |
|
|
│ ├── module_nn.py # Training script with CLI |
|
|
│ ├── nnTrainer.py # Training loop and optimization |
|
|
│ ├── model_base.py # Neural network architecture (EmbedMLP) |
|
|
│ ├── mechanism_base.py # Fourier analysis and decomposition |
|
|
│ ├── utils.py # Configuration and helpers |
|
|
│ └── configs.yaml # Default hyperparameters |
|
|
├── precompute/ # Batch training and plot generation |
|
|
│ ├── run_pipeline.sh # Full pipeline for one modulus |
|
|
│ ├── run_all.sh # Batch pipeline for all odd p |
|
|
│ ├── train_all.py # Train 5 configurations |
|
|
│ ├── generate_plots.py # Generate model-based plots + JSONs |
|
|
│ ├── generate_analytical.py # Analytical ODE simulation plots |
|
|
│ └── prime_config.py # Configurations and sizing formula |
|
|
├── hf_app/ # Gradio web application |
|
|
│ └── app.py # Interactive visualization app |
|
|
├── precomputed_results/ # Pre-computed plots and data |
|
|
│ ├── p_015/ # Results for p=15 |
|
|
│ ├── p_023/ # Results for p=23 |
|
|
│ ├── p_029/ # Results for p=29 |
|
|
│ └── p_031/ # Results for p=31 |
|
|
├── notebooks/ # Analysis and visualization notebooks |
|
|
├── requirements.txt # Python dependencies |
|
|
└── README.md |
|
|
``` |
|
|
|
|
|
## Citation |
|
|
|
|
|
```bibtex |
|
|
@article{he2025modular, |
|
|
title={On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking}, |
|
|
author={He, Jianliang and Wang, Leda and Chen, Siyu and Yang, Zhuoran}, |
|
|
journal={arXiv preprint arXiv:2602.16849}, |
|
|
year={2025} |
|
|
} |
|
|
``` |
|
|
|
|
|
## License |
|
|
|
|
|
[MIT License](LICENSE) |
|
|
|
