We thank the reviewer for their constructive feedback.

Insufficient Exploration of the H Matrix Design

We appreciate the suggestion to clarify our design choices.

Empirically, we found it necessary to set $\beta \in (0, 2)$. Fixing $\beta = 2$ enforces norm preservation but causes severe numerical instability under BF16, degrading performance and training stability. Despite significant effort, we were unable to stabilize this setting. Using $\beta \in (0, 1)$, as in DeltaNet, led to a loss of state-tracking, consistent with findings in “Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues.”

We also found that sharing projection parameters between $k$ and $w$ (as in DeltaNet and Gated DeltaNet) significantly hurts performance. Unsharing them was beneficial. We plan to include a full ablation study in the revision to support these findings.

Motivation for Using Householder Matrices

Our primary goal is to improve state tracking, especially for problems like the S5 permutation task, which is NC$^1$-complete under AC$^0$ reductions. Solving S5 implies the model can simulate a wide class of state-tracking computations.

This motivates us to look for transformations that naturally model compositions of swaps. Householder matrices are appealing in this context: each matrix reflects a vector across a hyperplane and can mimic a pairwise swap under certain conditions, while preserving other coordinates.

By cumulatively composing these Householder matrices (our PaTH scheme), we simulate multiple such swaps—aligning structurally with the S5 group.

This motivation is formalized in Theorem 2.1 (Appendix C), though currently only briefly mentioned in lines 625–626. We agree this should be made clearer earlier in the paper and will revise the introduction and Section 2 accordingly.

Efficiency Evaluation

Triton Kernel Latency Comparison (Forward + Backward, ms)

| Batch | SeqLen | HeadDim | Parallel Attn | PaTH Attn | Forgetting Attn |
|-------|--------|---------|----------------|-----------|-----------------|
| 16    | 2048   | 128     | 5.07           | 17.59     | 8.79            |
|       |        | 64      | 6.55           | 11.95     | 8.54            |
| 8     | 4096   | 128     | 8.47           | 21.88     | 11.16           |
|       |        | 64      | 11.01          | 16.64     | 14.50           |
| 4     | 8192   | 128     | 15.28          | 30.90     | 20.16           |
|       |        | 64      | 19.95          | 26.06     | 26.45           |
| 2     | 16384  | 128     | 28.76          | 48.49     | 37.92           |
|       |        | 64      | 37.64          | 43.89     | 50.05           |

Since FFN dominates training time in practice (esp. in MOE training), the end-to-end training speed difference will be smaller than the kernel-level gap shown here. We are still actively optimizing the kernel for headdim128.

Inference Latency

Due to time constraints, we did not implement inference-time optimization. We note that decoding latency is highly implementation-dependent, and efficient fused implementations are possible.

In particular, the rank-1 refinement on past key vectors involves a fixed pattern of memory access, making it amenable to a fused kernel that performs Flash-decoding with integrated key refinement. The main overhead is writing updated keys back to HBM, which we believe can be pipelined efficiently. We plan to explore inference latency more thoroughly in future work.

Model Architecture Details

Thank you for highlighting this. We apologize for the overly brief description. In our experiments, we use the standard LLaMA architecture, rather than the modified “Pro” version in FoX. This was an intentional choice: we aimed to minimize architectural changes, so the effect of replacing RoPE with PaTH can be isolated and fairly compared. All our baselines—RoPE, FoX, and PaTH—use this consistent setting.

The only additional components used by PaTH are those needed to compute w and β, following “Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues”: Linear projection → Short convolution → L2 normalization (per head).

To reduce parameter overhead and ensure fair comparison, we apply low-rank projections, similar in spirit to LoRA.

We will make these architectural choices clearer in the main text.

Q: Why report acc for ARC-e but acc_norm for ARC-c?

We follow the evaluation setup from Mamba (Section E.2.3 of the Mamba paper), which reports acc for ARC-e and acc_norm for ARC-c.

Q4: How does the softmax affect performance?

While linear RNNs are efficient, they struggle with retrieval and contextual selectivity. On RULER, removing softmax (reducing to DeltaNet) significantly hurts performance. We will clarify this in the revision and plan to include retrieval comparisons in future work.

