irregular-parameter-sharing

Experiments for irregular parameter sharing in Transformers.

The current experiment tests hard irregular sharing of exact Transformer MLP hidden-chunk parameter blocks. A standard dense MLP can be written as a sum over hidden-width chunks:

MLP_l(x) = sum_c GELU(x W1[id(l, c)]) W2[id(l, c)]

The unshared baseline gives every (layer, hidden_chunk) position a unique block ID, which is equivalent to an ordinary dense Transformer MLP without biases. The hard-sharing variants reuse exact block IDs across depth and hidden chunks using balanced random layouts, a maximum-depth-distance layout, and a best-of-N random-layout search.

PYTHONPATH=src python scripts/run_hard_block_lm.py \
  --out runs/hard_block_smoke \
  --device cuda \
  --token-path /home/davwis/main/exploration/irregular_token_lm/cache/openwebtext_gpt2_200000.npy \
  --vocab-size 50257 \
  --val-tokens 20000 \
  --search-steps 3 \
  --final-steps 3 \
  --eval-every 3 \
  --dim 128 \
  --depth 2 \
  --chunks 4 \
  --shared-blocks 4 \
  --search-budget 2 \
  --final-variants unshared random max_distance best_random \
  --final-seeds 0,1

The prior butterfly experiment asked whether sharing small butterfly-style parameter blocks across both depth and width could outperform random sharing under the same parameter budget. That code trains decoder-only language models on local OpenWebText GPT-2 tokens and compares:

• an unshared butterfly-block Transformer reference,
• randomly sampled depth-width sharing layouts,
• the best layout from a fixed random search budget,
• and a maximum-depth-distance sharing layout.

The current hard block-sharing task arc is in docs/hard-block-sharing/. Earlier exploratory arcs are in docs/butterfly-sharing/, docs/shared-basis-field/, and docs/shared-basis-sharing/.

Current result

The completed hard block-sharing run uses a 16-layer, width-1024 Transformer LM on GPT-2-tokenized OpenWebText, with 5,000 optimizer steps per final run and three final seeds. The sharing variants use 128 exact MLP hidden-chunk blocks for 256 (layer, chunk) positions, cutting the MLP-bank parameters from 134,217,728 to 67,108,864. All sharing variants in the main comparison are parameter-matched and use balanced two-use hard ties.

Variant	Val CE mean	Val CE sd	Delta vs unshared	Delta vs random
unshared dense MLP	4.5928	0.0384	0.0000	0.0284
sequence cross-layer sharing	4.5723	0.0781	-0.0206	0.0078
cycle cross-layer sharing	4.5322	0.0733	-0.0606	-0.0322
sequence width sharing	4.5621	0.0366	-0.0307	-0.0023
cycle width sharing	4.5043	0.0537	-0.0885	-0.0601
regular diagonal depth-width sharing	4.5504	0.0086	-0.0424	-0.0140
balanced random hard sharing	4.5645	0.0601	-0.0284	0.0000
maximum-depth-distance hard sharing	4.5279	0.0243	-0.0650	-0.0366
best-of-12 random hard sharing	4.6057	0.0523	0.0129	0.0413

The corrected comparison changes the interpretation: hard sharing still outperforms the unshared dense MLP at this scale, but the best matched baseline is regular cycle width sharing, not the maximum-depth-distance irregular layout. Fully depth-tied chunks are excluded from the main comparison because they use only 16 blocks instead of the 128-block sharing budget.

The combined artifact is docs/hard-block-sharing/artifacts/combined_summary.md.

Quick smoke test

python scripts/run_butterfly_lm.py \
  --out runs/smoke \
  --max-tokens 200000 \
  --val-tokens 20000 \
  --steps 20 \
  --search-budget 2 \
  --final-variants unshared random max_distance \
  --dim 256 \
  --depth 4 \
  --groups 8 \
  --batch-size 8 \
  --eval-batch-size 4

Data

The runner expects OpenWebText parquet shards at /home/davwis/main/data/openwebtext/plain_text/ by default.