REVIEW 4 major objections 6 minor 2 cited by
This paper claims that rewriting the Kalman filter in information form makes its nonlinear per-token updates parallelizable, giving a sequence mixer with explicit uncertainty that outperforms linear state-space and gated-linear-attention la
Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →
KLA reparameterises information-form Kalman filtering as an associative Möbius scan, giving a parallelisable nonlinear state-space mixer that tracks A5 permutations with one or two layers.
T0 review reviewed 2026-08-03 challenge →
load-bearing objection The Möbius-scan theory is sound and the primitive is genuinely interesting, but the A5 evidence as presented is too weak to carry the expressivity claim, and the abstract oversells what the body delivers. the 4 major comments →
Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
Core claim
The central claim is that Kalman filtering can serve as a drop-in sequence-mixing primitive for language modelling. In a diagonal linear-Gaussian state-space model, the posterior precision evolves by a Möbius transformation, so the sequence of precision updates is equivalent to a chain of 2x2 matrices under multiplication; the posterior information mean evolves affinely with a forget gate set by the precision ratio. These two recursions can be computed by associative parallel scans in O(T) work and O(log T) depth. The nonlinearity—the fractional-linear denominator—is exactly what makes the gate history-dependent and strictly more expressive than linear/affine gating, while the explicit preci
What carries the argument
The Möbius (fractional-linear) transform of the precision recursion: λ_t = M_t(λ_{t−1}) = (α_t λ_{t−1} + β_t)/(γ_t λ_{t−1} + δ_t), with M_t = [[1 + p_t ϕ_t, a_t^2 ϕ_t],[p_t, a_t^2]], where ϕ_t is token-evidence precision weighted by the observation operator. Because composing Möbius transforms is 2x2 matrix multiplication, the entire precision path is a parallel prefix scan; the affine information-mean recursion runs as a second scan. This object carries the argument: it is the point where nonlinear Bayesian filtering meets associative scan parallelism and where the uncertainty-ratio gate emerges.
Load-bearing premise
The argument assumes a diagonal, per-coordinate state model: every feature channel runs an independent scalar Kalman filter, so any advantage rests on tasks where the needed information lives within coordinates rather than in correlations between them.
What would settle it
Train a single-layer KLA on a task whose answer requires combining two different state coordinates, such as adding two numbers encoded in separate channels or a selective-copy variant where the token to copy is the XOR of two channel values. If accuracy stays at parity with a full-covariance Bayesian filter, the diagonal assumption is benign; if it collapses while a transformer solves it, the central expressivity claim is bounded to per-coordinate state tracking.
If this is right
- KLA is a drop-in mixer with O(T) work and O(log T) depth, matching the training profile of linear-time state-space and gated-attention layers.
- The precision-ratio forget gate is history-dependent and nonlinear, so the model can selectively down-weight new evidence as accumulated confidence grows—mechanistically similar to gating but derived from inference.
- KLA solves the A5 permutation-composition task with 1-2 layers, while linear state-space models and transformers need depth that grows with sequence length.
- On long-context associative recall at sequence length 2048, KLA exceeds baseline accuracy and reaches near-perfect performance at model dimension 256.
- Uncertainty-weighted updates reduce state saturation under heavy key-value load, and Monte Carlo decoding under the learned posterior gives further gains on noise-corrupted synthetic tasks.
Where Pith is reading between the lines
- Editorial: since all recursions are per-coordinate, the layer is effectively a bank of independent scalar filters; tasks that require binding or correlating across feature coordinates would likely need a full-covariance extension, which this paper does not provide.
- Editorial: the explicit posterior variance is a naturally interpretable uncertainty signal that could serve hallucination detection, abstention, or out-of-distribution prompting; the authors note these as untested future work.
- Editorial: the 2x2-matrix Möbius associativity is a general transferable mechanism—any nested fractional-linear recursion, such as time-varying or controlled dynamics or Bayesian smoothing, could inherit the same parallel-scan trick.
- Editorial: a clean empirical test of whether the nonlinearity matters is to fix the precision path to an affine approximation and measure A5 accuracy; if the result persists, the extra expressivity has a different source than the fractional-linear denominator.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces Kalman Linear Attention (KLA), a sequence-mixing layer that reparameterises a diagonal linear-Gaussian Kalman filter in information form. The central theoretical claim is that the posterior precision recursion is a Möbius (fractional-linear) map and the information-mean recursion is affine, so both update families compose associatively and can be evaluated with parallel prefix scans in O(T) work and O(log T) depth, matching Mamba/GLA-style efficiency while retaining a nonlinear, uncertainty-driven gating mechanism. The paper also claims empirical gains on synthetic language-modelling tasks (MAD suite), long-context associative recall (MQAR), and the A5 permutation-composition state-tracking benchmark, and states in the abstract that KLA is among the first stacked Bayesian-filtering primitives trained at billion-token scale.
Significance. If the central claims hold, KLA is a conceptually appealing and practically relevant primitive: it offers a nonlinear recurrence that remains scan-parallel, explicit belief-state uncertainty, and a principled probabilistic interpretation of gating. The theoretical core (Theorems 1 and 2, Corollaries 1 and 2) is self-contained and appears correct: the derivation of the Möbius precision update from standard information-filter equations is clean, and the scan parallelisation follows directly from associativity of matrix multiplication. The paper also provides a useful comparison with Särkkä and García-Fernández's lifting construction, showing that no 5-tuple augmentation is needed in this case. However, the empirical evidence for the headline expressivity claim—that KLA 'solves' A5 tasks that linear SSMs and attention cannot—is currently weak, and the manuscript's own limitations section disclaims the web-scale claims made in the abstract. The central theoretical contribution is sound, but the empirical demonstration supporting 'strictly more expressive' requires substantial strengthening before the paper's claims are fully credible.
major comments (4)
- [§5.5, Fig. 1, Appendix G.5] The A5 state-tracking result is the only direct evidence for the paper's central expressivity claim, but as reported it cannot be independently verified. Success is defined as '≥90% accuracy in at least one of five seeds' (Appendix G.5), a criterion that can be satisfied by a single lucky run and is not a standard measure of reliable performance. Figure 1 has no numerical table, no error bars, and no per-seed results. Baselines are taken from Merrill et al. (2024) rather than trained under the same protocol, parameter counts, data, and early-stopping. Given that the entire 'strictly more expressive' claim rests on this experiment, a matched re-evaluation under identical training conditions with all five seeds reported (mean, std, min/max) is required before the claim is supportable.
- [Section 3 'Notational conventions'; Theorem 1 proof (Appendix D)] The derivation assumes a diagonal/per-coordinate model: 'we identify diagonal matrices with their diagonal terms and apply scalar recursions elementwise.' Consequently, the KLA layer is a collection of independent scalar filters per coordinate, and the Möbius associativity argument does not extend to cross-coordinate interactions. The A5 permutation-composition task is non-abelian and may require cross-coordinate mixing in the state update, yet the paper does not test or justify whether diagonal structure suffices. This compounds the fragility of the A5 evidence: the empirical claim is doing critical work that neither theory nor robust experiment currently supports. Please add an explicit discussion and, ideally, an ablation with non-diagonal/matrix-state updates or a concrete argument for why per-coordinate independence still captures the required composition.
- [Abstract vs. Section 7 (Limitations)] The abstract states that KLA is 'among the first stacked Bayesian-filtering primitives trained at the billion-token scale,' but Section 7 explicitly says 'we do not claim state-of-the-art performance on web-scale corpora' and the body contains no billion-token experiments. This is an internal inconsistency between the abstract and the manuscript's own limitations. The abstract should be tempered to match the actual empirical scope, or the billion-token experiments should be included.
- [§5.3, Table 3, Appendix G.2] Table 3 reports a single accuracy number per task and per method, with no error bars, despite Appendix G.1 stating results are averaged over 5 seeds. Moreover, KLA+ uses 10 Monte Carlo samples for probabilistic decoding (Appendix G.2), which is a different inference/training procedure than the baselines; comparing KLA+ against standard cross-entropy baselines may conflate the benefit of the Bayesian primitive with the benefit of marginalisation. Please report per-seed variance for all methods and clearly separate 'KLA with posterior-mean decoding' from 'KLA with Monte Carlo marginalisation', and ensure baselines receive the same decoding treatment where applicable.
minor comments (6)
- [§5.3] 'Section 5.3 shows that KLA performs competitively...' is a self-referential phrase; it should say 'Table 3 shows...'.
- [Appendix A] The text mentions QK-Norm and expansion factor 1, but this is not described in the main method or hyperparameter tables. Please add a sentence explaining where QK-Norm is applied and why.
- [§5.2, Fig. 5] The runtime benchmark compares KLA implementations against each other but not against Mamba/GLA on the same hardware. Since the paper claims 'same computational cost' as those models, a direct wall-clock comparison with at least one strong baseline would make the claim concrete.
- [References] The workshop name 'EurIPS 2025' appears in the footnote; this is likely a typo for NeurIPS. Please verify.
- [Appendix E.3/E.4] The 'Kalman Attention Matrix' visualisations are qualitative. A quantitative comparison with the attention matrices of GLA/Mamba on the same tasks would strengthen the claimed interpretability.
- [Theorem 3 / Appendix D] The convolutional form for deterministic LTI systems is derived correctly, but the statement 'computable in O(T log T)' and 'O(log T) parallel time' conflates serial and parallel complexity. Please clarify that FFT gives O(T log T) serial work and O(log T) parallel depth.
Circularity Check
No significant circularity: the Kalman/information-filter derivation is self-contained; empirical caveats exist but do not reduce to construction.
full rationale
The paper's central derivation is mathematically self-contained. Theorem 1 derives the Möbius precision recursion directly from the standard information-form Kalman update (predict: λ_prior = λ_{t-1}/(a^2 + p λ_{t-1}), update: λ_t = λ_prior + k^2 Λ^v_t), and Theorem 2 derives the affine mean update from the same equations. These are not fitted to any target result; they are algebraic consequences of the stated diagonal linear-Gaussian model. The associativity enabling parallel scans follows from the 2×2 matrix representation of Möbius transformations, an external mathematical fact. No fitted parameter is later called a prediction; the MAD, MQAR, and A5 results are empirical evaluations, not derived from the trained parameters in a way that forces the outcome. Self-citations (Shaj et al. 2021, Shaj Kumar et al. 2023) appear only in related work on deep Kalman filters and are not load-bearing: the key theorems and proofs do not rely on them. The claim of 'strictly more expressive' nonlinear updates rests on the algebraic form (nonlinear vs. linear recurrence) plus the A5 experiment; while the experiment itself has validity concerns—Appendix G.5 defines success as ≥90% accuracy in at least one of five seeds, and baselines are copied from Merrill et al. rather than re-run under the same protocol—this is an empirical robustness issue, not a construction-level circularity. The diagonal (per-coordinate) assumption is explicitly stated and limits the theory's generality, but again does not make the derivation circular. Neither the abstract's unsupported 'billion-token scale' claim nor the discrepancy with Section 7's disclaimer are circularity; they are consistency/support issues. Overall, the derivation chain is not equivalent to its inputs, and no step reduces to its own conclusion.
Axiom & Free-Parameter Ledger
free parameters (6)
- process_noise_scale_p =
learned; initial 0.01
- OU_decay_a =
learned, time-invariant
- discretization_timestep_delta =
learned in [0.001, 0.1]
- value_precision_lambda_v =
learned from input
- observation/readout_operators_k_q =
learned from input
- state_expansion_N =
8 (MAD) / 16 (MQAR, A5)
axioms (7)
- standard math Standard Kalman/information-filter recursions for linear-Gaussian state-space models.
- standard math Möbius transformations compose through 2x2 matrix multiplication, and matrix multiplication is associative.
- standard math Parallel prefix (associative) scan has O(T) work and O(log T) depth.
- domain assumption Diagonal parameterisation is sufficient for language-model sequence mixing.
- domain assumption OU process is an appropriate continuous-time prior for the latent state.
- domain assumption A5 permutation composition is a valid proxy for sequential state tracking and distinguishes TC0 from NC1.
- domain assumption Observations vt are generated as noisy linear measurements of the latent state; training uses discriminative cross-entropy, so the generative interpretation is metaphorical.
Cite this review
Pith. "Pith review of Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking." pith.science (2026). https://pith.science/paper/C6YVXQIK
@misc{pith2026260210743,
author = {Pith},
title = {Pith review of: Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking},
year = {2026},
howpublished = {\url{https://pith.science/paper/C6YVXQIK}},
note = {Machine review of arXiv:2602.10743}
}
read the original abstract
State-space language models such as Mamba and gated linear attention (GLA) offer linear-complexity, parallelisable alternatives to transformers, but their linear state updates limit expressivity and robust state tracking. We close this gap from a probabilistic angle, casting sequence mixing as exact Bayesian filtering with the Kalman filter as the core primitive. Classical Kalman filters give principled state and uncertainty estimates but are viewed as inherently sequential; we show that reparameterising them in information form turns their updates into an associative scan - so the per-token recurrent update is non-linear (a M\"obius/precision recursion) yet remains temporally parallel. The resulting Kalman Linear Attention (KLA) layer is a drop-in sequence mixer that performs time-parallel probabilistic inference, carries an explicit belief-state uncertainty, and is strictly more expressive than GLA-style linear updates at the same computational cost. This expressivity translates directly into stronger state tracking: KLA solves permutation-composition ($A_5$) tasks that linear SSMs and attention cannot, while staying scan-parallel. As a drop-in primitive it also matches or improves on modern SSMs and GLAs across synthetic token-manipulation and zero-shot commonsense benchmarks, and is among the first stacked Bayesian-filtering primitives trained at the billion-token scale.
Figures
Forward citations
Cited by 2 Pith papers
-
Memory by Design: Probabilistic Sequence Layers
The design-model framework unifies sub-quadratic sequence models as Bayesian filters and introduces a covariance-tracking Bayesian Layer that improves retrieval robustness beyond training regimes on MQAR and RULER benchmarks.
-
Memory for Large Language Models
The paper organizes LLM memory architectures along three axes—implicit vs explicit representation, offline vs online update, and short vs long-term persistence—plus finer-grained update rules.
Reference graph
Works this paper leans on
-
[3]
More extensive exploration of scaffolding configurations optimally suited to the KLA filtering mixer can be a direction for future work
and is set to 16 (N or sometimes called d_state in SSM codebases) as default similar to (Gu & Dao, 2023). More extensive exploration of scaffolding configurations optimally suited to the KLA filtering mixer can be a direction for future work. We implemented KLA with PyTorch (Paszke et al.,
2023
-
[8]
Nielsen, F. and Garcia, V . Statistical exponential families: A digest with flash cards.arXiv preprint arXiv:0911.4863,
-
[9]
Yang, S., Wang, B., Shen, Y ., Panda, R., and Kim, Y . Gated linear attention transformers with hardware-efficient train- ing.arXiv preprint arXiv:2312.06635,
-
[11]
Yu, A. and Erichson, N. B. Block-biased mamba for long-range sequence processing.arXiv preprint arXiv:2505.09022,
-
[13]
Each transform is represented (up to multiplication by a nonzero scalar) by a 2×2 matrix M= a b c d and acts on a scalar z via the linear–fractional rule M(z) := az+b cz+d
It generalises the familiar affine form (c= 0 gives az+b ) by introducing an additional denominator term. Each transform is represented (up to multiplication by a nonzero scalar) by a 2×2 matrix M= a b c d and acts on a scalar z via the linear–fractional rule M(z) := az+b cz+d . Composing multiple transforms amounts to multiplying their representing matri...
2009
-
[14]
pointer-like
Figure 10.Process noise ablation on MADLAB tasks.Comparison of KLA with learnable process noise versus a variant with process noise fixed to zero (pt = 0). Removing process noise leads to severe performance degradation, with an average accuracy drop of 49.8 percentage points. E.2. Runtime Scaling Setup.We benchmark forward-only runtimes for three KLA impl...
2048
-
[16]
Model # Blocksd model dstate OtherS KLA (Ours) 1 128 8 expand=1 2048 Mamba 1 128 16 expand=1 2048 GatedDeltaNet 1 128 -e k=0.5,e v=1.0,H=4 2048 GLA 1 128 -e k=0.5,e v=1.0,H=4 2048 mLSTM 1 128 -H=16, proj_factor=2.0 2048 Batch size:We use a batch size of 172 for all MAD-Lab experiments. G.4. MQAR Hyperparameters For the MQAR (hard) experiments, we evaluate...
2048
-
[17]
A5 State Tracking Hyperparameters For the A5 state tracking experiments (Section 5.5), we use dmodel = 1024 and dstate = 16 for KLA
KLA 2 256 16 expand=1 8,192 Mamba 2 256 16 expand=2 8,192 GatedDeltaNet 2 256 -e k=0.75,e v=1.5,H=8 9,216 GLA 2 256 -e k=0.25,e v=1.0,H=4 8,192 G.5. A5 State Tracking Hyperparameters For the A5 state tracking experiments (Section 5.5), we use dmodel = 1024 and dstate = 16 for KLA. We train for up to 500 epochs with early stopping (patience 50), and a lear...
2024
-
[1957]
10 Parallel, Scalable and Efficient Bayesian Filters for Language Modelling Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality.arXiv preprint arXiv:2405.21060,
-
[2005]
Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., and Ré, C. Zoology: Measuring and improving recall in efficient language models.arXiv preprint arXiv:2312.04927,
-
[2015]
Krishnan, R. G., Shalit, U., and Sontag, D. Deep Kalman filters.arXiv preprint arXiv:1511.05121,
-
[2016]
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces.arXiv preprint arXiv:2312.00752,
-
[2017]
Luis, C. E., Bottero, A. G., Vinogradska, J., Berkenkamp, F., and Peters, J. Uncertainty representations in state- space layers for deep reinforcement learning under partial observability.arXiv preprint arXiv:2409.16824,
-
[2019]
KalMamba: Towards efficient probabilistic state space models for RL under uncertainty
Becker, P., Freymuth, N., and Neumann, G. KalMamba: Towards efficient probabilistic state space models for RL under uncertainty. InICML 2024 Workshop: Aligning Reinforcement Learning Experimentalists and Theorists,
2024
-
[2023]
Gated delta net- works: Improving Mamba2 with delta rule.arXiv preprint arXiv:2412.06464,
Yang, S., Kautz, J., and Hatamizadeh, A. Gated delta net- works: Improving Mamba2 with delta rule.arXiv preprint arXiv:2412.06464,
-
[2024]
The illusion of state in state-space models.arXiv preprint arXiv:2404.08819,
Merrill, W., Petty, J., and Sabharwal, A. The illusion of state in state-space models.arXiv preprint arXiv:2404.08819,
-
[2048]
Long-Context MQAR Table 7 lists the data configuration for the long-context MQAR benchmark task (Arora et al., 2023)
F.2. Long-Context MQAR Table 7 lists the data configuration for the long-context MQAR benchmark task (Arora et al., 2023). Table 7.Long-Context MQAR data parameters. Long-Context Multi-Query Associative Recall (MQAR) Setting Seq. length Vocab size Training seqs Eval seqs CR (hard) 2,048 256 12,800 1,280 27 Parallel, Scalable and Efficient Bayesian Filters...
2023
This paper was first reviewed by deepseek-v4-flash on August 3, 2026.
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.