Investigating Recurrent Transformers with Dynamic Halt

TL;DR

This work systematically compares depth-wise and chunk-wise recurrent augmentations of Transformer architectures and introduces two extensions, GUT and GUTLB, that integrate gating and a global mean-based dynamic halting mechanism. Through tasks like ListOps, Logical Inference, Flip-flop, and Long Range Arena, the authors reveal that gating and global halting can enhance depth-wise recurrence, while chunk-wise recurrence offers robustness to length generalization and distractors, with TLB-based variants often excelling in OOD settings. The results highlight complementary strengths and trade-offs between recurrence styles, demonstrating compute benefits from dynamic halting and suggesting future directions such as Mixture-of-Experts and alternative attention mechanisms. Overall, the study provides a nuanced understanding of how recurrences and dynamic halting shape inductive biases in transformers for algorithmic and long-range reasoning tasks.

Abstract

In this paper, we comprehensively study the inductive biases of two major approaches to augmenting Transformers with a recurrent mechanism: (1) the approach of incorporating a depth-wise recurrence similar to Universal Transformers; and (2) the approach of incorporating a chunk-wise temporal recurrence like Temporal Latent Bottleneck. Furthermore, we propose and investigate novel ways to extend and combine the above methods - for example, we propose a global mean-based dynamic halting mechanism for Universal Transformers and an augmentation of Temporal Latent Bottleneck with elements from Universal Transformer. We compare the models and probe their inductive biases in several diagnostic tasks, such as Long Range Arena (LRA), flip-flop language modeling, ListOps, and Logical Inference. The code is released in: https://github.com/JRC1995/InvestigatingRecurrentTransformers/tree/main