Why Depth Matters in Parallelizable Sequence Models: A Lie Algebraic View
Authors: Gyuryang Heo, Timothy Ngotiaoco, Kazuki Irie, Samuel J. Gershman, Bernardo L. Sabatini
Organizations: Howard Hughes Medical Institute, Department of Neurobiology, Harvard Medical School, Boston, MA, USA · Kempner Institute for the Study of Natural and Artificial Intelligence, Harvard University, Cambridge, MA, USA · Department of Psychology and Center for Brain Science, Harvard University, Cambridge, MA, USA
Scalable sequence models, such as Transformer variants and structured state-space models, often trade expressivity power for sequence-level parallelism, which enables efficient training. Here we examine the bounds on error and how error scales when models operate outside of their expressivity regimes using a Lie-algebraic control perspective. Our theory formulates a correspondence between the depth of a sequence model and the tower of Lie algebra extensions. Echoing recent theoretical studies, we characterize the Lie-algebraic class of constant-depth sequence models and their corresponding expressivity bounds. Furthermore, we analytically derive an approximation error bound and show that error diminishes exponentially as the depth increases, consistent with the strong empirical performance of these models. We validate our theoretical predictions using experiments on symbolic word and continuous-valued state-tracking problems.
Figures & tables
Figure 1: Geometric intuition of Lie theory. Initialized at e , consider actions A and B sequentially, and undo actions B−1 then A−1 . Composition of actions ABB−1A−1 would return to e . However, switching the order of undoing, by using ABA−1B−1 , might not return to e , incurring a discrepancy by landing on point e′ . Lie theory provides a measure of the potential offset between e and e′ .
Class
Finite group example
Related problems
Abelian
C2,C3,C60
Parity; Cyclic gridworld
Nilpotent
D8,H3
Quantum spin
Solvable
S3,A4,S4
2D affine transformation
Non-solvable
A5,S5
Rubik’s cube; 3D rotation
Table 1: Class of representative word problems on finite group.
Abelian
Nilpotent
Solvable
Model
L
C2
C3
D8
H3
S3
S4
Transformer
1
0.00
0.00
0.00
0.00
0.00
0.00
(+-only) Mamba
1
0.00
0.00
0.00
0.00
0.00
0.00
(+-only) GLA
1
0.00
0.00
0.00
0.00
0.00
0.00
2
0.00
0.00
0.00
0.00
0.00
0.00
Signed Mamba
1
1.00
0.00
0.00
0.00
0.00
0.00
Table 2: Length generalization performance measured as sequence-level accuracy for various categories of word problems. Models are trained on length 128 and tested on 256. L denotes the number of layers. DeltaProduct uses the theoretically necessary number of Householder products depending on the problem. The blue color cell highlights perfect generalization. The hyper-parameter search space is in App. D.1 . The number of training sequences is 500K except for C2 which used 100K.
Figure 2: Maximum sequence length on the A5 WP for each model with varying number of layers to achieve >90% training sequence accuracy. All models are trained on length up to 128. Deep models ( > 4 layers) that failed to achieve a longer sequence length than shallower models are not shown. T in the legend (⌊log2T⌋+1) is sequence length (i.e., y-axis). The reference curve is the upper bound from Prop. 3.7 .
Figure 3: Mean squared error on the rotated vector prediction task at various sequence lengths for transformer, GLA, and signed Mamba, with different numbers of layers. Left and right columns illustrate the results on training and test sets, respectively. Standard error is shown using 3 seeds. Performance of DeltaProduct with 4 Householder products is shown as a reference.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Parameters
Values
Hidden size
{64, 128, 256}
Learning rate
{1e-2, 1e-3, 3e-4, 1e-4}
Batch size
{256, 2048}
Appendix
Table 3: Hyper-parameter search space. The number of heads is set to 8 for transformer variants, and d_state (i.e., the “head dimension”) is set to 16 for SSMs. Each configuration is run with 3 random seeds.
Figure 4: Mean squared loss on the rotated vector prediction task at various sequence length for AUSSM, with different numbers of layers, on train (left) and test (right) sets. Standard error is shown using 3 seeds. Performance of DeltaProduct with 4 Householder products is shown as a reference.
Probes are routinely paired with an intervention: ablate the direction the probe found, run the model, and read the change in task accuracy, taking a large drop as evidence that the computation depends on what the probe read and a near-zero drop as evidence that it does not. Either inference requires that the ablation have removed the target from the layer. We find that the ablation does not remove what it targets. A probe refitted on the ablated activations recovers its original accuracy in every cell we test, and keeps recovering when the probe's entire row space is deleted rather than a single axis, because the quantity survives in the orthogonal complement. Because a refitted probe recovers, neither a large task drop nor a near-zero one establishes whether the model needed the target, and one probe fit detects this. Replacing the ablation with iterative nullspace projection, scored against random subspaces of matched dimension, reverses the conclusion: representations that looked causally inert carry most of the task. The correction also separates where a variable is most readable from where deleting it does most damage, and those are not the same layer in any pretrained model we study. The erasure is defined by a linear probe family, so removing a nonlinearly encoded quantity remains open.
Language modeling using Transformers has become commonplace despite their fixed computational depth and quadratic runtime with respect to input tokens. Recurrent models on the other hand offer linear depth but no parallel execution. In this work, we extend balanced-tree recursive operators from sequence encoding to autoregressive prediction, enabling all prefix representations to be computed with logarithmic depth and linear runtime. Our experiments provide an initial characterization of this model class, demonstrating robust length extrapolation and performance approaching that of ALiBi-based Transformers, highlighting its potential as an alternative architecture for language modeling.
Yiqin Wang, Nuri Cingillioglu, Charles Pert
Department of Computing Imperial College London London, United Kingdom
Deep neural networks are widely believed to derive their expressive power from their ability to form hierarchical representations, capturing progressively more abstract and compositional features across layers. In language modeling, transformers have emerged as the dominant architecture, with early layers capturing local syntactic patterns and later layers encoding more complex clause-level dependencies. While this intuition has shaped model design, there remains a lack of rigorous theoretical work demonstrating how deep transformers represent such hierarchical structures. In this work, we analyze the expressiveness of deep transformer models through the formal lens of bounded-depth, non-recursive context-free grammars. For this class of grammars, we explicitly construct transformers with positional attention whose depth grows linearly with grammar depth, while the neuron count scales with the number of derivation-tree shapes and quadratically with the number of production rules. Our theoretical results support the linear representation hypothesis by demonstrating that these architectures possess the structural capacity to encode abstract grammatical states into low-dimensional, linearly separable subspaces within the residual stream.
Vinoth Nandakumar, Qiang Qu, Pramod Thebe +2
University of Sydney · San Francisco State University · IIT Madras