Shared Functional Bases Across Transformer Layers

Motivation

The question was whether there is a mathematical shortcut for cutting the memory footprint of language models. Transformer layers perform structurally similar operations; if they also shared parametric structure, it would suffice to store a common basis once and encode each layer’s deviation cheaply. The extreme form of this idea is a Rubik’s-cube analogy: one canonical matrix plus a short sequence of cheap transformations (permutation, scaling, rotation) per layer.

We tested this on the Q projections of GPT-2 small (12 layers, d = 768), later extending to all matrix types and to GPT-2 medium, on a single RTX 3090.

A chain of nine tests

Each test was designed to cheaply falsify or narrow the previous hypothesis, always against a null control (a random matrix or random subspace subjected to the identical procedure) and with pre-registered decision thresholds.

1. Whole-matrix alignment. Each Q matrix is aligned to a reference layer through progressively stronger transformations: row/column permutation (Hungarian algorithm), diagonal scaling, one-sided orthogonal rotation (Procrustes), two-sided rotation. The metric is residual relative energy.
2. Shared weight subspaces. Top-k singular subspaces per layer (k = d/4); we measure energy captured by the layer’s own basis (ceiling), by a shared energy-weighted basis (candidate), and by a random basis (floor). Score is (shared − null)/(own − null).
3. Functional space. The same test, but on matrices M_l = W_l C_l^{1/2}, where C_l is the input-activation covariance of the layer estimated from 32k calibration tokens — directions weighted by how real inputs actually excite them.
4. End-to-end reconstruction. Actual compression: shared whitening (average covariance), shared basis U, V, a per-layer k × k core, reconstruction W = U A Vᵀ S⁺, inserted into the model, evaluated by perplexity on WikiText-2 against a per-layer whitened SVD baseline.
5. Fine-tuning cores. Q projections replaced by a factorization with frozen bases; only the cores A_l (442k parameters) are trained with the ordinary language-modeling loss.
6. Matched-budget control. Per-layer factorization B_l·C_l with rank 40 (the same ~737k stored parameters), identical training protocol — this decides whether sharing adds value or is merely a parameter constraint.
7. Full model and attribution. Extended to six matrix types (Q, K, V, O, MLP-in, MLP-out), each with its own shared basis; attention-only, MLP-only, and combined variants.
8. Scaling. The same attention-compression test on GPT-2 medium (24 layers, d = 1024) at the same relative rank k = d/4.
9. Hybrid with residual. The shared basis plus a per-layer low-rank correction W_l ≈ U A_l Vᵀ S⁺ + B_l C_l with residual rank r = 32, initialized from the top-r SVD of the basis error — targeting from the start what the basis failed to capture.

Results

Details the table does not capture. In test 1, nearly the entire drop achieved by “cheap moves” came from plain norm scaling; no real structure was found, only a weak locality gradient (neighboring layers more similar than distant ones). Test 3 also revealed that Q matrices look high-rank in the raw weights, but real inputs excite only a fraction of the directions — layer 2 carries 90% of its output energy in 14 of 768 dimensions. Test 4 showed that captured energy is not the same as function: 96% captured energy cost +4 ppl, but 78% cost +376 — a strongly nonlinear relationship, plausibly because attention’s softmax exponentially amplifies score errors. Test 6 also revealed an optimization advantage for the shared basis beyond the compression ratio itself: training cores between orthonormal bases converged smoothly, while training free factors B·C suffered from scale degeneracy and instability.

Three main findings

Sharing across layers exists, but only in functional space and only through optimization. Layers do not share weights or weight subspaces (tests 1-2: fail). They share functional structure — after reweighting by the covariance directions of real input activations (test 3: score 0.70-0.86) — which analytic projection cannot find (test 4: reconstruction collapses the model to ppl 408), but gradient search inside a frozen shared basis can (test 5: ppl 41.7, ~97% of the loss recovered). At matched memory budget, the shared basis demonstrably carries more function than independent factorizations (test 6: 41.7 vs 46.6 ppl, three seeds, non-overlapping ranges).

Attention can be shared, MLP cannot — and this is not a method artifact. Compressing attention matrices at 9.6x is recoverable to a usable range after fine-tuning (ppl 77.7); compressing MLP matrices destroys the model outright before fine-tuning (ppl 60,748) and training does not fix it (128 at a milder 6.6x, test 7). A hybrid test with a per-layer residual (test 9) sharpened the failure: a rank-32 residual covered only 12-17% of the basis error energy for MLP versus 23-33% for attention — MLP layer individuality is high-rank and distributed, so a small residual cannot capture it. This independently and quantitatively supports the picture from the interpretability literature: attention implements generic, transferable computation across layers, while MLP acts as content-specific memory.

Weights carry a signature of a matrix’s role relative to the residual stream — but this finding did not survive outside the GPT-2 family. On GPT-2 small and medium, a classifier trained only on attention matrices recognized “writers” (O, MLP-out) among MLP matrices with 0.875 accuracy (permutation p = 0.027), and 0.938 in reverse. A later replication across eleven models spanning four architecture families showed the claim does not hold in general: the transfer classifier’s accuracy on the other models sits at or below the majority-class baseline. The signature is carried mainly by weight outlier statistics, which are strongly architecture-dependent and vanish where outliers are weaker. We keep this as a correction to an earlier conclusion — the claim holds only within the GPT-2 family, not across transformers generally.

Context within the broader project

The shared functional basis was the strongest positive signal among nine structural compression approaches tested in the project, and yet, in the sense of repeatable scalability, it did not survive either: it works on attention matrices but not on MLP matrices, so it cannot be deployed across a whole model without further modification. It shows that structure sharing is possible, but not universal enough to replace simpler approaches such as weight quantization or layer merging, which turned out to be more reliable and easier to scale within this project.

Limits

The main measurements were run on GPT-2 small, with a single confirmatory run on GPT-2 medium restricted to attention compression. Calibration, training, and evaluation all used a single dataset (WikiText-2), so part of the recovered quality may be domain adaptation — out-of-domain evaluation was not performed. Fine-tuning was short (3-8k steps), so the final numbers reported here are best read as upper-bound estimates of achievable quality. We did not measure compressed models on downstream tasks, only via perplexity, and we never measured the inference speed of the factorized form. Tests 4-9 ran on raw Q/K matrices without accounting for the gauge freedom of the product W_qᵀW_k (attention scores depend only on this product, so raw matrices are gauge-ambiguous) — part of the layer individuality measured in tests 1-3 may reflect this unaccounted degree of freedom.