lylybig/resnet-moe

MOE-ResNet — Möbius Interaction Analysis of Trained ResNets

Empirical study of how a trained ResNet allocates functional contribution across its residual branches, using the exact Möbius interaction decomposition of the residual-gate function. Each experiment lives in its own folder with a fixed-seed pipeline, written outputs, a per-folder REPORT.md analysis, and figures.

The codebase wraps every BasicBlock of a torchvision ResNet-18 in a GatedBasicBlock that exposes a per-sample binary mask on the residual branch:

$$x_{l+1}(m)=\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}\text{\hspace{0.17em}⁣}(x_l(m)+m_l\cdot F_l(x_l(m)))\mathrm{.}x\_\{l+1\}(m)\; =\; \backslash mathrm\{ReLU\}\backslash !\backslash bigl(x\_l(m)\; +\; m\_l\; \backslash cdot\; F\_l(x\_l(m))\backslash bigr).$$

With $L = 8$ residual branches and $2^L = 256$ binary masks, every experiment reduces to: enumerate masks → collect logits $h_x(\mathbf 1_S) \in \mathbb R^{1000}$ → compute Möbius coefficients $\Delta_S h_x = \sum_{T \subseteq S}(-1)^{|S|-|T|},h_x(\mathbf 1_T)$ via a length-$L$ butterfly → derive an experiment-specific summary statistic over the 256 coefficients.

Five experiments

#	Question	Result
1	Does naive subset summation reconstruct the full ResNet output?	Naive sums fail by ~80×; Möbius is exact to float32 epsilon.
2	How is Möbius energy distributed across interaction orders?	Peak at $k = 5$, $\kappa \approx 5$ — no low-order bias at $\lambda = 1$.
3A/3B	Does $\kappa$ migrate to low orders under residual scaling?	Yes — $\kappa$ falls 5.00 → 1.14 monotonically as $\lambda \to 0.1$.
3 (top-$K$)	Is the decomposition sparse at the level of individual subsets?	Magnitude top-128 hits 88% top-1 vs <1% for any control.
4	Are expert sets input-dependent and class-aligned?	Yes (scalar ranking only): NN-acc 24.2% on 32-bit signature, 240× chance.
5	Does difficulty modulate expert usage?	Mostly null. Hard samples need ~1–3 more experts but no order shift.

Each result lives at experiment_{1..5}/results/REPORT.md. A cross-experiment numeric index is at tables.md.

Repository layout

MOE-ResNet/
├── README.md                   ← this file
├── tables.md                   ← consolidated results across experiments
├── environment.yml             ← conda env spec (kept from the proposal)
├── requirements.txt            ← pip alternative
├── experiment_plan/            ← one .md per experiment with the formal plan
├── experiment_1/               ← naive subset summation vs Möbius reconstruction
├── experiment_2/               ← Möbius interaction order spectrum
│   └── results_lambda/         ← Experiment 3A/3B residual-scaling sweep
├── experiment_3/               ← top-K sparse residual-expert reconstruction
├── experiment_4/               ← input-dependent residual expert sets
├── experiment_5/               ← sample difficulty vs interaction complexity
├── imagenet-1k/data/           ← HF parquet shards (validation + test splits)
├── checkpoints/                ← torchvision ResNet-18 weights (auto-cached)
├── figures/                    ← shared figure outputs (rare)
├── paper/                      ← LaTeX sources for the writeup
└── src/                        ← legacy proposal scaffolding (predates these experiments)

The experiments share three core primitives, all defined inside experiment_1/ and re-imported by later experiments:

