Reproduce “Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior” on OpenResearch

About this paper

Zeyi Huang, Xuehai He, LiLiang Ren, Yiping Wang, Baolin Peng, Hao Cheng, Shuohang Wang, Pengcheng He and 3 more

A new Latent Recurrent Transformer (LRT) architecture reuses a high-level hidden state from the previous token as recurrent memory for the current token, enabling efficient autoregressive decoding with a single forward pass. It demonstrates improved compute-quality trade-offs, achieving lower Bits Per Byte and higher in-context learning scores compared to standard transformers, and matches or exceeds the performance of multi-forward inference models while maintaining single-forward efficiency.

Key insights

• High-level hidden states from preceding tokens, already computed during standard autoregressive decoding, contain rich latent context that can be effectively repurposed as recurrent memory without adding new inference steps.
• Combining both KV Projection and Residual Injection for recurrent memory delivery offers complementary benefits, leading to better performance than either method used in isolation.
• An upper-middle layer (e.g., layer 12 in a 20-layer model) serves as the optimal source for recurrent memory, suggesting that high-level but not overly specialized representations are most beneficial for cross-token recurrence.

Summary by alphaXiv.