通过结构化设计,让模型在速度与表达力间自由权衡。
Trading Complexity for Expressivity Through Structured Generalized Linear Token Mixing

- 提出统一框架,分离输入影响与历史信息传播机制
- 新结构使计算复杂度可调控,理论证明表达力提升
- 适用于注意力与状态空间模型,适合高效架构设计者
Token mixing 层在语言模型中对学习和生成长程依赖关系起关键作用。其效率取决于解码速度与内存开销、缓存大小之间的权衡。针对因果生成,本文提出一个统一框架,将两个核心特性解耦:(i) 单步生成中输入对输出的直接作用;(ii) 历史输出间的信息递归传播。该框架涵盖主流架构如注意力机制与状态空间模型,并通过引入结构化设计,允许每个状态依赖多个过去状态而非仅前一状态,从而推广递推方程。我们设计的新递推模式在理论上可实现目标复杂度,同时提供表达力的分析依据——以运行时为代价换取更强表达能力。在合成任务与语言建模上的实证验证表明,该方法为跨模型家族的高效且高表达力 token mixer 设计提供了统一工具集。
原文摘要 · Abstract (English)
Token mixing layers play a key role in how language models can learn and generate long-range dependencies. Their efficiency relies on the necessary trade-off between decoding speed and the memory requirements, along with the cache size. Considering causal generation, this paper explores new trade-offs thanks to a unified framework which separates two crucial features: (i) the direct influence of inputs on outputs in one generation step; (ii) the recurrent propagation of information through past outputs. This framework encompasses major architectures such as attention and state-space models, but also generalizes the recurrence equations by allowing each state to depend on multiple past states rather than only the immediate predecessor. By introducing structure, we design new recurrence patterns that provably achieve the desired complexity, while providing theoretical insights on their expressivity -- trading runtime for expressivity in a principled way. Empirical validation is performed on synthetic tasks, along with language modeling. Together, these results provide a unified toolkit for the understanding and design of efficient and expressive token mixers across model families.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。