Chain and Causal Attention for Efficient Entity Tracking (2024)

Abstract: This paper investigates the limitations of transformers for entity-tracking tasks in large language models. We identify a theoretical constraint, showing that transformers require at least \log_2 (n+1) layers to handle entity tracking with n state changes. To address this issue, we propose an efficient and frugal enhancement to the standard attention mechanism, enabling it to manage long-term dependencies more efficiently. By considering attention as an adjacency matrix, our model can track entity states with a single layer.Empirical results demonstrate significant improvements in entity tracking datasets while keeping competitive performance on standard natural language modeling. Our modified attention allows us to achieve the same performance with drastically fewer layers. Additionally, our enhanced mechanism reveals structured internal representations of attention. Extensive experiments on both toy and complex datasets validate our approach. Our contributions include theoretical insights, an improved attention mechanism, and empirical validation.

## introduction transformers @xcite have deeply renewed state-of-the-art approaches in many natural language processing (nlp) tasks. more specifically, the attention mechanism can readily handle long-range dependencies and complex linguistic phenomena. this kind of architecture surpasses the previous ones, like markov models and recurrent networks. for instance, transformers exhibit an expressivity in fluent language generation, even for long text, that was not possible before. moreover, self-attention is very powerful in deriving rich and meaningful representations of languages. this major paradigm shift paved the way for large language models (llm). starting with gpt @xcite and bert @xcite , many successors have emerged, witnessing a massive scaling up of model sizes by several orders of magnitude along with unprecedented improvements in performance. nowadays, state-ofthe-art approaches mostly rely on llms for most of the nlp, setting a new standard for the development and application of language models in both academia and industry @xcite @xcite @xcite . however, the amount of resources needed to train and use them has also increased by several orders of magnitude. this raises important concerns regarding their ecological impact and limits their use and applications @xcite . this sustainability issue has motivated different lines of work to make llms more frugal by optimizing architectures and reducing computational overhead without a drop in performance. various approaches have been proposed, such as pruning, quantization, and efficient training techniques, to reduce the carbon footprint of llms @xcite @xcite @xcite . most of these methods address this call for frugality in a task-agnostic way, while the complexity and computational resources may greatly differ depending on the downstream task under study. some tasks like entity tracking imply long-range dependencies with chain-like memorization @xcite . more precisely, for entity tracking, the model has to maintain and update the states of entities as described through a sequence of operations. while it looks fairly easy for humans, this is challenging for standard transformer models. @xcite indeed show that current llms have underwhelming performance on causal reasoning of event plausibility and entity states, despite the crucial part it plays in natural language reasoning. moreover, the complexity drastically increases with the sequence length. therefore, in this work, we focus on the entity tracking task, which is critical in many applications like questionanswering, dialogue systems, text-based games, and reasoning @xcite @xcite . in this paper, we introduce a modified attention layer to overcome this limitation. our goal is to enhance the standard attention mechanism with a new design, to better cope with long-term and chainlike dependencies. for this purpose, our approach relies on a different interpretation of the attention matrix as an adjacency matrix. our contributions are the following: ## background and related work in this work, we focus on the decoder-only architecture @xcite based on transformers @xcite . this kind of generative architecture is at the core of causal language modeling (clm). while there are many variants @xcite @xcite @xcite , they all rely on the same framework. after a tokenization step of the texts, the embedding layer converts the pre-processed input in a sequence of vectors, i.e a matrix. then, l layers of transformers process this matrix, each made of a self-attention layer followed by a feedforward network. the output vectors are finally projected into logits for prediction purposes. ## theoretical limits of transformers for entity tracking following the introduction of entity tracking in section 2.2, we here more formally define this task to specify the theoretical limitations of transformers. ## experiments throughout all of our experiments, we base our models on the gpt-2 @xcite architecture, which is the base of most modern llms. we compare standard transformer models and our proposed enhanced attention layer on 3 experiments: an abstract toy dataset, a textual entity tracking task, and fine-tuning on a standard language modeling task. while our method introduces a new hyperparameter γ, we used γ = 0.9, as we show in appendix c that any value between 0.8 and 1 seems to work fine. detailed experimental information is in appendix d. figure 3: test accuracy of transformer models during training on our toy dataset. the standard transformer architecture struggles to learn the task and needs 4 to 5 layers to reach 100%, while our enhanced attention consistently solves the task with only one layer. we show the average, min and max values over 4 runs for each model. ## discussion our results correlate with the theoretical prediction that transformers require at least log 2 (n + 1) layers for entity tracking tasks with n jumps. indeed, our toy dataset has @xmath0. similarly, our advanced version of the boxes dataset has @xmath1. this confirms the chacal model, with our improved attention mechanism, successfully learns entity tracking tasks using the minimal number of layers, outperforming by far standard transformers on these specific datasets. this demonstrates the frugality and efficiency of our method, as it manages complex tasks with fewer resources compared to standard transformers. we however notice a slower convergence at the beginning of training and a greater variance in the training curves. we believe that chacal could greatly benefit from understanding better its training dynamics (e.g. a more suited initialization, or regularization), which we leave as future work. while our modified attention mechanism does not significantly enhance general language modeling tasks, this observation is in line with several previous studies, such as @xcite , @xcite . entity tracking is a relatively rare subtask within language modeling, indicating that pre-training on natural language alone does not foster robust entitytracking capabilities. we also explain the absence of immediate benefits in language modeling by the large size of the models, which have an abundance of layers that allow them to handle most entity tracking tasks in theory. as chacal was specially designed to reduce the number of attention layers, one may try to explore other architectural choices more suited to our purpose, involving shallower but wider models for instance. a notable finding is that our model with γ = 0.9 across all layers reaches performances similar to standard transformers. instead of fine-tuning a model that was pre-trained using standard attention, pre-training from scratch with our novel attention mechanism could lead to a more efficient attention structure optimized for entity tracking. finally, the high computational efficiency of our method mainly relies on the attention matrix being lower-triangular, which is only the case in decoder transformers for causal language modeling. further research may attempt to generalize chacal to noncausal transformers such as bert @xcite and many others @xcite @xcite . ## conclusion in this study, we have theorized that transformers require at least log 2 (n + 1) layers for effective entity tracking tasks with n jumps. multiple experiments on entity tracking tasks, with both a toy dataset and a more complex natural language dataset, validate our claims and consistently demonstrate that transformers need at least this many layers to achieve perfect performance. the chacal model, leveraging our novel attention mechanism, excels in entity tracking tasks while drastically lowering layer requirements compared to standard transformers, and without adding more parameters nor requiring special tuning. this underscores the frugality and computational efficiency of our approach. at last, fine-tuning experiments of gpt-2 with our enhanced attention on a natural language modeling task suggests that cha-cal can at least match the standard transformer models. several key questions emerge: how does pretraining from scratch with chacal compare to fine-tuning existing models? what are the optimal settings for γ across different layers and tasks, and can it be dynamically adjusted during training to enhance the model's adaptability and performance? in conclusion, this study advances our understanding of attention mechanisms in transformers, particularly in optimizing models for entitytracking tasks. by bridging theory with practical insights, we lay a foundation for future research aimed at developing more efficient transformer architectures tailored to specific application domains.