Recurrent transformers trade depth for time with latent thinking steps
Trading Depth for Time in Recurrent Transformers
Machine Learning
Summary
Transformers are a type of computer program used to understand and generate language, and they work by applying many layers of computation. The study looks at whether it’s better to add more layers of computation in sequence (depth) or to spread computation across steps over time by adding special "thought" tokens between words. The authors find that adding these thought tokens lets a shallower model nearly match the performance of a deeper one while using fewer parameters, meaning it can be more efficient. This suggests that thinking more over time could replace simply going deeper in these models.
What this means in practice
- •For language model engineers: Design language models that achieve high accuracy with fewer parameters by integrating temporal thought tokens instead of stacking more Transformer layers.
- •For chatbot developers: Create efficient chatbots with improved response quality by using latent recurrent layers to refine token predictions without significantly increasing model size.
Authors
Zeyi Huang, Xuehai He, Yong Jae Lee, Yelong Shen
Abstract
Recurrent Transformers increase computational depth through temporal recurrence, feeding each token's high-level hidden state into the computation of the next. This raises a natural question: is additional computation better spent on more temporal steps or greater physical depth? We investigate this question using Latent Recurrent Transformers (LRTs), which retain one backbone forward pass per vocabulary token during decoding and provide a controlled setting for comparing these two ways of adding computation. Specifically, we insert a latent thought token between consecutive vocabulary tokens. Each thought token passes through the same $L$ layers as a vocabulary token, sharing the backbone parameters and providing an additional stage of hidden-state refinement before predicting the next token. We compare this $L$-layer LRT against a $2L$-layer LRT without thought tokens. Both execute $2L$ Transformer blocks per vocabulary token during decoding, but the thought-token model uses fewer parameters. On 16- and 20-layer mixture-of-experts NanoChat backbones, one thought token brings the shallower model within 0.006 and 0.004 bits per byte of its double-depth counterpart, recovering 67% and 81% of the improvement with approximately 48% fewer total parameters. These results suggest that temporal thinking offers a parameter-efficient alternative to increasing physical depth in recurrent Transformers.