The Key to Going Linear: Analysis-Driven Transformer Linearization

LLM MQIA Long Context ML TAVR Transformer
因果自注意力机制的二次方计算开销,严重制约了长上下文Transformer模型的推理效率。尽管目前已存在大量事后线性化方法(post hoc linearization pipelines),但其中哪些组件真正能够维持模型性能却难以甄别。本研究在严格冻结主干网络(frozen-backbone)的前提下,单独剥离并考察状态更新机制(state update design)的影响。我们发现,Softmax操作本质上依赖于键向量(key)驱动的、秩为1的正交投影;这一发现清晰解释了为何“增量式”(delta-style)网络架构优于纯粹基于门控机制的累加方式。我们进一步识别出一种潜在的近似误差来源,并据此提出若干结构化改进策略,具体包括:引入“汇点标记”(sink tokens)、采用短卷积(short convolutions),以及实施固定预算的缓存路由(fixed-budget cache routing),从而显著缩小剩余性能差距。我们将该线性化方案扩展应用于LLaMA与Qwen系列模型,参数规模最高达320亿,其在MMLU基准测试上的表现超越了此前所有事后线性化基线方法,同时在长上下文检索任务上的性能亦可媲美复杂的自适应缓存框架。
The quadratic cost of causal self-attention severely bottlenecks long-context transformer inference. While numerous post hoc linearization pipelines exist, it is difficult to identify which components preserve model quality. This work isolates the effect of state update design in a strict frozen-backbone regime. We show that softmax relies on key-dependent, rank-1 orthogonal projections, elucidating why delta-style networks outperform purely gated accumulation. We identify a potential source of approximation errors and introduce structural interventions, specifically sink tokens, short convolutions, and fixed-budget cache routing, which reduces the remaining gap. We scale this linearization approach across LLaMA and Qwen models up to 32B parameters, outperforming prior post hoc baselines on MMLU and matching the long-context retrieval of complex adaptive-caching frameworks.
许愿