💬 NLP

PaTH Attention: Position Encoding via Accumulating Householder Transformations

This paper introduces PaTH, a flexible, data-dependent position encoding scheme based on accumulated Householder transformations that outperforms the standard Rotary Position Encoding (RoPE) in language modeling tasks while offering efficient parallel training and the ability to convert pretrained RoPE models.

Original authors: Songlin Yang, Yikang Shen, Kaiyue Wen, Shawn Tan, Mayank Mishra, Liliang Ren, Rameswar Panda, Yoon Kim

Published 2026-02-05

📖 4 min read☕ Coffee break read

CC0 1.0

Original authors: Songlin Yang, Yikang Shen, Kaiyue Wen, Shawn Tan, Mayank Mishra, Liliang Ren, Rameswar Panda, Yoon Kim

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.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 "Blind" Memory

Imagine a large language model (like the one you are talking to) as a very smart librarian who reads a long book to answer your questions. To understand the story, the librarian needs to know not just what the words are, but where they appear in the book.

Current models use a system called RoPE (Rotary Position Encoding) to remember word positions. Think of RoPE like a fixed, pre-printed ruler.

If you move a word from page 1 to page 2, the ruler shifts it by a set amount.
The Flaw: This ruler is "blind" to the content. It treats the word "apple" and the word "rocket" exactly the same way just because they are in the same spot. It doesn't care if the story is about a fruit or a spaceship. Because of this, the librarian sometimes gets confused when the story requires complex, step-by-step logic (like tracking who owns which item in a long list).

The Solution: PaTH (The "Smart" Ruler)

The authors introduce PaTH, a new way to remember positions. Instead of a fixed ruler, PaTH is like a dynamic, shape-shifting ruler that changes its shape based on the story it is reading.

How it works: As the librarian reads each word, PaTH applies a special mathematical "twist" (called a Householder transformation) to the memory of previous words.
The Analogy: Imagine you are walking through a hallway with mirrors.
- RoPE: The mirrors are fixed. No matter what you wear, your reflection looks the same.
- PaTH: The mirrors are smart. If you are wearing a red hat, the mirror twists your reflection one way. If you are wearing a blue hat, it twists it another way. The "twist" depends on the actual content (the hat) you are holding.