module	what it provides
experiment_1/gated_resnet18.py	GatedBasicBlock, set_per_sample_masks, build_gated_resnet18.
experiment_1/mobius.py	mask enumeration, fast Möbius transform, full / empty mask indices.
experiment_1/core.py	evaluate_all_masks: replicates input across the mask dim and forwards.
experiment_1/data.py	parquet-backed ImageNet dataset reading the HF mirror.
experiment_2/spectrum.py	scalar Möbius (predicted-class), order-energy reductions.
experiment_3/topk.py	magnitude rankings, cumulative reconstruction, baselines, $K_\text{eff}$.
experiment_4/signatures.py	top-$K$ binary / weighted signatures, Jaccard, NN-classification.
experiment_5/difficulty.py	five difficulty scores from full-mask logits.
experiment_5/complexity.py	six complexity measures + $K_\tau^\text{err}$, $K_\eta^\text{mass}$.

Sampling and reproducibility

Two committed sample lists pin every image used downstream:

• experiment_1/sample_index.json — 10000 test images (uniform, unlabeled, seed 0). Used by Experiments 1, 2, 3 (top-$K$), and 3A/3B.
• experiment_1/sample_index_val.json — 10000 validation images (per-class, 10/cls × 1000, seed 0, with ground-truth labels). Used by Experiments 4 and 5 (which need labels for same/different-class comparisons, loss, correctness).

Determinism cross-check: full-mask top-1 on sample_index_val.json is 69.94%, matching torchvision's reported 69.76% for IMAGENET1K_V1. Any divergence indicates a software / weight regression.

Every experiment script seeds Python random, numpy, PyTorch CPU + CUDA RNGs from --seed, and sets torch.backends.cudnn.deterministic = True, torch.backends.cudnn.benchmark = False. ResNet-18 stays in eval() mode so BN uses running statistics.

Environment setup

Option A: conda (recommended)

conda env create -f environment.yml
conda activate moe-resnet
pip install pyarrow             # parquet reader (used in step_2 / step_3)

