murphysics999/stoch_norm_gpt2_small

GPT-2 small with a stochastic LayerNorm channel (σ_g = 0.5)

GPT-2 small (124M), fine-tuned so that every LayerNorm read is noisy. Each of the 25 LayerNorms (ln_1/ln_2 in all 12 blocks, plus ln_f) normalizes its input, adds isotropic Gaussian noise with a learned per-LayerNorm scale σ, then normalizes again before its gain and bias:

x̂ = normalize(x)
x̂ = normalize(x̂ + σ · ε),   ε ~ N(0, I)
out = x̂ · weight + bias

The σ values were learned during training under a fixed total information budget, so each σ says how much information the network chose to let through that read. The weights are ordinary GPT-2 weights; the noise lives in the LayerNorm forward pass, shipped as custom code.

The noise is part of the model. It is on by default, including in eval() and generate(), and the model was only ever trained with it on.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "murphysics999/stoch_norm_gpt2_small"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).eval()

torch.manual_seed(0)  # the noise draws come from torch's RNG
inputs = tok("The history of the city begins", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=40, do_sample=True, top_p=0.95)
print(tok.decode(out[0], skip_special_tokens=True))

trust_remote_code=True is required: it loads modeling_noisy_gpt2.py (about 80 lines; read it first). Pass it to AutoTokenizer as well, or transformers asks interactively. Tested with transformers 4.57.1 and 5.9.0, which give bit-identical outputs for the same seed on CPU.

Things to know:

• Every forward pass is a random draw. Even greedy decoding varies across seeds; fix torch.manual_seed for reproducibility. To get the model's predictive distribution for a context, average the softmax over several draws (e.g. repeat the context along the batch).
• KV cache. With the cache (the default), each position's noise is drawn once and reused on later steps, which is exactly the joint distribution of a single training forward pass. With use_cache=False, every step redraws the noise for the whole prefix.
• Precision. The LayerNorm arithmetic always runs in float32, whatever dtype the model is loaded in, so fp16/bf16 loading is safe.
• Per-tap noise levels are in model.config.tap_sigmas, keyed by module name, and in policy.json with κ and rate.
• model.set_noise(False) switches to a deterministic pass through the same weights. That is not the trained model: training never ran without noise, and in our experiments the noise-off pass has misrepresented the noisy model's behaviour. Use it only if you specifically need it.

Evaluation

OpenWebText held-out split: the last 5,000 documents, never trained on; 5,666,990 tokens, 512-token blocks, bf16 autocast.

model	perplexity	next-token top-1 acc	KL(GPT-2 ‖ model), nats/token
this model (noise on, mean of 3 draws)	31.85 ± 0.01	0.370	0.270
GPT-2 small, deterministic	25.22	0.399	0

GPT-2 was pretrained on WebText, which OpenWebText recreates, so its row is not strictly held-out for GPT-2 itself. The "±" is the SD across full-split noise draws.

How the noise levels were set

Reads live on the unit sphere of the 767-dimensional mean-zero hyperplane (p = 767). Each noisy read is modelled as a von Mises-Fisher distribution with concentration κ, matched to σ on mean cosine. Its rate is KL(vMF(κ) ‖ uniform): the information, in nats per token position, that the read can carry about its input. For the actual Gaussian-then-renormalize channel this matches exact 1-D quadrature and a Monte Carlo estimate to within about 0.1%.

Training fixes a total budget B = 25 · rate(σ_g): what 25 reads at a uniform σ_g = 0.5 would carry, 15,418 nats (22,243 bits) per token position. A learned softmax over the 25 LayerNorms splits B between them, and σ at each LayerNorm follows from its share. The allocation is trained jointly with the weights.

Learned allocation, in depth order:

LayerNorm	σ	κ	rate (bits)
h.0.ln_1	0.456	4,052	972.9
h.0.ln_2	0.300	8,904	1,379.6
h.1.ln_1	3.097	260	54.9
h.1.ln_2 (floor)	10.101	76	5.4
h.2.ln_1	0.567	2,737	781.0
h.2.ln_2	0.585	2,593	755.5
h.3.ln_1	0.532	3,070	836.2
h.3.ln_2	0.454	4,076	975.9
h.4.ln_1	0.540	2,982	822.0
h.4.ln_2	0.471	3,813	942.6
h.5.ln_1	0.507	3,345	878.0
h.5.ln_2	0.440	4,321	1,005.2
h.6.ln_1	0.623	2,323	704.0
h.6.ln_2	0.420	4,719	1,049.8
h.7.ln_1	0.556	2,834	797.7
h.7.ln_2	0.411	4,903	1,069.3
h.8.ln_1	0.640	2,223	683.7
h.8.ln_2	0.394	5,301	1,109.2
h.9.ln_1	0.620	2,347	708.7
h.9.ln_2	0.392	5,357	1,114.5
h.10.ln_1	0.490	3,549	907.1
h.10.ln_2	0.381	5,646	1,141.6
h.11.ln_1 (floor)	10.101	76	5.4
h.11.ln_2	0.347	6,735	1,233.0
ln_f	0.123	50,638	2,320.6

2 LayerNorm(s) sit at the rate table's floor (σ = 10.1); the network all but switched those reads off. Such reads were charged about 0 by the budget but still carry the floor rate, which is the value shown. Total rate actually carried: 15,425 nats.

Training

• Initialization: pretrained GPT-2 small.
• Objective: distillation only. Minimize KL(frozen GPT-2 ‖ noisy model) over next-token distributions (T = 1.0); no cross-entropy term.
• Data: OpenWebText, GPT-2 tokenizer, an EOS token between documents, packed into 512-token blocks.
• Schedule: 25 epochs × 5000 steps × batch 32 = 2.05B tokens. AdamW, lr 0.0001 with cosine decay to 1e-06 after 0.2 epochs of warmup, weight decay 0.005 (none on biases or LayerNorm parameters), gradient clipping 1.0, bf16 autocast, dropout 0. Allocation logits trained with lr 0.006.
• Noise ramp: the budget was ramped geometrically from σ_g = 0.05 to 0.5 over the first 6 epochs.

Intended use and limitations

This is a research artifact for studying how much information a transformer routes through each read, not an improved language model: it is somewhat worse than GPT-2 by perplexity, and it inherits GPT-2's biases and failure modes, including offensive and false text. English only.

Citation

Tracing distinguishability through transformer processing with stochastic LayerNorm, Kieran Murphy, 2026 [arxiv]