Compositional Reasoning with Transformers, RNNs, and Chain of Thought

It is well understood that different neural network architectures are suited to different tasks, but is there always a single best architecture for a given task? We compare the expressive power of transformers, RNNs, and transformers with chain of thought tokens on a simple and natural class of tasks we term Compositional Reasoning Questions (CRQ). This family captures multi-step problems with treelike compositional structure, such as evaluating Boolean formulas. We prove that under standard hardness assumptions, none of these three architectures is capable of solving CRQs unless some hyperparameter (depth, embedding dimension, and number of chain of thought tokens, respectively) grows with the size of the input. We then provide constructions for solving CRQs with each architecture. For transformers, our construction uses depth that is logarithmic in the problem size. For RNNs, logarithmic embedding dimension is necessary and sufficient, so long as the inputs are provided in a certain order. For transformers with chain of thought, our construction uses n CoT tokens for input size n. These results show that, while CRQs are inherently hard, there are several different ways for language models to overcome this hardness. Even for a single class of problems, each architecture has strengths and weaknesses, and none is strictly better than the others. solve all CRQs of size n (Theorem 5.4). This ability depends on the inputs being arranged in a particular order (Algorithm 1); if they are ordered adversarially, RNNs require O(n) hidden dimension (Theorem 5.2).

4. In Section 6, we prove that transformers augmented with O(log n) CoT tokens cannot solve CRQs of size n, but transformers augmented with O(n) CoT tokens can (Theorem 6.1).

Each of these results fills a gap in the literature; taken together, they demonstrate the fundamental trade-offs between different models (Table 1). While deep transformers are highly parallelizable, they require O(log n) depth in the worst case, so the model size must depend (albeit mildly) on the problem size. Likewise, RNNs use little compute but must grow to handle larger problems, although their success depends on the order of the inputs. Chain of thought allows a single, logarithmic-size model to handle any CRQ, but it runs slowly and is not parallelizable. Overall, our work reveals a rich complexity landscape for an important class of reasoning problems.