If you do not have CUDA 11.8, edit environment.yml and replace the pytorch-cuda=11.8 line with the channel/spec that matches your driver (see https://pytorch.org/get-started/locally/).

Option B: pip + venv

python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt
pip install pyarrow

A CUDA-enabled GPU is required. Each experiment fits comfortably on a single H100 80 GB at B=8, mask_chunk=256 (~13 GB peak, ~3 min wall on N = 10000); on 8–16 GB GPUs reduce to B=2, mask_chunk=128 (see each README.md for the memory-tight invocation).

Dataset

ImageNet-1k via the Hugging Face mirror ILSVRC/imagenet-1k. One-time download:

bash experiment_1/step_1.sh

This pulls 14 validation-*.parquet (50k labeled images) and 28 test-*.parquet (100k unlabeled images) into imagenet-1k/data/ (~19 GB total). Pretrained ResNet-18 weights are fetched automatically by torchvision.

Running the experiments

Each experiment is a self-contained pipeline (step_1, step_2, …) with its own README.md describing every flag. Headline commands:

Experiment 1 — naive subset summation fails

python -m experiment_1.step_2_sample \
    --pattern 'test-*.parquet' \
    --mode uniform --num-images 10000 --seed 0 \
    --output experiment_1/sample_index.json

python -m experiment_1.step_3_evaluate \
    --sample-index experiment_1/sample_index.json \
    --batch-size 8 --mask-chunk 256 --num-workers 8 --seed 0 \
    --output-dir experiment_1/results

python -m experiment_1.step_4_plot \
    --summary experiment_1/results/summary.json

Experiment 2 — interaction order spectrum

python -m experiment_2.step_1_evaluate \
    --sample-index experiment_1/sample_index.json \
    --batch-size 8 --mask-chunk 256 --num-workers 8 --seed 0 \
    --output-dir experiment_2/results

python -m experiment_2.step_2_plot \
    --summary experiment_2/results/summary.json

Experiment 3A/3B — residual scaling sweep

python -m experiment_2.step_3_lambda_sweep \
    --sample-index experiment_1/sample_index.json \
    --batch-size 8 --mask-chunk 256 --num-workers 8 --seed 0 \
    --output-dir experiment_2/results_lambda

python -m experiment_2.step_4_lambda_plot \
    --summary experiment_2/results_lambda/summary_all.json

Experiment 3 — top-$K$ sparse reconstruction

python -m experiment_3.step_1_evaluate \
    --sample-index experiment_1/sample_index.json \
    --batch-size 8 --mask-chunk 256 --num-workers 8 --seed 0 \
    --num-random-seeds 5 \
    --output-dir experiment_3/results

python -m experiment_3.step_2_plot \
    --summary experiment_3/results/summary.json

Experiment 4 — input-dependent expert sets

python -m experiment_1.step_2_sample \
    --pattern 'validation-*.parquet' \
    --mode per-class --per-class 10 --seed 0 \
    --output experiment_1/sample_index_val.json   # one-time

python -m experiment_4.step_1_signatures \
    --sample-index experiment_1/sample_index_val.json \
    --batch-size 8 --mask-chunk 256 --num-workers 8 --seed 0 \
    --output-dir experiment_4/results

python -m experiment_4.step_2_overlap \
    --signatures experiment_4/results/signatures.pt \
    --output-dir experiment_4/results --max-pairs 45000 --nn-K 32

python -m experiment_4.step_3_plot \
    --summary experiment_4/results/overlap.json \
    --signatures experiment_4/results/signatures.pt

Experiment 5 — difficulty vs complexity

python -m experiment_5.step_1_evaluate \
    --sample-index experiment_1/sample_index_val.json \
    --batch-size 8 --mask-chunk 256 --num-workers 8 --seed 0 \
    --output-dir experiment_5/results

python -m experiment_5.step_2_analyze \
    --per-image experiment_5/results/per_image.pt

python -m experiment_5.step_3_plot \
    --analysis experiment_5/results/analysis.json \
    --per-image experiment_5/results/per_image.pt

Wall-time budget

Single H100 80 GB, settings as above:

experiment	wall	gates
Exp 1 (256-mask sweep)	~3 min	one pass
Exp 2 (spectrum)	~3 min	one pass
Exp 3A/3B (λ sweep)	~21 min	9 lambdas × one pass
Exp 3 (top-K)	~3 min	one pass + 5 random seeds
Exp 4 (signatures)	~3 min	one pass
Exp 5 (per-image)	~3 min	one pass

Steps 2/3 of each experiment are CPU-only and take seconds.

Headline cross-experiment story

1. Möbius is the right algebraic unit (Exp 1: naive 80×, Möbius 1e-6).
2. Energy concentrates at mid-orders $k \approx 5$ (Exp 2: $\kappa^v = 5.0$, $C_{\le 3} = 14%$).
3. Residual scaling gives the predicted low-order migration (Exp 3A/3B: $\kappa$ 5.0 → 1.14 monotonically; narrow regime $\lambda \in [0.8, 1]$ already shows $\kappa$ drop while top-1 agreement is still 73–100%).
4. Within those mid-orders, energy is sparse at the subset level (Exp 3 top-$K$: $\Gamma_{64} = 94%$, magnitude top-128 hits 88% top-1, ~130× over random / order-matched).
5. Which subsets dominate is class-aligned, but only in the predicted-class projection (Exp 4: vector ranking is near-global, scalar ranking gets 24% NN-acc on 32 bits = 240× chance).
6. Difficulty modulates how broadly the network spreads mass, not which orders it uses (Exp 5: $\bar k$, $\kappa$, $C_{\le 3}$ invariant; $N_\text{eff}^v$ +1.4, $K_\tau^\text{err}$ +2.8 on hard quartile).

This composes into a single picture: a trained ResNet behaves like an implicit MoE with a near-fixed mid-order expert pool, class-aligned input-dependence in the predicted-class projection, and a small difficulty-driven broadening of expert support.

License & citation

See paper/ for the LaTeX writeup and citation details. The Möbius machinery, residual-gate setup, and reproducibility pipeline in experiment_{1..5}/ are released under the same terms as the paper.