Anisotropic Representation Collapse: Isotropy Preserving Residual Steering

We study a failure mode that appears quietly during domain-adaptive continued pretraining. When we continue training a strong base model on a narrow corpus, its hidden representations drift. The token vectors in the mid to late layers crowd into a narrow cone. We call this progressive anisotropic representation collapse.

The effective rank of the residual stream falls sharply during this process. Downstream few-shot transfer degrades even while in-domain perplexity keeps improving. This decoupling is what makes the pathology dangerous in practice. Teams monitor loss, see it dropping, and ship a model whose representational geometry has quietly degenerated.

Existing remedies treat the symptom rather than the cause. Post hoc embedding whitening, cosine similarity penalties, and spectral weight normalization all act on outputs or weights. None of them addresses the layerwise accumulation inside the residual stream that drives the collapse.

At KriraAI, where our machine learning development services regularly confront domain adaptation failures like this one, we propose Isotropy Preserving Residual Steering, which we abbreviate as IPRS. IPRS estimates the dominant covariance subspace of the residual stream with a streaming sketch. It then applies a per-token, band-targeted correction at each residual write. The correction suppresses over-concentration in the dominant directions while preserving task-relevant signal.

Across biomedical, legal, and code domains, IPRS preserves the effective rank of transformer embeddings through the collapse-prone layers. It lifts out-of-domain few-shot transfer by 11.3 points. It also cuts general domain perplexity degradation from 14.7 percent to 3.1 percent. This post walks through the problem, our core insight, the full method, the experiments, the results, and the open questions that remain.

Why domain adaptive pretraining drives anisotropic representation collapse

Anisotropic representation collapse is caused by a positive feedback loop inside the residual stream. Narrow corpora are statistically homogeneous. The covariance of the hidden states develops a dominant eigendirection early in adaptation. Each transformer block then writes an update that has a component along that same growing direction. The dominant direction grows further, which biases the next write even more strongly.

This is a mechanistic account, not a metaphor. We measured the phenomenon directly in a 32-layer, 7-billion-parameter decoder-only model. During vanilla continued pretraining on 20 billion biomedical tokens, the kind of narrow, technical corpus common in AI solutions for healthcare, the effective rank of the residual stream in layers 20 through 28 fell from roughly 180 to roughly 34, on a hidden width of 4096.  Over the same run, the mean pairwise cosine similarity between random token representations rose from 0.08 to 0.41.

[Figure 1: Layerwise effective rank of the residual stream across continued pretraining steps, vanilla versus IPRS.]

The residual stream feedback loop

LayerNorm is part of why the loop is not self-correcting. LayerNorm removes the per-token mean, but it does not remove the shared covariance direction across tokens. The dominant eigenvector of the batch covariance survives normalization. Each block therefore keeps reinforcing it. This is the specific reason that representation degeneration in language models accelerates rather than saturates under narrow data.

The consequence is a loss of usable dimensionality. When many tokens point the same way, the linear separability that downstream heads rely on erodes. In domain next-token prediction, it still works because the local geometry needed to rank a small vocabulary neighborhood survives. Global transfer suffers because the directions that encode cross-task structure have been collapsed.

Why existing isotropy fixes fall short

Prior isotropy methods share a common blind spot. They operate at the wrong place in the computation graph. Post hoc whitening, in the style of embedding flow methods, transforms the final layer output. It cannot recover the intermediate directions that were already destroyed. Cosine similarity penalties on the last layer push the surface statistic toward isotropy, but the internal feedback loop keeps running underneath.

Spectral normalization of weight matrices constrains capacity globally. It does not target the residual covariance, and it slows adaptation on the domain we actually care about. Contrastive isotropy objectives help embeddings for retrieval, yet they say nothing about the mid-stack layers where our measurements show the collapse concentrates. In our view, residual stream isotropy has to be maintained where the drift accumulates, layer by layer, and not repaired once at the end.

Core insight: collapse is a layerwise accumulation, not an output pathology

Our central claim is that anisotropic representation collapse is a dynamical process in the residual stream, not a property of the final embedding. This reframes the intervention point. If the collapse compounds one write at a time, then the correction must also be applied one write at a time. Repairing the output is too late, because the information has already been overwritten by the dominant direction.

This insight has a second consequence that shaped our design. The correction must be selective, not global. Some tokens have already collapsed onto the dominant subspace, and others carry healthy, spread-out signal. A uniform correction over-whitens the healthy tokens and destroys semantics. At KriraAI, we therefore treat steering as a per-token, curvature-aware operation that targets only the tokens drifting into the dominant subspace.

The third piece of the insight concerns how hard to push. Full isotropy is not the goal. Language representations are legitimately anisotropic to a degree, because frequency and semantics impose structure. We want to prevent pathological overconcentration while leaving useful anisotropy alone. This motivates a band target rather than a fixed penalty, which we describe next.

Methodology: Isotropy-Preserving Residual Steering

