This codebase tests whether P_SGD(f|S) ≈ P_B(f|S) holds for Transformers — i.e., whether SGD-trained transformers approximate the Bayesian posterior over functions, as Mingard et al. (2021) showed for FCNs, CNNs, and LSTMs.
Mingard et al. showed that the distribution of functions found by SGD tracks the Bayesian posterior:
P_SGD(f|S) ≈ P_B(f|S)
The key insight is that this agreement is driven not by SGD itself, but by the prior P(f) of randomly initialised networks — which strongly favours low-complexity (low-sensitivity) functions. Their methodology uses the NNGP (Neural Network Gaussian Process) kernel for exact GP inference to compute P_B.
Why transformers are harder: Standard NNGP doesn't apply because softmax attention is non-Gaussian. Hron et al. (2020) derived a closed-form NNGP for linear attention; softmax attention requires Monte Carlo estimation.
The gap this project fills: No one has yet directly compared P_SGD(f|S) to P_B(f|S) for transformers.
We use the 7-bit Boolean function testbed: inputs are {0,1}^7 (128 total), and a function f: {0,1}^7 → {0,1} is represented as the 128-bit string of its outputs. This is tractable because:
- The full input space can be enumerated
- Functions are hashed as integers for O(1) comparison
- Training sets of size |S| = 64 leave |E| = 64 evaluation points
The experiment measures correlation between empirical histograms of which functions SGD converges to (P_SGD) and which the Bayesian posterior predicts (P_B), using Spearman rank correlation as the primary metric.
| Leg | Method | Architecture | Tool |
|---|---|---|---|
| 1 | Closed-form Hron NNGP kernel | Linear attention | neural-tangents |
| 2 | Monte Carlo NNGP | Softmax attention | nt.monte_carlo_kernel_fn |
| 3 | SGD ensemble | Softmax attention | PyTorch |
Legs 1 and 2 both produce P_B; Leg 3 produces P_SGD. Comparing Legs 1 and 2 also tells us whether linear attention is a good proxy for softmax attention in this regime.
mingard-transformers/
│
├── src/mingard_transformers/
│ ├── boolean/
│ │ ├── enumeration.py # Input space, function hashing, train/eval splits
│ │ └── complexity.py # Boolean sensitivity, LZ complexity
│ │
│ ├── architectures/
│ │ ├── nt_fcn.py # neural-tangents FCN (validation baseline)
│ │ ├── nt_transformer.py # neural-tangents transformer (Legs 1 + 2)
│ │ └── pt_transformer.py # PyTorch transformer (Leg 3)
│ │
│ ├── bayesian/
│ │ ├── gp_inference.py # MSE-path GP posterior sampling (exact, no EP)
│ │ ├── hron_kernel.py # Leg 1: Hron closed-form kernel wrapper
│ │ └── mc_kernel.py # Leg 2: Monte Carlo kernel wrapper
│ │
│ ├── sgd_ensemble/
│ │ ├── trainer.py # Single-network train-to-zero-error loop
│ │ └── ensemble.py # Parallel ensemble with checkpointing
│ │
│ ├── metrics/
│ │ ├── correlations.py # Spearman ρ with bootstrap CI, top-k overlap
│ │ └── divergences.py # KL, TV, JS divergences
│ │
│ └── visualization/
│ ├── scatter.py # log P_B vs log P_SGD scatter (Figure 1a style)
│ └── simplicity_bias.py # P(f) vs sensitivity (Figure 1b style)
│
├── experiments/
│ ├── 00_validate_fcn.py # FCN validation gate — run this first
│ ├── 01_leg1_linear_attn.py # Compute P_B via Hron kernel
│ ├── 02_leg2_mc_softmax.py # Compute P_B via MC NNGP
│ ├── 03_leg3_sgd_ensemble.py # Compute P_SGD via ensemble
│ └── 04_compare_all.py # Final figures and metrics
│
├── slurm/
│ ├── run_kernels.sh # SLURM: Legs 1 + 2 (CPU)
│ └── run_ensemble.sh # SLURM: Leg 3 (GPU, array job)
│
├── tests/
│ └── test_boolean.py # 23 unit tests (all passing)
│
├── environment.yml # Pinned dependencies for HPC
└── results/ # Output directory (gitignored)
The key constraint: neural-tangents 0.6.5 requires JAX 0.4.x. This is incompatible with JAX 0.9+ if already installed elsewhere. Create an isolated conda environment on the HPC:
conda env create -f environment.yml
conda activate mingard-transformers
# GPU JAX (replace with your CUDA version):
pip install "jax[cuda12_pip]==0.4.26" \
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
# GPU PyTorch:
pip install torch==2.2.1 --index-url https://download.pytorch.org/whl/cu121
pip install -e .For local CPU-only development (no neural-tangents needed for Leg 3 testing):
pip install -e . # uses whatever JAX/torch is already installed
python -m pytest tests/ # should all passpython experiments/00_validate_fcn.py --M 1000Expected: Spearman ρ > 0.7 between FCN P_B and FCN P_SGD. This reproduces Mingard et al.'s qualitative result and confirms the pipeline is correct.
sbatch slurm/run_kernels.sh
# or locally:
python experiments/01_leg1_linear_attn.py --seed 0
python experiments/02_leg2_mc_softmax.py --seed 0 --check_convergenceResults saved to results/leg1_*/ and results/leg2_*/.
sbatch slurm/run_ensemble.sh # 10-task array, 500 networks each = 5000 total
# or locally (small run):
python experiments/03_leg3_sgd_ensemble.py --M 200 --leg1_results results/leg1_L2H4d16_seed0Note: Transformers require Adam (SGD does not converge reliably). The default is --optimizer adam --lr 1e-3. This is a valid choice — Mingard et al. study optimizer effects as a secondary phenomenon.
python experiments/04_compare_all.py \
--leg1_dir results/leg1_L2H4d16_seed0 \
--leg2_dir results/leg2_L2H4d16_mc500_seed0 \
--leg3_dir results/leg3_transformer_L2H4d32_adam_M5000_seed0 \
--output_dir figures/Produces:
fig1_scatter_pb_vs_psgd.pdf— log P_B vs log P_SGD, colored by sensitivityfig2_simplicity_bias.pdf— probability mass vs Boolean sensitivity for all three distributionsfig3_prior_vs_sensitivity.pdf— prior P(f) vs sensitivity (Bhattamishra-style simplicity bias)
MSE loss, labels {-1, +1} — gives an exact closed-form GP posterior (no expectation propagation approximation needed). Both the neural-tangents GP inference and the PyTorch ensemble use this.
Function hashing — functions are packed as integers via np.packbits, giving O(1) comparison and storage. A function on 64 evaluation points is a 64-bit integer.
Stopping criterion — training stops at the first epoch achieving 100% training accuracy, exactly matching Mingard et al.'s protocol. Non-converged runs are discarded.
Input format for transformers — the 7 Boolean bits are treated as a length-7 sequence with 1 channel: (N, 7, 1). Each bit is a token, which is the principled interpretation for attention.
Validation gate — the FCN experiment must pass (ρ > 0.7) before transformer experiments are trusted. This catches pipeline bugs early.
If the P_SGD ≈ P_B result holds for transformers, you should see:
- Scatter plot (Fig 1a): points tightly clustered along the diagonal in log-log space
- Simplicity bias (Fig 2): P_B, P_SGD, and the prior P(f) all peaked at low Boolean sensitivity
- Spearman ρ > 0.7 between P_B (Hron kernel) and P_SGD
If the result doesn't hold cleanly, the most informative diagnostic is the Hron vs MC comparison (Leg 1 vs Leg 2): if those disagree (ρ < 0.8), the linear attention prior is a poor proxy for the actual softmax transformer.
- Mingard et al. (2021). Is SGD a Bayesian sampler? Well, almost. JMLR.
- Hron et al. (2020). Infinite attention: NNGP and NTK for deep attention networks. ICML. arXiv:2006.10541
- Bhattamishra et al. (2023). Simplicity Bias in Transformers. ACL.
- Yang & Hu (2021). Feature Learning in Infinite-Width Neural Networks. ICML. (Tensor Programs IV — explains why NNGP results apply to lazy/NTK regime only, not pretrained LLMs.)