This allows the model to build a "path" of memory that adapts to the specific data it is processing, making it much better at solving puzzles that require tracking states (like remembering a variable's value in a computer code).

The Magic Trick: Doing It Fast

Usually, making a ruler that changes shape for every single word would be incredibly slow and computationally expensive. It would be like trying to recalculate the entire hallway's mirror angles every time you take a step.

The paper's second major contribution is a speed-up algorithm (similar to something called FlashAttention).

The Analogy: Instead of recalculating the whole hallway every time, the authors found a clever shortcut. They realized that if you group the mirrors into small blocks, you can calculate the "twist" for the whole block at once using a compact mathematical formula.
The Result: They built a system that runs almost as fast as the old fixed-ruler system (RoPE) but keeps the "smart, changing" benefits of the new system.

What Did They Prove?

The authors tested this new system in two ways:

Synthetic Puzzles (The "Training Wheels"):
They gave the model simple logic games, like a "Flip-Flop" game where it has to remember the last time a switch was flipped, or a "Word Problem" game involving complex math rules.
- Result: The old models (RoPE) often failed these puzzles, getting confused by the sequence. The new PaTH model solved them almost perfectly, even with fewer layers of "brain" than the others.
Real Language Tasks:
They trained models on real text (books, code, conversations).
- Result: PaTH models were better at understanding long contexts (reading very long books without forgetting the beginning) and performed better on reasoning tasks, especially in coding and math.
- Bonus: They showed that you can take an existing model trained with the old "fixed ruler" (RoPE) and convert it to use the "smart ruler" (PaTH) with just a little bit of extra training, without starting from scratch.

The Bottom Line

PaTH is a new way for AI to remember the order of words. Instead of using a rigid, one-size-fits-all system, it uses a flexible system that changes based on the words themselves. The authors figured out how to make this flexible system run fast enough to be practical, resulting in AI models that are better at logic, tracking information over long distances, and understanding complex sequences.

Technical Summary: PaTH Attention

1. Problem Statement

The attention mechanism is fundamental to modern Large Language Models (LLMs), yet it is inherently permutation-invariant. To model structured domains like language, position encoding is required. While Rotary Position Embedding (RoPE) has become the de facto standard, it possesses a critical limitation: the transformation applied to keys and queries is data-independent. In RoPE, the interaction between two tokens depends solely on their relative distance, not on the content of the tokens themselves.

This data-independence restricts the expressivity of RoPE-based transformers. Theoretical analysis indicates that RoPE-based models remain within the $TC^0$ complexity class (equivalent to ordinary transformers with absolute position embeddings), limiting their ability to solve tasks requiring sequential reasoning or state tracking (e.g., flip-flop language modeling, state-tracking problems). Consequently, these models struggle with synthetic tasks that demand dynamic adaptation to input sequences.

2. Methodology: PaTH Attention

The authors propose PaTH (Position Encoding via Accumulating Householder Transformations), a flexible, data-dependent position encoding scheme.

Core Mechanism

Unlike RoPE, which applies a fixed rotation matrix, PaTH computes attention logits using a cumulative product of data-dependent Householder-like transformations along the path between query and key positions.

Attention Logit Formulation: The attention score between query $q_i$ and key $k_j$ is defined as:
$A_{ij} \propto \exp\left(k_j^\top \left(\prod_{s=j+1}^{i} H_s\right) q_i\right)$
Data-Dependent Transition Matrix ( $H_s$ ): Each matrix $H_s$ in the sequence is a Householder-like identity-plus-rank-one matrix:
$H_s = I - \beta_s w_s w_s^\top$
Here, $w_s$ and $\beta_s$ are functions of the input token $x_s$ . Specifically, $\beta_s = 2 \times \text{sigmoid}(u^\top x_s + b)$ , allowing the transition matrix to have negative eigenvalues, which is crucial for state tracking performance.

Theoretical Expressivity

The paper establishes that PaTH extends the theoretical capabilities of transformers.

Complexity Class: A one-layer PaTH transformer with two attention heads can solve an $NC^1$ -complete problem under $AC^0$ reductions. This demonstrates that PaTH can break the $TC^0$ barrier that constrains RoPE-based models.
Relation to Linear RNNs: PaTH is shown to be closely related to expressive linear RNNs (like DeltaNet), inheriting their ability to track states while retaining the associative recall benefits of softmax attention.

Extensions: PaTH-FoX

The authors demonstrate that PaTH is compatible with other attention modifications. They propose PaTH-FoX, which combines PaTH's multiplicative position encoding with the Forgetting Transformer (FoX)'s additive, data-dependent forget gates. This hybrid approach further enhances performance, particularly in length extrapolation.

3. Efficient Training and Inference

A major challenge in PaTH is computing the cumulative product of Householder matrices efficiently without materializing the full $L \times L$ attention matrix.

UT Transform Representation: The authors utilize the UT transform to represent the product of Householder matrices compactly. A product of $L$ transformations can be expressed as $P = I - W^\top T^{-1} W$ , where $W$ contains the vectors $w_s$ and $T$ is a triangular matrix derived from the $\beta$ values.
FlashAttention-Style Algorithm: To enable hardware-efficient training, the authors develop a blockwise algorithm:
- They pre-compute "boundary-adjusted" query and key matrices for each block.
- They use a right-to-left scan over key/value blocks to accumulate the cumulative Householder products in a streaming manner.
- This approach maintains the I/O efficiency of FlashAttention (operating within SRAM) while incorporating dynamic positional encoding.
Inference: For inference, historical keys are updated in-place using the current timestep's transition matrix, eliminating the need for a separate cache of transformation vectors. This allows compatibility with existing inference kernels like FlashDecoding.

4. Experimental Results

Synthetic Benchmarks

PaTH was evaluated on tasks designed to probe sequential reasoning and state tracking:

Flip-Flop Language Modeling (FFLM): PaTH-based models (1 layer, 2 heads) achieved near-perfect accuracy on this task, whereas RoPE-based models struggled significantly.
Word Problems ( $A_5$ Group): PaTH solved $NC^1$ -complete word problems with fewer layers than baselines (RoPE, SBA, FoX).
Multi-query Repeated Associative Recall (MQRAR-N): PaTH successfully tracked variable values for $N$ -back recall ( $N < 4$ ), outperforming recent baselines that failed to track ordered histories.

Language Modeling (760M Parameters)

Experiments were conducted on the Fineweb-Edu corpus (50B tokens):

Standard Benchmarks: PaTH consistently outperformed RoPE and FoX on zero-shot commonsense reasoning tasks (LAMBADA, PIQA, HellaSwag, etc.) and achieved lower perplexity on Wikitext.
Length Extrapolation: Models trained on 4K context generalized well to longer sequences.
- RoPE: Failed abruptly beyond 4K.
- PaTH: Maintained stable performance up to 32K.
- PaTH-FoX: Generalized robustly up to 64K tokens, showing particular strength in code generation (CodeParrot) where state tracking is critical.
Long-Context Benchmarks: On RULER, BABILONG, and PhoneBook, PaTH-FoX achieved the highest retrieval and state-tracking scores, significantly outperforming RoPE and FoX in Multi-Needle-In-A-Haystack and variable tracking tasks.

Model Conversion

The authors explored converting pretrained RoPE models to PaTH:

SmolLM2-1.7B: Converting a model before it was exposed to high-quality math/code data and then continuing pretraining resulted in PaTH outperforming the original RoPE model on math (GSM-8K) and coding (HumanEval, MBPP) benchmarks.
Qwen2.5-7B: Distilling from an already overtrained RoPE model yielded mixed results, with some performance gains but also degradation on certain benchmarks, suggesting that conversion is most effective before the model "ossifies."

5. Significance and Claims

The paper claims that PaTH addresses the theoretical and empirical limitations of current position encoding schemes by introducing data-dependent multiplicative transformations.

Expressivity: It provides a principled way to move transformers beyond the $TC^0$ complexity class, enabling them to solve $NC^1$ -complete problems.
Practicality: Through the UT transform and blockwise algorithms, PaTH achieves efficiency comparable to standard attention mechanisms, making it viable for large-scale training.
State Tracking: The data-dependent nature of the Householder transformations allows the model to dynamically adapt to input sequences, solving state-tracking tasks that are intractable for RoPE-based models.

6. Limitations

The authors acknowledge several caveats:

Numerical Stability: Training can be sensitive to precision; specifically, BF16 may cause instability in the cumulative product, requiring clipping of the scaling factor $\beta$ .
Compute Overhead: While efficient, larger head dimensions increase the computational and memory overhead of the Householder product calculations compared to standard attention.
Geometric Biases: PaTH does not natively model rotations (as Householder matrices are reflections). This may limit certain geometric inductive biases present in RoPE, such as structured dependence on relative position, though the authors suggest future work could approximate rotations via compositions of reflections.