y-agent/modular-addition-feature-learning

1---2title: Modular Addition Feature Learning3emoji: "🔢"4colorFrom: blue5colorTo: yellow6sdk: gradio7sdk_version: "6.5.1"8app_file: hf_app/app.py9pinned: false10---11 12# On the Mechanism and Dynamics of Modular Addition13 14### Fourier Features, Lottery Ticket, and Grokking15 16**Jianliang He, Leda Wang, Siyu Chen, Zhuoran Yang**17*Department of Statistics and Data Science, Yale University*18 19[[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)]20 21---22 23## Overview24 25This 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:26 271. **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.282. **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.293. **Grokking** — Under partial data with weight decay, the network first memorizes, then suddenly generalizes through a three-stage process: memorization → sparsification → cleanup.30 31An [**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$.32 33### Launch Locally34 35```bash36pip install -r requirements.txt37python hf_app/app.py38# Opens at http://localhost:786039```40 41### Deploy to Hugging Face Spaces42 43We 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.44 45**First-time setup:**46 47```bash48pip install huggingface_hub hf_xet49```50 51Log in (get a **write** token from https://huggingface.co/settings/tokens):52 53```bash54huggingface-cli login55```56 57**Upload to the Space:**58 59```bash60python deploy_to_hf.py61# or with a custom commit message:62python deploy_to_hf.py --message "Update app"63```64 65The deploy script prepends the required HuggingFace Space metadata (SDK config, app path, etc.) to `README.md` before uploading, so the GitHub README stays clean.66 67**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.68 69**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.70 71> **Tip:** For GPU-accelerated on-demand training, select a GPU runtime in your Space settings.72 73## Pre-computation Pipeline74 75The `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.76 77### Quick Start78 79```bash80# Full pipeline for a single modulus (train → plots → analytical → verify)81bash precompute/run_pipeline.sh 2382 83# With custom d_mlp84bash precompute/run_pipeline.sh 23 --d_mlp 12885 86# Delete checkpoints after generating plots (saves disk space)87CLEANUP=1 bash precompute/run_pipeline.sh 2388 89# Batch: all odd p in [3, 99]90bash precompute/run_all.sh91 92# Or up to p=19993MAX_P=199 bash precompute/run_all.sh94```95 96### Manual Steps97 98```bash99# Step 1: Train all 5 configurations100python precompute/train_all.py --p 23 --output ./trained_models --resume101 102# Step 2: Generate model-based plots (21 PNGs + 7 JSONs)103python precompute/generate_plots.py --p 23 --input ./trained_models --output ./precomputed_results104 105# Step 3: Generate analytical simulation plots (2 PNGs, no model needed)106python precompute/generate_analytical.py --p 23 --output ./precomputed_results107```108 109### Output110 111Each modulus produces ~33 files in `precomputed_results/p_XXX/`:112 113| Category | Files | Description |114|----------|-------|-------------|115| Overview (Tab 1) | 2 PNGs + 1 JSON | Loss, IPR, phase scatter |116| Fourier Weights (Tab 2) | 3 PNGs + 1 JSON | DFT heatmaps, cosine fits, neuron spectra |117| Phase Analysis (Tab 3) | 3 PNGs | Phase distribution, alignment, magnitudes |118| Output Logits (Tab 4) | 1 PNG + 1 JSON | Logit heatmap, interactive explorer |119| Lottery Mechanism (Tab 5) | 3 PNGs | Magnitude race, phase convergence, contour |120| Grokking (Tab 6) | 5 PNGs + 3 JSONs | Loss/acc curves, memorization, weight evolution |121| Gradient Dynamics (Tab 7) | 4 PNGs | Phase alignment + DFT for Quad and ReLU |122| Decoupled Simulation (Tab 8) | 2 PNGs | Analytical ODE integration |123| Metadata | 2 JSONs | Config + training log |124 125> **Note:** Grokking results (Tab 6) require $p \geq 19$. Smaller values of $p$ have too few data points for a meaningful train/test split.126 127## The 5 Training Configurations128 129| Config | Activation | Optimizer | LR | Weight Decay | Data | Epochs | Used In |130|--------|-----------|-----------|-----|-------------|------|--------|---------|131| `standard` | ReLU | AdamW | 5e-5 | 0 | 100% | 5,000 | Tabs 1–4 |132| `grokking` | ReLU | AdamW | 1e-4 | 2.0 | 75% | 50,000 | Tabs 1, 6 |133| `quad_random` | Quad | AdamW | 5e-5 | 0 | 100% | 5,000 | Tab 5 |134| `quad_single_freq` | Quad | SGD | 0.1 | 0 | 100% | 10,000 | Tab 7 |135| `relu_single_freq` | ReLU | SGD | 0.01 | 0 | 100% | 10,000 | Tab 7 |136 137## Running a Single Experiment138 139For custom experiments outside the pre-computation pipeline:140 141```bash142cd src143 144# Train with default config (p=97, d_mlp=1024, ReLU, 5000 epochs)145python module_nn.py146 147# Train with specific parameters148python module_nn.py --p 23 --d_mlp 512 --num_epochs 5000 --lr 5e-5149 150# Dry run: see config without training151python module_nn.py --dry_run --p 23 --d_mlp 512152```153 154## Notebooks155 156Interactive analysis notebooks in `notebooks/`:157 158| Notebook | Description |159|----------|-------------|160| `empirical_insight_standard.ipynb` | Fourier weight analysis, phase distributions, output logits |161| `empirical_insight_grokk.ipynb` | Grokking stages, weight dynamics, IPR evolution |162| `lottery_mechanism.ipynb` | Neuron specialization, frequency magnitude/phase tracking |163| `interprete_gd_dynamics.ipynb` | Phase alignment under single-frequency initialization |164| `decouple_dynamics_simulation.ipynb` | Analytical gradient flow simulation |165 166## Setup167 168### Requirements169 170- Python 3.8+171- PyTorch 2.0+172- CUDA-capable GPU (recommended for $p > 50$; CPU works for small $p$)173 174### Installation175 176```bash177git clone https://github.com/Y-Agent/modular-addition-feature-learning.git178cd modular-addition-feature-learning179pip install -r requirements.txt180```181 182## Project Structure183 184```185modular-addition-feature-learning/186├── src/                          # Core source code187│   ├── module_nn.py             # Training script with CLI188│   ├── nnTrainer.py             # Training loop and optimization189│   ├── model_base.py            # Neural network architecture (EmbedMLP)190│   ├── mechanism_base.py        # Fourier analysis and decomposition191│   ├── utils.py                 # Configuration and helpers192│   └── configs.yaml             # Default hyperparameters193├── precompute/                   # Batch training and plot generation194│   ├── run_pipeline.sh          # Full pipeline for one modulus195│   ├── run_all.sh               # Batch pipeline for all odd p196│   ├── train_all.py             # Train 5 configurations197│   ├── generate_plots.py        # Generate model-based plots + JSONs198│   ├── generate_analytical.py   # Analytical ODE simulation plots199│   └── prime_config.py          # Configurations and sizing formula200├── hf_app/                       # Gradio web application201│   └── app.py                   # Interactive visualization app202├── precomputed_results/          # Pre-computed plots and data203│   ├── p_015/                   # Results for p=15204│   ├── p_023/                   # Results for p=23205│   ├── p_029/                   # Results for p=29206│   └── p_031/                   # Results for p=31207├── notebooks/                    # Analysis and visualization notebooks208├── requirements.txt              # Python dependencies209└── README.md210```211 212## Citation213 214```bibtex215@article{he2025modular,216  title={On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking},217  author={He, Jianliang and Wang, Leda and Chen, Siyu and Yang, Zhuoran},218  journal={arXiv preprint arXiv:2602.16849},219  year={2025}220}221```222 223## License224 225[MIT License](LICENSE)226