🤖 machine learning

Length Generalization with Log-Depth Recurrent Units

The paper introduces MLP-LDRU, a log-depth recurrent unit that approximates recurrence via parallel reduction to achieve near-perfect length generalization on regular-language tasks and competitive performance on broader benchmarks, effectively addressing the positional biases of recurrent models and the depth constraints of transformers.

Original authors: Charles Pert, Dalal Alrajeh, Alessandra Russo

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Charles Pert, Dalal Alrajeh, Alessandra Russo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Short Memory" of AI

Imagine you are teaching a child to count. If you only ever practice counting up to 10, the child might get confused when asked to count to 100. They haven't learned the rule of counting; they've just memorized the first ten numbers.

In the world of Artificial Intelligence (neural networks), this is called the Length Generalization problem.

RNNs (Old School): Think of these like a relay race where the baton is passed one person at a time. If the race is long, the first runner has to wait a long time for the baton to come back around. They get tired (forget things) by the time the race ends.
Transformers (Modern Giants): Think of these like a group of people standing in a circle, all shouting at once. They can hear everyone immediately, but if the circle gets too huge, the noise becomes chaotic, and they struggle to understand the pattern if the group size changes from what they practiced.

Both struggle when asked to handle sequences (like sentences or code) that are much longer than what they saw during training.

The Solution: The "Balanced Tree" (MLP-LDRU)

The authors propose a new model called MLP-LDRU. To understand how it works, imagine you have a stack of 8 books and you want to find the total weight of all of them.

The Old Way (RNN): You pick up Book 1, then add Book 2, then add Book 3, and so on, one by one. This takes a long time, and the first book gets "forgotten" by the time you reach the end.
The New Way (MLP-LDRU): You pair up the books.
1. Pair (Book 1 + Book 2) and (Book 3 + Book 4) and (Book 5 + Book 6) and (Book 7 + Book 8).
2. Now you have 4 pairs. Pair them up again: (Pair 1 + Pair 2) and (Pair 3 + Pair 4).
3. Now you have 2 groups. Pair them up one last time to get the final answer.

This is called a Log-Depth Reduction. It's like a balanced tree. No matter how many books you have, everyone gets to the finish line in roughly the same amount of time. The first book isn't "tired" because it didn't have to wait for the last book to be processed first.

The Secret Sauce: The "Magic Glue"

The paper introduces a special "glue" (a mathematical operator) used to combine these pairs. The authors designed this glue to act like associative math.

Associativity means the order of grouping doesn't matter. $(A + B) + C$ is the same as $A + (B + C)$ .
The authors forced their "glue" to behave this way. By teaching the AI that "grouping doesn't matter," the AI learns the underlying rule of the sequence rather than just memorizing specific positions.

The Experiments: The "Grammar Test"

To test this, the authors didn't use messy real-world language (which is hard to measure). Instead, they used Regular Languages.

Analogy: Imagine a strict robot that only accepts sentences following a perfect, simple grammar rule (like "every 'A' must be followed by a 'B'").
They created 21 different grammar puzzles. Some were easy (like checking if a number is even), and some were hard (like keeping track of nested parentheses, similar to balancing a checkbook).
They also invented a new puzzle called Prefix Languages. This is like a game where the first few words determine the entire outcome, but the rest of the sentence is just noise. It tests if the AI can remember the beginning while ignoring the middle.

The Results: The "Perfect Score"

The results were impressive:

The Champion: The MLP-LDRU model got 100% accuracy on 18 out of 21 puzzles, even when the test sentences were 10 to 12 times longer than the training sentences.
Beating the Giants: It outperformed standard Transformers and older RNNs, which often failed completely when the sentences got too long.
The "Why": The authors found that the model failed on the few remaining puzzles not because the "tree" structure was wrong, but because the training data didn't show the AI enough types of combinations. It's like practicing math with only even numbers; when you finally get an odd number, you're stuck. The model needed to see more variety in the "combinations" to master the rule.

Beyond the Grammar: The "List" Test

They also tested the model on ListOps, a task involving nested lists (like a recipe inside a recipe inside a recipe).

While specialized "tree-structured" models were slightly better at this, the MLP-LDRU still performed very well, beating standard Transformers and LSTMs.
They also tested it on standard text classification (like sorting news articles), where it performed competitively, showing that this "balanced tree" idea works even outside of strict grammar rules.

The Takeaway

The paper argues that to make AI that can handle long sequences reliably, we shouldn't just make the model bigger. Instead, we should change how it processes information. By using a balanced tree structure and forcing the model to learn associative rules (where order of grouping doesn't matter), the AI can generalize to lengths it has never seen before, much like a child who understands the concept of counting can count to a million even if they only practiced up to 10.

Technical Summary: Length Generalization with Log-Depth Recurrent Units

Problem Statement

Generalization to out-of-distribution (OOD) sequence lengths remains a central challenge in neural modeling. Existing architectures exhibit distinct limitations in this domain:

Recurrent Neural Networks (RNNs): While offering length-adaptive computation, they suffer from positional bias where earlier tokens undergo more recurrent steps than later ones, leading to long-range memory issues and difficulty in modeling dependencies over extended sequences.
Transformers: These models apply a fixed stack of layers independent of sequence length. While they enable parallel computation, they tend to generalize less reliably to OOD lengths and are constrained by fixed computational depth.

The paper posits that regular languages provide a rigorous testbed for evaluating length generalization because they allow for precise mathematical formulation and exact validation of correctness at any sequence length, directly connecting to automata theory.

Methodology: MLP-LDRU

The authors propose MLP-LDRU (Log-Depth Recurrent Unit), a specific type of binary operator designed to approximate recurrence through parallel reduction. The core hypothesis is that learning operators with an explicit bias toward associativity and identity behavior can improve length extrapolation.

