让递归Transformer通过记忆缓冲提升性能,参数少33%却更优。
MeSH: Memory-as-State-Highways for Recursive Transformers
- 用显式记忆缓冲和轻量路由,动态分化每轮计算
- 在1.4B模型上超越更大非递归模型,准确率+1.06%
- 适合追求高效低参模型的开发者和研究者
递归Transformer通过复用参数并多次迭代隐藏状态,将计算深度与参数深度解耦。然而,在匹配计算量下,参数较少的递归模型常落后于非递归模型。通过探测隐藏状态,我们发现性能差距源于两个主要瓶颈:计算模式同质化(每轮采用相似计算方式)和信息过载(长期与瞬时信息共存于单一隐藏状态)。为此,我们提出记忆作为状态高速公路(MeSH),将状态管理外化至显式内存缓冲,并使用轻量级路由器动态分化各轮计算。可视化探测证实,MeSH成功通过迭代间功能专一化解决上述问题。在Pythia系列(160M-6.9B)上,增强后的递归模型持续优于基线,并在1.4B规模上超越更大非递归模型,平均下游准确率提升+1.06%,同时非嵌入参数减少33%。分析表明,MeSH是构建更强递归模型的可扩展且原理清晰的架构。代码已开源。
原文摘要 · Abstract (English)
Recursive transformers reuse parameters and iterate over hidden states multiple times, decoupling compute depth from parameter depth. However, under matched compute, recursive models with fewer parameters often lag behind non-recursive counterparts. By probing hidden states, we trace this performance gap to two primary bottlenecks: undifferentiated computation, where the core is forced to adopt a similar computational pattern at every iteration, and information overload, where long-lived and transient information must coexist in a single hidden state. To address the issues, we introduce a Memory-as-State-Highways (MeSH) scheme, which externalizes state management into an explicit memory buffer and employs lightweight routers to dynamically diversify computation across iterations. Probing visualizations confirm that MeSH successfully resolves the pathologies by inducing functional specialization across iterations. On the Pythia suite (160M-6.9B), MeSH-enhanced recursive transformers consistently improve over recursive baselines and outperforms its larger non-recursive counterpart at the 1.4B scale, improving average downstream accuracy by +1.06% with 33% fewer non-embedding parameters. Our analysis establishes MeSH as a scalable and principled architecture for building stronger recursive models. Our code is available at https://github.com/LivingFutureLab/MeSH/ .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。