Technical Summary: Gated Subspace Inference for Transformer Acceleration

1. Problem Statement

The primary bottleneck in large language model (LLM) inference, particularly at batch size one, is memory bandwidth rather than computational throughput. As formalized by the roofline model, the forward pass of a transformer is entirely memory-bandwidth-bound because the arithmetic intensity of linear layer operations ($y = Wx$) is approximately 1 FLOP/byte. Reading weight matrices from High-Bandwidth Memory (HBM) dominates the latency.

Existing acceleration strategies face significant limitations:

• Quantization reduces the size of weight elements but does not reduce the number of elements read.
• Pruning reduces the number of elements but typically requires retraining or fine-tuning.
• Static Low-Rank Factorization (e.g., LoRA, SVD-LLM) decomposes weights offline but fails to adapt to the specific input distribution at inference time.
• Activation Sparsity methods (e.g., CATS, PowerInfer) exploit discrete neuron subsets or element-wise sparsity but do not leverage the continuous low-dimensional structure of the activation manifold.

The paper posits that the activation vectors at any given layer lie approximately in a low-dimensional subspace ($k \ll d$). Consequently, the interaction between weights and activations has a low effective rank. Reading the full weight matrix for every token is wasteful because the input distribution does not visit all directions in the activation space.

2. Methodology: Gated Subspace Inference (GSI)

GSI is a lossless acceleration method that exploits the low effective rank of the weight-activation interaction without retraining, architectural modification, or attention approximation. It extends the Skyline Subspace Inference (SSI) framework from MLP layers to the entire transformer (including QKV and output projections).

2.1. Core Decomposition

For a linear layer with weight matrix $W \in \mathbb{R}^{d_{out} \times d}$ and activation $x \in \mathbb{R}^d$, the activation is decomposed exactly into a subspace component and a residual:
$$x = V_k V_k^\top x + r$$
where $V_k \in \mathbb{R}^{d \times k}$ is a rank-$k$ orthonormal basis representing the activation subspace, and $r$ is the residual. The output $y = Wx$ is computed as:
$$y = (W V_k)(V_k^\top x) + Wr = Mg + Wr$$
Here, $M = W V_k$ is a precomputed, cached "low-rank weight image," and $g = V_k^\top x$ is the projection of the activation onto the subspace.

2.2. The Gated Residual Passthrough

The critical innovation is a per-token gate that determines whether the residual correction $Wr$ is computed.

• Fast Path: If the residual ratio $\rho_t = \|r\| / \|x\|$ is below a threshold $\epsilon$, the output is approximated as $y \approx Mg$. This requires reading only $k$ columns of the weight matrix (or the cached image $M$) instead of the full $d$ columns.
• Slow Path: If $\rho_t \ge \epsilon$, the full computation $y = Wx$ is performed to ensure exactness.

This gate ensures that the output distribution is preserved within a controllable tolerance. Theoretical analysis (Theorem 4) bounds the per-token error on the fast path by $\|W\|_2 \cdot \epsilon \cdot \|x\|$. Crucially, the paper demonstrates that discarding the residual entirely (static projection) leads to catastrophic output degradation, whereas the gated passthrough maintains lossless quality.

2.3. Depth Cascade and Calibration

The method leverages "depth coherence," observing that activation subspaces at consecutive layers are highly correlated (mean cosine similarity $> 0.90$ in deeper layers).

• Cascade Initialization: The basis $V_k$ for layer $l+1$ is initialized from the basis of layer $l$, requiring only minor rank-1 corrections (DGKS updates) rather than a full Singular Value Decomposition (SVD) for every layer.
• Calibration: The basis is computed once via a single forward pass on a representative calibration sequence. For a 6B model, this takes approximately 30 seconds.
• Storage Overhead: The cached images $M$ and bases $V_k$ add less than 4% to the model's memory footprint (approx. 469 MB for a 13 GB model).

3. Key Contributions

1. Gated Subspace Inference (GSI): A method that accelerates all linear maps in a transformer (QKV, output, and MLP projections) by projecting activations onto a learned low-rank subspace and conditionally skipping residual corrections.
2. Lossless Quality via Gating: The paper establishes that the gated residual passthrough is essential. Without it, static projection causes generation collapse. With it, the method achieves character-for-character identical output to the baseline at specific operating points.
3. Depth Cascade: A mechanism that exploits subspace coherence across layers to reduce calibration costs by 96%, revealing that the network operates on a coherent low-dimensional manifold propagating through depth.
4. Orthogonal Acceleration: GSI operates on the hidden dimension ($d$-axis), complementing existing attention acceleration methods like ADA (which operate on the token dimension $T$-axis). Together, they cover the full forward pass.

4. Experimental Results

Experiments were conducted on AMD MI300X GPUs using GPT-2 (124M), GPT-J (6B), and OPT (6.7B) models.

• Speedup: The method achieves effective speedups of 3.0× to 10.5× on linear-layer weight reads.
  • GPT-J 6B at operating point $k=256, \epsilon=0.10$: 15.6× speedup on weight reads with 99.8% of tokens taking the fast path.
  • OPT 6.7B at $k=256, \epsilon=0.10$: 10.5× speedup.
  • GPT-2 124M at $k=256, \epsilon=0.10$: 3.0× speedup.
• Quality Preservation:
  • Perplexity: Ratios are consistently below 1.00 (e.g., 0.991 for GPT-J), indicating the accelerated model is statistically indistinguishable from or slightly better than the baseline.
  • Token Agreement: Top-1 token agreement exceeds 98% across all configurations.
  • Generation Fidelity: At the target operating point ($k=256, \epsilon=0.05$ or $0.10$), GPT-J and GPT-2 produce 100% character-for-character identical output to the baseline over 50 greedy tokens.
• Effective Parameter Count: The inference cost is proportional to the "effective parameter count" (the parameters corresponding to visited directions in activation space). For GPT-J 6B, the effective count drops from 6B to ~375M (a 16× reduction) at the optimal operating point.
• Negative Control: Static projection (no residual) at $k=128$ increased perplexity by 52% and caused generation collapse, validating the necessity of the gate.

5. Significance and Claims

The paper claims that GSI provides a lossless acceleration for transformer inference that fundamentally shifts the cost model from the total parameter count to the effective parameter count.

• No Retraining: The method works on pre-trained models without any fine-tuning or architectural changes.
• No Approximation of Attention: It does not alter the attention mechanism, making it compatible with other acceleration techniques like FlashAttention or ADA.
• Structural Insight: The results confirm that the low effective rank of the weight-activation interaction is a structural property of trained transformers, holding across different architectures and model sizes.
• Practical Impact: By reducing the memory bandwidth requirement for linear layers by a factor of $d/k$ (e.g., 16×), GSI addresses the primary bottleneck in batch-size-one inference, enabling significantly faster decoding times (e.g., reducing a 1.3 ms forward pass to 0.22 ms when combined with attention acceleration).

The authors conclude that the gated residual passthrough is the distinguishing factor that allows for high compression ratios without sacrificing the mathematical correctness of the output, making it a viable path for lossless inference acceleration.