Architecture and Mechanism

Log-Depth Reduction: Unlike standard RNNs which process sequences sequentially ( $O(n)$ depth), MLP-LDRU processes sequences using a balanced binary tree structure. Given a sequence of length $n$ , the model applies a binary operator $\odot_\theta$ in parallel to adjacent pairs for $\lceil \log_2 n \rceil$ steps. This results in a uniform computational depth of $O(\log n)$ across all tokens, mitigating positional bias while enabling parallel computation.
The Operator ( $\odot_\theta$ ): The core operator is an MLP-based function designed to induce approximately associative behavior.
- It takes two input embeddings $h_i, h_j$ and concatenates them.
- An MLP produces gating vectors $g_i, g_j$ .
- These gates modulate the inputs via element-wise multiplication, followed by learned linear projections and biases.
- The final output is a weighted sum of the projected inputs.
- Initialization Bias: The operator is initialized such that the projection matrices are identity matrices and biases are zero. Combined with the element-wise sum structure, this biases the operator toward behaving as a simple sum initially, approximating an associative operation.
Padding Handling: If a sequence length is odd during reduction, it is padded with a zero embedding. The operator is designed such that if either input is zero, the output is the other input, treating zero as a neutral element (identity) to avoid learning an identity map from scratch.
Complexity: MLP-LDRU achieves $O(nd^2)$ work complexity (linear in sequence length) with $O(\log n)$ computational depth, contrasting with the $O(n)$ depth of RNNs and the $O(n^2)$ work complexity of standard Transformers.

Experimental Setup

The authors evaluated MLP-LDRU on:

21 Regular Language Tasks: Including standard benchmarks (e.g., Parity Check, Modular Arithmetic, Tomita languages, Dyck languages $D_n$ ) and a new class of Prefix Languages ( $P_{p,q}$ ) designed to test long-range dependency handling under length extrapolation.
ListOps: A hierarchical task for compositional extrapolation beyond regular languages.
NLP Classification: Benchmarks including GLUE, AG's News, and DBPedia to probe behavior on non-regular data.

Key Contributions

Proposal of MLP-LDRU: A specific Log-Depth Recurrent Unit that uses a binary operator biased toward approximate associativity and identity handling to induce sequential recurrence via parallel reduction.
Introduction of Prefix Languages: A new class of regular tasks ( $P_{p,q}$ ) specifically designed to isolate and test the handling of long-range dependencies under length extrapolation.
Empirical Evaluation: Comprehensive testing on 21 regular tasks, ListOps, and NLP benchmarks, demonstrating superior or competitive performance compared to RNNs, Transformers, and other recursive neural networks.
Analysis of Training Requirements: An analysis linking generalization failures to insufficient coverage of equivalence class (EC) compositions in the training data, suggesting that limited training lengths prevent models from observing rare monoid compositions required for perfect generalization.

Results

Regular Language Tasks

Performance: When trained on sequences of length 1–40, MLP-LDRU achieved 100% OOD accuracy on 18 of 21 tasks (evaluated on lengths 41–500) and at least 99.9% on the remaining 3 when increasing the maximum training length.
Comparison: MLP-LDRU outperformed all Transformer variants (including those with ALiBi, NoPE, and randomized RoPE), Gated DeltaNet, and other baselines on the majority of tasks. It matched or exceeded the capabilities of RNNs and LSTMs, which often struggle with long-range dependencies in these tasks.
Operator Sensitivity: Ablation studies showed that the choice of reduction operator is critical. While simple element-wise sums or linear projections failed on complex tasks like Modular Arithmetic, the MLP-based operator with gating achieved perfect generalization. This suggests that while associativity is necessary, the specific capacity to learn complex compositions is also required.

ListOps and NLP

ListOps: On hierarchical tasks, tree-based baselines (BBT-GRC, RIR-GRC) generally achieved the highest accuracy. However, MLP-LDRU remained competitive, surpassing LSTM and Transformer baselines at larger dataset sizes (500k and 1M samples). The addition of an explicit associativity regularization term further improved performance, indicating that encouraging approximate associativity is beneficial even outside regular languages.
NLP Classification: MLP-LDRU performed competitively with similarly trained Transformer baselines on GLUE, AG's News, and DBPedia, achieving the highest accuracy on DBPedia (98.6%).

Interpretability and Monoid Analysis

Equivalence Classes: Analysis of the D6 task revealed that MLP-LDRU learns to cluster sequence embeddings according to their corresponding equivalence classes (ECs) in the syntactic monoid.
Coverage Gap: Heatmaps of EC compositions showed that training on short sequences (length 10–40) results in sparse coverage of the monoid space, missing rare compositions found in longer sequences (length 480–500). This sparsity gap explains why models trained on short sequences fail to generalize perfectly to longer ones; they have not been exposed to all necessary composition rules.

Significance and Claims

The paper claims that MLP-LDRU offers a robust solution to length generalization on regular tasks by leveraging the mathematical properties of monoids. By inducing an approximately associative operator, the model can approximate a sequential recurrence through parallel reduction, effectively combining the length-adaptive nature of RNNs with the parallel efficiency of tree-based structures.

The authors emphasize that maximum training length is the primary driver of OOD performance. Their analysis suggests that for complex languages, achieving 100% generalization requires sufficient exposure to all possible equivalence class compositions, which may necessitate training on longer sequences than the target test lengths.

While the model demonstrates strong performance on regular languages and competitive results on NLP tasks, the authors remain modest regarding its scope. They note that while regular languages provide exact verification, the extent to which these ideas extend to richer compositional problems under distribution shift in natural language requires further investigation. Additionally, they identify the need for theoretical bounds on the minimum training lengths or sample complexity required to achieve full coverage of monoid compositions.