Source-linked AI summary
Tracing distinguishability through transformer processing with stochastic LayerNorm
Kieran Murphy
TL;DR
Point-valued representation geometry is not intrinsically tied to downstream function, motivating a stochastic treatment of representations. The paper adds noisy, rate-limited residual-stream reads to pretrained transformers and uses information-theoretic distinguishability to trace preserved input distinctions. Experiments on ViT-S and GPT2-small show depthwise perturbation propagation and head-specific sensitivity, while the framework remains a proof of concept with approximation and scale-transfer limits.
Problem
Point-valued representation distances are not intrinsically tied to downstream behavior, creating a gap between geometric similarity and functional distinguishability.
Method
The paper adds stochastic LayerNorm-based residual-stream reads, jointly fine-tunes transformer weights under a shared rate budget, and traces distinguishability using information-theoretic measures.
Results
Experiments on ViT-S and GPT2-small reveal depthwise propagation of visual and linguistic perturbation distinguishability and head-specific sensitivity in query, key, and value projections.
Takeaways & Limitations
The framework provides a functionally grounded lens for analyzing which input distinctions remain available through transformer reads, MLPs, and attention projections.
Takeaways & Limitations
The experiments are a proof of concept, and conclusions may not transfer to unconstrained larger models; effective posterior mixtures also remain difficult to characterize.
Abstract
from arXiv · showhide
Representational similarity is foundational to analyses of deep networks, yet distances between point-valued representations are not intrinsically tied to downstream function: nearby states may produce different behaviors, while distant states may behave similarly. We instead give representations volume, turning similarity into statistical distinguishability. Overlapping stochastic representations necessarily induce overlapping downstream distributions, grounding latent comparison in model function and bringing it under information-theoretic tools such as the data-processing inequality. We realize this idea in pretrained transformers through a light-touch modification to LayerNorm: at each residual-stream read, we normalize the state, add isotropic Gaussian noise, and renormalize. During distillation fine-tuning, one learned allocation parameter per residual-stream read distributes a fixed global rate budget across the processing stack. The resulting model can be viewed as transformer blocks reading the residual stream with learned finite precision under a shared global rate budget. Using the Bhattacharyya coefficient, we trace which counterfactual distinctions are preserved through MLP blocks or selectively exposed to the query, key, and value computations of individual attention heads. Experiments on ViT-S and GPT-2 small reveal the depthwise propagation of continuous visual perturbations and head-specific sensitivity to token distinctions aligned with known attention motifs. These results establish distinguishability as a functionally grounded lens on transformer computation that complements existing interpretability approaches.
1 INTRODUCTION
The paper reframes representation comparison around stochastic volume and functional distinguishability, then implements this perspective in transformers with noisy, rate-limited residual-stream reads. Experiments use these models to trace information preservation across depth and attention-head projections.
- Point-valued representations lack a general functional grounding for geometric similarity because nearby states can support different downstream behavior.
- Stochastic representations make behavior over neighborhoods part of training and create explicit channels for measuring information transmission and pairwise distinguishability.
- Noise is added inside every LayerNorm with one optimized magnitude per location, so attention and MLP blocks pay for residual-stream read precision.
- The method starts from a pretrained deterministic solution and jointly adapts model weights and stochastic channels so noisy neighborhoods become part of computation.
- Experiments on ViT-S and GPT2-small measure augmentation distinguishability across depths, perturbation propagation, and the sensitivity of individual attention heads’ query, key, and value projections.
2 METHOD
The method turns transformer residual-stream reads into rate-limited stochastic channels and uses information-theoretic overlap to trace which distinctions survive processing. It applies noisy LayerNorm at each read, allocates a fixed global rate budget, and evaluates distinguishability through MLPs and attention projections.
- Stochastic residual reads: Each pre-norm attention or MLP sublayer reads its residual stream through a stochastic LayerNorm channel.The state is normalized, perturbed with isotropic Gaussian noise, renormalized, and then passed through the learned affine transformation.
- Information limits: The channel imposes a local information limit: downstream computation cannot contain more information about the raw state than the stochastic read provides.Because the normalized state is a deterministic function of the raw state, the local mutual information reduces to I(Z; V), and the data-processing inequality bounds later computation through V.
- Rate allocation: A fixed global rate bound is distributed across taps with learned softmax allocations, and each allocated rate is converted to a noise scale through a calibrated lookup table.The total rate constraint is enforced directly rather than through an added rate-loss term.
- Distillation and interpretation: The noisy model is inserted into a pretrained transformer and fine-tuned by distillation, making stochastic neighborhoods part of the learned computation.The functionally relevant object at each tap becomes the conditional distribution observed by the downstream sublayer, called the local posterior.
- Distinguishability metric: Bhattacharyya overlap quantifies residual-state distinguishability, with lower overlap indicating greater distinguishability under the stochastic read.The local posterior is approximated with a von Mises–Fisher distribution, whose concentration is calibrated to the Gaussian noise scale.
- Tracing through blocks and heads: The data-processing inequality guarantees that MLPs cannot increase distinguishability, while attention outputs are bounded by the joint distribution of all visible reads.Linear query, key, and value projections separately reveal which residual-stream distinctions remain available to each head component; noninvertible projections can only discard distinctions.
- Common random numbers: For depthwise comparisons, common random numbers match upstream noise across inputs, and the parameter k controls how many draws are shared before remaining noise is marginalized separately.The fully coupled estimate is a lower bound on the Bhattacharyya coefficient between effective posteriors.
3 EXPERIMENTS
Experiments use distilled stochastic ViT-S and GPT-2 models to trace how noise, perturbation distinguishability, and information allocation evolve through depth and attention heads. Results show modality- and depth-dependent stochastic structure, augmentation-specific visual resolution, and head-specific preservation of token distinctions.
- Experimental setup: Distillation used deterministic ViT-S and GPT-2 models with stochastic reads, ramping noise during fine-tuning on ImageNet-1k or OpenWebText.The student was trained solely with a distillation loss, with noise increased during the first 6 of 25 epochs.
- Rate allocation and compounded stochasticity: B ∼10^4 nats marked a behavioral knee: above it, models stayed close to their base behavior, while below it, distillation loss and validation performance degraded.Learned rate allocations were relatively flat across depth, generally assigning lower resolution to attention than MLP taps and protecting the final readout.
- Rate allocation and compounded stochasticity: ViT patch streams accumulated more upstream-induced variance than the CLS stream, while the final tap showed a large inherited-variance fraction because its local read was unusually precise.Variance was decomposed into locally introduced and upstream-induced components at MLP taps.
- Rate allocation and compounded stochasticity: Scrambling cross-stream residual structure reduced ViT distillation loss increasingly at later depths, whereas GPT-2 showed small early degradations and mildly beneficial later scrambling.The effect was strongly model- and depth-dependent, consistent with differences between redundant visual measurements and nested language-model dependencies.
- Vision resolution: For ViT-S, increasing blur magnitude smoothly decreased the Bhattacharyya coefficient, with distinguishability varying substantially across depth and augmentation type.Across 200 images, blurring was the most sensitive augmentation and fading the least sensitive at both an intermediate MLP and the final readout.
- Natural-language attention selectivity: In GPT-2, single-token perturbations produced depthwise-growing Bhattacharyya distance, while individual layer-5 heads selectively exposed the distinction through their query and key projections.The larger-budget model distinguished the pair earlier, and head-specific projection behavior appeared immediately after the perturbed position.
- Natural-language attention selectivity: Seven middle-layer heads preserved more key-projection distinguishability than the random-projection baseline, and six had independent associations with induction.For highlighted heads other than L4H7, selectivity was specific to keys; L4H7 retained unusually high distinguishability across queries, keys, and values.
4 DISCUSSION
The framework gives distinguishability an operational, functionally grounded meaning by processing overlapping stochastic residual-stream reads. Its conclusions remain deliberately one-sided and rely on approximations, while experiments are presented as proof of concept requiring broader validation.
- Overlapping stochastic representations constrain the distinctions transmitted through a read, while low overlap only indicates that a distinction is available.An attention block may still access the distinction through other streams.
- The method does not imply that unmodified transformer point representations possess an intrinsic probabilistic geometry.It instead creates a controlled finite-precision extension of a pretrained model.
- The behaviorally anchored regime can remain within hundredths of a nat of original behavior while making every residual-stream read stochastic and rate-limited.Its lower-rate boundary provides a probe of the original model’s operating resolution.
- The information-rate expression is an upper bound, and effective posteriors induced by upstream stochasticity are generally mixtures that may be substantially broader than local vMF posteriors.Common random numbers reduce variance but do not characterize the full mixture.
- Experiments on ViT-S and GPT-2 small are a proof of concept, with future work needed at larger scales and for joint distinctions available to attention blocks.The authors also identify transfer to unconstrained counterparts as an open question.
A IMPLEMENTATION
The implementation centers on a simple LayerNorm modification, with project code publicly available online.
- The project code is available at murphyka.github.io/stoch layernorm.
- The paper identifies a simple modification to LayerNorm as its implementation crux.
- The LayerNorm modification is presented in Algorithm 1, apart from the analysis.
A.1 THE CHANNEL AND ITS CALIBRATION
The channel models normalized stochastic reads on a sphere using a von Mises–Fisher approximation, with concentration controlling rate and a fixed budget distributed across taps.
- Each LayerNorm read is modeled as vMF(µ, κ) on the sphere induced by zero mean and unit RMS normalization.
- rate(κ) is defined as the KL divergence from vMF(µ, κ) to the uniform distribution.
- The Gaussian noise scale is calibrated through the mean resultant length ρ = A_p(κ), matching the channel’s mean read-cosine to ρ.This calibration supports a reparameterization-clean channel with gradients to weights and rate allocation.
- A single global budget B is distributed across taps through softmax allocation logits, with allocated rates converted to noise scales using a precomputed monotonic lookup.
- Rate conservation avoids the degenerate solution of compressing every tap and paying a penalty once; the distillation KL alone drives allocation.Both architectures contain 25 taps.
A.3 THE NOISE RAMP
Training gradually increases the stochasticity budget from a near-clean initialization to its target during the first six of twenty-five epochs.
- The noise budget starts at σg = 0.05 and reaches its target over the first 6 of 25 epochs, approximately 24% of training.The update occurs every optimizer step rather than once per epoch.
- Two ramp shapes are implemented: linear in log-rate and linear in log σg.GPT-2 uses the log-rate schedule, whereas ViT uses the log-σg schedule.
A.4 GPT-2
GPT-2-small is fine-tuned on OpenWebText with explicit document boundaries and non-overlapping 512-token blocks.
- GPT-2-small has 124M parameters and residual dimension D = 768.
- Documents receive an explicit EOS before packing, preserving GPT-2’s context-reset convention at document boundaries.
- Training uses non-overlapping 512-token windows, eliminating the need for padding or attention masks.
A.5 VIT
ViT-S is fine-tuned on full ImageNet-1k at 224^2 resolution using a fixed augmentation configuration across the sweep.
- The vision runs fine-tune timm’s vit small patch16 224 with residual dimension D = 384.
- Training uses full ImageNet-1k at 224^2 resolution.
- The augmentation configuration is held fixed across σg and model comparisons to avoid confounding from augmentation changes.
B GAUSSIAN TO VON MISES-FISHER APPROXIMATION
The noisy LayerNorm channel is calibrated against a von Mises–Fisher approximation by comparing angular-deviation distributions under rotational symmetry.
- The comparison reduces to the scalar angular deviation from the distribution mean, µ, because of rotational symmetry.
- Histograms of angular deviations are compared with analytic vMF densities across several noise levels.
- GPT-2 and ViT-S dimensionalities are shown, with differences in angular-deviation means and spreads measured across noise levels.
C PUSHFORWARD DENSITY UNDER QUERY, KEY, AND VALUE PROJECTIONS
The paper derives projected densities for stochastic LayerNorm reads under query, key, and value projections, enabling Bhattacharyya-coefficient evaluation after projection.
- The noisy LayerNorm read is approximated by a vMF distribution on a unit-radius sphere of effective dimension p = d − 1.
- A query, key, or value projection, including LayerNorm affine parameters, maps the stochastic read into a projected state.
- After restricting to the image of a rank-deficient projection, the transformed coordinates have an invertible scale and translation that preserve the Bhattacharyya coefficient.
- With orthonormal projection rows and unit-norm spherical samples, the projected support is the unit ball, whose density is derived by integrating over spherical preimages.
- The projected Bhattacharyya coefficient is estimated by importance sampling from an equally weighted proposal formed from the two local vMF posteriors.
- The estimator is unbiased under exact sampling, and the arithmetic–geometric mean inequality bounds every Monte Carlo summand.