arXiv:2604.25930cs.CLcs.LG2026-04

用结构化循环状态实现高效语言建模,支持精确检索。

Associative-State Universal Transformers: Sparse Retrieval Meets Structured Recurrence

  • 引入共享循环模块与混合状态更新机制,提升参数效率。
  • 在字节级维基上达5.083比特/字节,参数减少53.8%仍保持性能。
  • 稀疏槽路由+直接指针融合显著提升检索准确率至99.2%。

我们探讨结构化循环状态能否作为语言建模的紧凑关联基础,同时支持精确检索。提出UniMatrix系列模型,采用通用变换器架构,在深度上复用共享循环块,并引入混合状态更新、ROSA风格残差路径和词元条件嵌入调制。在字节级WikiText-2、合成关联回忆、Apple MPS吞吐量分析及三词交互修正基准上评估。小规模下,UniMatrix-Core与UniMatrix-ROSA在参数匹配时略优于Transformer(5.084 vs 5.124比特/字节),但原始UniMatrix在关联回忆任务中仅接近随机水平(25.4%),而变压器可达25.4%。后续改进模型UniMatrix-Assoc提升有限。相比之下,添加稀疏槽路由与直接指针对数融合的UniMatrix-SparsePointer,在原实验中达到75.6%,无丢弃优化后达99.2%,且比基准变压器少用53.8%参数。消融实验表明,性能提升源于充足的槽容量与精确的指针级输出路由。总体而言,结构化循环状态具潜力且参数高效,但强长程行为仍需显式稀疏检索与更优核函数。

原文摘要 · Abstract (English)

We study whether a structured recurrent state can serve as a compact associative backbone for language modeling while still supporting exact retrieval. We introduce UniMatrix, a Universal Transformer style family that reuses a shared recurrent block across depth and augments it with hybrid state updates, a ROSA-style residual path, and token-conditioned embedding modulation. We evaluate these models on byte-level WikiText-2, synthetic associative recall, throughput profiling on Apple MPS, and a corrected benchmark for triple-token interactions. At small scale, UniMatrix-Core and UniMatrix-ROSA slightly outperform a parameter-matched Transformer on WikiText-2 while using many fewer parameters, reaching 5.084 and 5.083 bits-per-byte versus 5.124. The main negative result is equally important: on associative recall, the original UniMatrix family remains near chance while the Transformer reaches 25.4 percent, showing that compressed recurrent state alone is not enough for exact lookup. A retrieval-oriented follow-up, UniMatrix-Assoc, helps only marginally. By contrast, UniMatrix-SparsePointer, which adds sparse slot routing and direct pointer-logit fusion, reaches 75.6 percent on the original pilot recipe and 99.2 percent on a no-dropout follow-up while using 53.8 percent fewer parameters than the Transformer baseline. Ablations show that the gain comes from sufficient slot capacity and exact pointer-level output routing. Overall, structured recurrent state is promising and parameter-efficient, but strong long-range behavior still requires explicit sparse retrieval and better kernels.

循环神经网络稀疏检索参数效率语言建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。