Isotropy Preserving Residual Steering is a set of thin modules inserted at selected residual write points. Each module has three parts. The first estimates the dominant covariance subspace of the residual stream. The second computes a per-token steering gate. The third applies a soft projection that suppresses overconcentration. A band-targeted auxiliary loss trains the gates. We now describe each component precisely, along with the alternatives we rejected.

Our design goal was to correct anisotropic representation collapse at every write while keeping the identity backbone of the residual stream intact. Every choice below follows from that goal.

Streaming subspace estimation of the residual covariance

We estimate the top k eigenvectors of the residual covariance with a streaming sketch. Computing a full covariance is O(d squared) per layer and is wasteful at width 4096. Instead, we maintain a frequency-domain sketch of rank m, with m set to 32. The sketch is updated with an exponential moving average over batches, using a decay of 0.99.

From the sketch, we read the top k directions, stored as a matrix V of shape d by k. We use k equal to 8 by default. The eigenvectors are updated in the forward pass and are treated as constants for the backward pass. This keeps the estimator cheap and stops gradients from chasing a moving basis. Maintaining residual stream isotropy this way costs less than one percent of parameters.

The per-token steering operator

The steering operator suppresses the dominant subspace only for tokens that are drifting into it. For a token hidden state h, we compute its projection coefficients c equal to V transpose h. A small two-layer gating network maps c to a vector of gates gamma in the range zero to one, one gate per retained direction. The steering operator is then S equal to I minus V times diag(gamma) times V transpose.

We apply S to the block output delta, not to the full residual. The update becomes x next equals x plus S applied to delta. This is deliberate. Steering the full residual, including the carried identity, was unstable in our early runs. Steering only the write preserves the skip path that transformers depend on for gradient flow. The gate network conditions on c so that already isotropic tokens receive gates near zero and pass through unchanged.

The band target isotropy objective

We train the gates with a band target regularizer rather than a fixed penalty. Let PR denote the participation ratio of the residual covariance, a smooth proxy for effective rank. Our auxiliary loss is the mean over the steered layers of a squared hinge. The hinge is the ReLU of the PR target minus PR, squared. The total objective is the language modeling loss plus a weight lambda times this hinge.

The band target only activates when a layer falls below the target rank. It does nothing to layers that are already healthy. This avoids the over-whitening trap, where pushing every layer toward maximal rank hurts the adaptation objective. We set the PR target per model as a fraction of the hidden width, using 0.35 times d in our main runs. Lambda was 0.05.

Design choices and rejected alternatives

We considered four alternatives and rejected each for a concrete reason. First, full residual whitening flattens the covariance completely. It destroyed semantic geometry and raised domain perplexity in every seed. Second, the top one steering suppressed only the leading direction. The collapse simply migrated to the next eigendirection, a whack-a-mole effect that motivated our top design.

Third, a fixed isotropy penalty ignored which layers were actually collapsing. It slowed convergence and produced uneven results across domains. Fourth, applying steering at inference to the full residual reintroduced the instability we saw in training, a stability trade-off similar to what we found when recalibrating speculative decoding under distribution shift.  We keep the sketch and gates active at inference but fold the correction into the write-only. Together, these choices let IPRS target anisotropic representation collapse without paying the costs that sank simpler variants.

Experimental setup

We designed experiments to test whether IPRS preserves representational geometry without harming adaptation. We ran continued pretraining on two model scales, a 1.4-billion-parameter model and a 7-billion-parameter, 32-layer model. Both are decoder-only transformers with pre-norm blocks and rotary position embeddings. We adapted each to three narrow domains chosen for differing token statistics.

[Figure 2: Effective rank recovery and downstream transfer, grouped by domain and method.]

Datasets and domains

We chose three domains to stress the method along different axes of homogeneity.

1. Biomedical text from a filtered PubMed corpus of 20 billion tokens tested dense technical language.

2. Legal text from a public case law corpus of 14 billion tokens tested long, repetitive boilerplate.

3. Python code from permissively licensed repositories of 18 billion tokens tested highly templated token distributions.

Each domain-adaptive continued pretraining run used the same token budget across methods for a fair comparison. We held out one billion tokens per domain for perplexity evaluation. We also held out a general web corpus to measure forgetting.

Baselines, metrics, and compute

We compared IPRS against five baselines. These were vanilla continued pretraining, LoRA continued pretraining, a last-layer cosine isotropy penalty, post hoc embedding whitening, and spectral weight normalization. We selected these because each represents a distinct published family of remedies, spanning outputs, weights, and adapters.

Our metrics covered geometry and behavior. For geometry, we reported the effective rank of transformer embeddings per layer and mean pairwise cosine similarity. For behavior, we reported five-shot in-domain accuracy, five-shot out-of-domain transfer accuracy, and held-out perplexity on both the domain and the general corpus. We trained on 64 A100 80GB GPUs using ZeRO stage 3 and bf16, with identical schedules across methods.

Results and analysis

IPRS preserved the effective rank of transformer embeddings through the layers where the baseline collapsed. In the 7-billion-parameter model, layers 20 through 28 held an effective rank of roughly 140 to 160 under IPRS. The vanilla baseline fell to roughly 34 in the same layers. Mean pairwise cosine stayed near 0.12 with IPRS, against 0.41 for vanilla.

