arXiv:2605.03953 · v3

SATFormer: Learning When to Use Early Value Residuals in Transformers

首次提交 2026-05-06 00:38

版本更新 2026-10-27 01:34

雷达观测 2026-10-29 00:35

摘要

Recent Transformer architectures expose later layers to representations computed in earlier layers, motivated by the observation that low-level features can become harder to recover as the residual stream is transformed through depth. One such method adds static value residuals: learned mixing coefficients that expose the first-layer value projection V_1 uniformly across tokens and heads. However, we theorize that the usefulness of $V_1$ is unlikely to be constant across tokens, heads, and contexts, and therefore treat early-representation reuse as a retrieval problem. Here, we introduce Selective Access Transformer (SATFormer), which preserves the first-layer value pathway while controlling access with a context-dependent gate. Across models from 130M to 1.3B parameters, SATFormer improves validation loss over the static value-residual and Transformer baselines at every evaluated scale, and improves average zero-shot accuracy over both at all four scales. Its strongest gains appear on retrieval-intensive benchmarks, where it improves over static value residuals by approximately 1.5 average points, while maintaining near baseline throughput and memory. Mechanistic analyses suggest sparse, depth-dependent, head-specific, and category-sensitive access patterns, supporting the interpretation that SATFormer learns selective reuse of early representations rather than uniform residual copying. Our code is available at https://github.com/skyegunasekaran/satformer.