Thank you for engaging!

I still don't understand the motivation for improving state tracking in a strong attention model.

State tracking remains a challenging and unresolved issue for RoPE-based attention models. A recent benchmark, LoCoDiff: Natural Long Context Code Bench, evaluates a key capability for coding agents—tracking the state of edited files—which is a prototypical state tracking task. Performance degrades sharply as context length increases: while some models reach near 100% accuracy on prompts under 5k tokens, accuracy drops significantly by 10k tokens and falls below 50% at 25k tokens. This highlights a fundamental limitation of RoPE-based attention: the number of layers required for effective state tracking grows logarithmically with sequence length, yet the number of layers in practice is fixed.

In contrast, NC¹-complete (under AC⁰ reductions) architectures such as PaTH and DeltaFormer may be capable of solving these state tracking tasks without the need for deeper networks as the input length grows. At a higher level, PaTH may also be better suited for chain-of-thought reasoning tasks, offering improved token efficiency due to its greater expressive power.

While we currently lack empirical results on real-world code tasks, this work focuses on synthetic tasks specifically designed to test state tracking. We view these results as a promising starting point, and acknowledge that extending PaTH to practical scenarios like test-time scaling and code modeling will require significant future work—likely the subject of a follow-up paper.

Efficiency

We acknowledge the current limitation of PaTH attention efficiency—its kernel is approximately 2× slower than standard attention (note: this refers to the attention kernel itself, not end-to-end training performance). In future work, we plan to explore partial PaTH—applying PaTH to only a subset of the attention dimensions while using NoPE for the rest—similar in spirit to partial RoPE or Multi-Latent Attention, to improve speed without sacrificing too much expressivity. We also aim to rewrite the kernel using more advanced GPU programming languages such as TileLang, since Triton’s current memory management is suboptimal for PaTH, which requires substantial SRAM to efficiently support transition matrices.

Regarding inference latency, we aim to implement PaTH in vLLM for a fair comparison. However, this would require several weeks of engineering effort, so we are unable to provide numbers at this time. We believe that inference speed comparisons using Hugging Face are somewhat unreliable, as HF is significantly slower and inference performance is highly implementation-dependent. That said, if you think an HF-based speed comparison would still be useful, we’re happy to run it and report the results. We also hope that our discussion in Section 3.3 partially addresses your concerns about inference efficiency. Also, as mentioned earlier, partial PaTH appears to be a promising direction for reducing inference complexity further.

Despite this efficiency limitation, we believe it’s important to explore more expressive—even if more expensive—attention mechanisms. While many recent works have focused on hardware-efficient linear attention for better throughput, they often fall short in expressivity and downstream task accuracy. As compute continues to become cheaper, we risk running out of data and model variants that meaningfully push the frontier. Our goal is not to make PaTH faster than RoPE, but to propose a practical, blockwise algorithm that makes PaTH trainable on modern hardware—enabling empirical investigation of its unique properties. With the current kernel, for example, we are able to fine-tune a 7B model with head dimension 128 without difficulty (see our response to Reviewer V5VP for results).

We remain optimistic about the direction of developing more expressive and powerful attention mechanisms that remain hardware-aligned—ensuring that wall-clock performance stays affordable even as theoretical complexity grows.

Given that PaTH is using the additional components, including short convolution and L2 normalization

We did not modify the QKV computation—it remains identical to what is used in LLaMA with RoPE attention. The L2 normalization of $w$ is necessary to ensure the transition matrix has a spectral norm less than 1, which is important for stability. As for the short convolution, it does not have a significant impact on performance. In fact, in our distillation experiment (see our response to Reviewer V5VP), we excluded the short convolution module entirely, and the model still performed well. We plan to clarify these architectural details and conduct ablation studies to make this more transparent.

Thank you for the constructive feedback.

Implementation detail

We aim to improve the clarity of the algorithm flow in the next iteration. For now, we provide a PyTorch reference implementation for illustration. (For simplicity, we omit the online softmax trick and explicitly construct the attention matrix A; in our actual Triton implementation, A is never materialized and online softmax is used.) Inputs are shaped [B, H, L, D] for q, k, v, w, and [B, H, L] for beta.