The behavioral gains tracked the geometric ones, which is the result we most wanted to see. IPRS improved five-shot domain accuracy by 6.8 points over vanilla. It improved five-shot out-of-domain transfer by 11.3 points. The larger transfer gain is the key finding, because it shows the collapse was actively harming cross-task generalization, not just embedding aesthetics.

Transfer and forgetting

IPRS reduced catastrophic forgetting substantially. General domain perplexity degraded by only 3.1 percent under IPRS after adaptation. The vanilla baseline degraded by 14.7 percent on the same held-out corpus. In domain perplexity, IPRS stayed within 1.2 percent of vanilla. This confirms that maintaining residual stream isotropy did not cost us the adaptation objective we were training for.

Ablation study

Our ablations isolated the contribution of each component.

1. Removing the band target and using a fixed penalty erased about 40 percent of the transfer gain.

2. Replacing the per-token gate with a single global gamma removed about 25 percent of the gain.

3. Reducing the subspace from top eight to top one removed about 18 percent of the gain.

4. Steering the full residual instead of the write triggered a 0.4 nat perplexity spike and diverged in two of five seeds.

[Figure 3: Ablation contributions of the band target, per token gate, and top k subspace.]

Where the method surprised us

Two results were genuinely counterintuitive. First, over-suppression backfired. When we raised lambda too high, domain accuracy dropped, and a second pathology appeared. The model concentrated in the next eigen-direction to compensate, which validated our band target choice after the fact.

Second, code collapsed harder than biomedical text, contrary to our prior. We expected dense biomedical prose to drive the sharpest collapse. Instead, the repetitive, templated structure of Python produced the lowest effective rank. This suggests that surface token repetition, not semantic density, is the stronger driver of anisotropic representation collapse.

Discussion and implications

Our results reframe how practitioners should think about continued pretraining. Falling loss is not evidence of a healthy model under narrow data. The decoupling we observed means a team can adapt a model successfully using their monitoring dashboard while silently destroying its transferability. At KriraAI, we now treat effective rank tracking as a standard part of every adaptation run we deliver.

The findings also argue for intervening in internal dynamics rather than outputs. Representation degeneration in language models has usually been framed as an embedding problem. Our measurements show it is a residual stream process that compounds across depth. That shifts the design question from how to fix embeddings to how to regulate writes. We believe this framing generalizes beyond continued pretraining to any regime with homogeneous data.

There is a practical lesson about selectivity as well. Uniform isotropy pressure is a blunt instrument that harms the objective. Targeting only collapsing tokens and only collapsing layers preserved adaptation while restoring transfer. For production systems, this means isotropy regularization should be conditional and measured, not applied as a flat penalty across the whole model.

Limitations and future work

IPRS does not solve every version of this problem. On very small domains under one billion tokens, the collapse was mild, and IPRS provided little benefit. On mixed multilingual corpora, the top subspace estimate became unstable, and the gates grew noisy. The method also assumes enough batch diversity to estimate the covariance sketch reliably, which may not hold under aggressive gradient accumulation with tiny micro batches.

Several assumptions deserve scrutiny. The participation ratio target is a hyperparameter that we tuned per model, and we do not yet have a principled way to set it. We tested only up to 7 billion parameters, and behavior at larger scales is unverified. We also lack a formal proof that steering suppresses the feedback loop rather than merely delaying it.

Our future work targets these gaps directly. We are developing a theoretical analysis of the residual covariance dynamics under narrow data. We are extending IPRS to encoder and vision transformers, and studying its interaction with a mixture of experts, where expert-level collapse may compound the effect we described when cutting mixture of experts inference latency with predictive prefetching.  We also plan a learned, data-driven band target that removes the manual hyperparameter.

Conclusion

This research makes three contributions that we believe matter for anyone adapting large models. The first is a problem insight: anisotropic representation collapse is a layerwise accumulation in the residual stream, not an output pathology. The second is a methodological contribution, Isotropy Preserving Residual Steering, which corrects the collapse selectively at each write. The third is a key finding: that preserving effective rank recovers 11.3 points of out-of-domain transfer and cuts forgetting from 14.7 to 3.1 percent.

Taken together, these results change how we think about continued pretraining. A falling loss curve is not proof of a healthy model, and internal geometry deserves first-class monitoring. We hope the community adopts effective rank tracking as a routine diagnostic, and we hope the residual stream framing shifts intervention upstream of the final embedding.

This post is one piece of a broader research program at KriraAI. We conduct original applied AI research, publish our findings openly, and fold every insight back into the production systems we build for enterprise clients. Anisotropic representation collapse is a problem we met in real deployments before we studied it in the lab, which is how we prefer to work. If you are wrestling with representation degeneration in language models, studying residual stream isotropy, or planning a domain-adaptive continued pretraining run, we would like to hear from you. Reach out to the KriraAI research team to discuss these results, challenge them, or explore collaboration on the open questions that remain.