跳到正文
arXiv cs.AI· Alexey Zemtsov, Albina Klepach, Vadim Lomshakov, Andrei Polubarov, Nikita Lyubaykin, Alexander Derevyagin, Vladislav Kurenkov, Alexander Nikulin·· 3 天前AI 评分22

循环 Transformer 的块状共形注入:让循环步骤学会追踪状态

Rotate the Loop: Blockwise Conformality for Looped Transformers

AI 导读

研究提出"块状共形注入"方法,将循环 Transformer 的残差拆分为坐标对齐的组并在各组内旋转,使每组拥有独立时间尺度而组内几何保持精确。实验显示该方法在末 token 准确率这一状态追踪任务和语言建模上均取得性能提升。论文被 NeurIPS 2026 Workshop "Transitioning from Pre-Training to Post-Training" 接收。

正文 · 原文

View PDF HTML (experimental)

Abstract:Looped transformers reuse one shared block across many recurrent steps, so depth becomes a quantity that can be spent at runtime rather than a stack of distinct layers. Those extra steps help only if a write made on one recurrent step is still readable on later ones, and that depends on the linear map applied to the residual between recurrent steps. We show that whether the loop learns to track state is a geometric question about the form of this map, and we propose blockwise conformal injections: the residual splits into coordinate-aligned groups and rotates inside each one, so every group has its own timescale while the geometry stays exact within it. Our experiments show a performance increase on a demanding form of state tracking - last-token accuracy when the rest of the prefix is already correct - and on language modelling, where a signal written earlier has to be retrieved at the end of the sequence.
Comments: Accepted at NeurIPS 2026 Workshop "Transitioning from Pre-Training to Post-Training"
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Machine Learning (stat.ML)
Cite as: arXiv:2610.14030 [cs.LG]
  (or arXiv:2610.14030v1 [cs.LG] for this version)
  https://doi.org/10.48550/arXiv.2610.14030

arXiv-issued DOI via DataCite (pending registration)

Submission history

From: Alexey Zemtsov [view email]
[v1] Wed, 30 Sep 2026 23:32:34 UTC (141 KB)

来源:arXiv cs.AI · arxiv.org