import torch
from einops import rearrange
def naive_path_attn(q, k, v, w, beta, scale, BT=64):
    b, h, l, d = q.shape 
    wb = w * beta[..., None]
    q, k, w, wb = [rearrange(x, 'b h (n c) d -> b h n c d', c=BT) for x in (q, k, w, wb)]
    m = torch.triu(torch.ones(BT, BT, dtype=torch.bool))
    T = -(wb @ w.transpose(-1, -2)).masked_fill(m, 0)
    # forward substituition
    for i in range(1, BT):
        T[..., i, :i] += (T[..., i, :, None] * T[..., :, :i]).sum(-2)
    T = T + torch.eye(BT)
    Twbk = T @ (wb @ k.transpose(-1, -2)).masked_fill(m, 0)
    qw = (q @ w.transpose(-1, -2)).tril()
    Twb = T @ wb
    Al = (q @ k.transpose(-1, -2)).tril() - qw @ Twbk
    q = q - qw @ Twb
    k = k - Twbk.transpose(-1, -2) @ w
    H = w.transpose(-1, -2) @ Twb
    q, k = [rearrange(x, 'b h n c d -> b h (n c) d') for x in (q, k)]
    A = torch.zeros(b, h, l, l)
    # FlashAttention-2 style blocking
    for i in range(0, l, BT):
        qi = q[:, :, i:i+BT]
        for j in range(i-BT, -1, -BT):
            kj = k[:, :, j:j+BT]
            A[:, :, i:i+BT, j:j+BT] = qi @ kj.transpose(-1, -2)
            qi = qi - qi @ H[:, :, j//BT]
    for i in range(l//BT):
        A[:, :, i*BT:(i+1)*BT, i*BT:(i+1)*BT] = Al[:, :, i]
    A = A.masked_fill_(~torch.tril(torch.ones(l, l, dtype=torch.bool)), float('-inf'))
    return (A * scale).softmax(-1) @ v

Regarding NoPE

Thank you for the concern. We view NoPE as orthogonal to our work, and also length extrapolation is not our primary focus but rather a byproduct of our data-dependent encoding. We didn’t compare with NoPE due to time constraints. Moreover, we want to highlight that NoPE shares RoPE’s theoretical limitations (both in TC⁰), while PaTH goes beyond TC⁰ to support parallelizable state tracking.

Regarding FFLM

Regarding the discrepancy between theory and experiment: this stems from the fact that, in our implementation, we apply a short convolution layer after w_proj to compute w. In the theoretical construction, however, w is used directly without convolution, requiring the first layer to perform explicit token shifting—hence, two layers are needed. With short convolution, the token shift is effectively handled by the convolution itself, so theoretically only one layer is sufficient to solve FLIP-FLOP, consistent with our experimental results. This is analogous to the case of induction heads: a single layer cannot implement the induction mechanism, but two layers can. However, if the short convolution is applied to the key/value projections, a single layer can implement the induction head, as shown in [1, 2].

We agree it’s unfair to compare shortconv-enhanced PaTH to others—thank you for pointing this out. As suggested, we report results across layer counts for fairness (best from LR ∈ {3e-4, 6e-4, 1e-3}):

Small OOD errors persist for RoPE/FoX even when stacking multiple layers (see FFLM paper appendix for theoretical insight) . Though stick-Breaking attention solves FFLM with 2 layers, it underperforms PaTH in A5 word problem and MQRAR-N tasks. Overall state tracking ability: RoPE ~= FoX < SBA < PaTH.

[1] KV Shifting Attention Enhances Language Modeling

[2] Physics of language model 4.1

Thanks for the review. Due to restrictions imposed by the NeurIPS committee, we are unable to provide new figures or code here. However, we commit to open-sourcing the highly optimized Triton implementation.

Replacing RoPE with PaTH from an Existing Checkpoint

Thank you for the insightful suggestion regarding distillation. We followed the Rapid Attention Distillation pipeline [1] to distill a teacher model using RoPE into a student model with PaTH attention.

We initialize PaTH using the teacher model’s weights, including all MLP and attention layer parameters. Since PaTH introduces an additional w_proj parameter, we initialize it with the weights from k_proj. Both k_proj and w_proj are then trained independently.

Following [1], we use the DCLM dataset for training. The distillation process involves two stages:

• Stage 1: Align the outputs of each attention layer between the teacher (RoPE) and student (PaTH) using MSE loss. Only the attention layers are trained, while MLP layers remain frozen. This stage is run for 100M tokens with a context length of 512.
• Stage 2: Perform standard knowledge distillation (using KL divergence on the vocab logits) with full finetuning across all layers. This stage is run for 3B tokens with a context length of 4K.

We summarize selected results below.

The teacher model is Qwen2.5-7B-Instruct. We use OpenCompass’s recommended settings for evaluating Qwen-Instruct series models (please refer to the corresponding markdown files in their repository for task-specific configurations).

Results using OpenCompass’s recommended configs:

We observe that the distilled student PaTH model largely recovers the performance of the teacher on these challenging tasks, and in some cases even outperforms it (e.g., GSM-8K, MBPP, TheoremQA, GPQA).

We believe PaTH has strong potential in coding, reasoning, and math tasks. As a next step, we plan to apply distillation followed by supervised fine-tuning (SFT) on coding-specialized models such as Qwen2.5-Coder and Qwen2.5-Math, using a mix of code and math data. Evaluating the reasoning capabilities of these distilled models is also of great interest, which we leave for future work.

We also evaluate on RULER using lm-evaluation-harness:

We find that PaTH largely recovers the teacher model’s performance at the training length of 4K. However, performance degrades when evaluating beyond this length, indicating that the student model does not fully retain the teacher’s long-context capability. This suggests the need for further sequence length extension experiments, which we leave for future work.

[1] RADLADS: Rapid Attention Distillation to Linear Attention Decoders at Scale

Thank you for raising the score! We plan to conduct more distillation and continual pretraining experiments over the next month and will make the weights publicly available.

Can you at least post the loss graph with tables?

We used wandb to log the loss and applied EMA smoothing with a coefficient of 0.6. Below is the reported loss at different training steps (We trained for 23,840 steps with a batch size of 2M tokens, amounting to approximately 50B tokens in total)

We observe a consistent trend: PaTH achieves ~0.02 lower loss than the RoPE baseline, while PaTH-FoX improves by ~0.04.

We thank the reviewer for their constructive feedback.

Re. Weaknesses

Practical implications of NC¹ vs. TC⁰

We believe NC¹ architectures may benefit tasks that require strong state tracking, such as code modeling & reasoning tasks. A weak empirical signal supporting this is shown in Figure 3, where PaTH achieves lower perplexity on the CodeParrot corpus. We plan to run more targeted code modeling experiments and explore distillation into coding models to further assess whether PaTH offers consistent advantages in this domain.

NoPE not included as a baseline

The primary goal of this work is to improve state tracking rather than length extrapolation. That said, we agree that NoPE is orthogonal and potentially complementary. For example, one could build a hybrid approach similar to MLA—using NoPE across most of the hidden dimensions and applying PaTH only to a subset. Exploring such partially structured designs is an interesting direction for future work.

Re. Questions

Q2: Why is data-dependency a benefit compared to additive positional embeddings?

We refer the reviewer to Section 5.2 of “The Illusion of State in State Space Models”, which discusses how data-dependent transition dynamics can significantly enhance expressivity. Without such data-dependence, PaTH cannot perform proper state tracking—which is the core motivation of this work.

Q4: Can PaTH be integrated into existing LLMs with post-training?

Yes—PaTH can be initialized from existing pretrained Transformer checkpoints and post-trained. We demonstrate this via a distillation experiment in our response to Reviewer V5VP, showing the feasibility of retrofitting PaTH onto pretrained models.

Q1, Q3, Q5: Attention patterns, lost-in-the-middle, and special tokens

We use a <bos> token at the beginning of each document but do not employ additional meta-tokens. We are interested in further analyzing PaTH’s attention behavior—e.g., whether it exhibits attention sinks or "lost-in-the-middle" effects—in future iterations.