REVIEW 3 major objections 5 minor 31 references
Kalman Bayesian Transformer
T0 review · 3 major / 5 minor · reviewed 2026-08-15 · deepseek-v4-flash
Pith's one-line read Sequential fine-tuning of a transformer can be run as a Kalman-style Bayesian update that stores one sample at a time and, in the paper's experiment, beats replay buffers of up to 50 samples.
desk verdict A fresh combination of Bayesian tools for transformer fine-tuning, but a likely zero-gain issue in the backward pass means the advertised full-head update may not be happening. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing identity is the recursive Bayes update $$p(W_B\mid D_k)\propto p(W_B\mid D_{k-1})\,p(X_k,Y_k\mid W_B),$$ which turns sequential fine-tuning into a one-sample-memory update. The machinery that makes this tractable is closed-form propagation of first and second moments through linear and ReLU layers, a first-order Taylor approximation of the softmax Jacobian to push moments through the final layer, and an RTS smoother in the backward pass that corrects the means and covariances of weights and intermediate activations layer by layer. The method also initializes the weight prior so that the pre-trained transformer's behavior is exactly reproduced before any data arrive.
What would settle it
On a small transformer where the true Bayesian posterior can be obtained by Monte Carlo sampling, run the same sequential fine-tuning task and compare the method's updated weight covariance after many samples with the covariance of samples from the true posterior; a large divergence in a regime with small hidden widths and strong weight-activation correlations would show that the independence and Gaussian assumptions are not benign.
Extended reading notes
Core claim
The central discovery is that the recursive Bayesian update $p(W_B|D_k) \propto p(W_B|D_{k-1}) p(X_k,Y_k|W_B)$ can be implemented in the moment domain for transformer output-head fine-tuning. Means and covariances of activations propagate in closed form through the linear and ReLU layers, softmax moments are approximated by a first-order Taylor expansion, and a Kalman smoothing backward pass updates the weight means and covariances one layer at a time. The paper reports that, on an inverted-pendulum decision-transformer task, this single-sample-memory sequential update outperforms warm-started retraining with replay buffers of 10, 20, 25, and 50 samples in stabilization success rate, does not show the success-rate drops that replay-based retraining exhibits, and produces uncertainty estimates that increase with data noise.
Load-bearing premise
The formulas assume each layer's weights and the previous layer's activations are independent and that pre-activations are Gaussian, so the closed-form ReLU moment propagation is exact only under those conditions; the paper itself notes that off-diagonal covariance terms have no closed-form backward propagation through nonlinear activations.
Editorial extensions
If this is right
- Fine-tuning can proceed online with a one-sample memory, eliminating replay buffers and their storage and management costs.
- Predictive uncertainty is available in closed form at every step, so a downstream controller or operator can know when the model's confidence is low.
- The per-sample update is a single closed-form pass rather than an iterative retraining loop, which suits latency-critical hardware.
- Because the prior at each step is the previous posterior, the method is designed to avoid catastrophic forgetting, and the experiment shows no forgetting-induced success-rate drops.
- The same recursive update naturally extends as new data arrive indefinitely, without needing to revisit old data.
Reading between the lines
- The demonstrated setting adapts only the linear output head; the moment-propagation cost of full-covariance updates would grow quickly if applied to attention or embedding parameters, so scaling to deeper adaptation needs further work.
- In regimes with strong correlations between weights and activations, the independence and Gaussian assumptions make the updated distribution an assumed-density filter rather than an exact posterior, so the uncertainty numbers should be read as calibrated heuristics until tested against sampling.
- A natural testable extension is to apply the same update to a language-model fine-tuning task with high-dimensional softmax outputs, where the first-order Taylor approximation may be the limiting component.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes Kalman Bayesian Transformer, a method for sequential fine-tuning of transformers under distribution shift and memory constraints. The method freezes the transformer body, replaces the output linear layer with a Bayesian neural network (BNN) initialized to reproduce the pre-trained model, and then updates the BNN weights one data sample at a time. The forward pass propagates means and covariances of activations through linear, ReLU, and softmax layers using closed-form moment propagation and a first-order Taylor approximation; the backward pass applies Rauch-Tung-Striebel smoothing to update weight distributions. The experimental evaluation on a decision transformer for inverted-pendulum stabilization reports higher stabilization success rate than warm-started retraining with memory buffers of 10-50 samples, lower per-sample computation time, and uncertainty estimates that grow with data noise.
Significance. The paper connects Kalman filtering/smoothing with transformer fine-tuning and provides a closed-form, single-pass alternative to sampling-based Bayesian methods. If the algorithm were correct as written, the contribution would be significant: memory-efficient sequential adaptation with explicit uncertainty quantification at low latency. The paper also gives a clear algorithmic description and a concrete experimental setup. However, the central derivation contains a load-bearing inconsistency in the backward pass that, unless resolved, means the method does not implement the full-head Bayesian update it claims. No machine-checked proofs are provided; the theoretical claims rest on moment-propagation formulas from prior work.
major comments (3)
- [Section IV-C and Algorithm 2, Eqs. (43)-(45)] The backward pass defines K_{z_{i-1}} = Σ_{z_{i-1},w_i} Σ^{-1}_{w_i,w_i}. Section IV-B explicitly assumes w_i and z_{i-1} are independent, which gives Σ_{z_{i-1},w_i}=0 and hence K_{z_{i-1}}=0. Then (44)-(45) leave µ_{z_{i-1}} and Σ_{z_{i-1},z_{i-1}} unchanged, so in the next lower layer the innovation µ^+_{z_{i-1}} - µ_{z_{i-1}} in (38) is zero and no update reaches layers L-1 and below. The remark that the cross-covariance "can be easily calculated similar to (21)" is not sufficient, because (21) computes Σ_{u_i,u_i}, not Σ_{z_{i-1},w_i}, and under the stated assumption this cross-covariance is exactly zero. Section VI-A's admission that off-diagonal covariance terms have no closed-form backward propagation does not fix the inconsistency. Either the algorithm silently updates only the final layer's weights, contradicting the claim of full-head sequential Bayesian fine-tuning, or it relies on an unstated non-independence model that is never derived. This issue must be resolved before the method can be accepted as described.
- [Section IV-C and Algorithm 2 lines 22-24] The recursive update (13) is only a valid Bayesian recursion if p(X_k,Y_k|W_B) is a genuine likelihood. The paper never defines this likelihood; instead, Algorithm 2 sets µ^+_{z_L} to the one-hot target and Σ^+_{z_L,z_L} to Σ_data. This is equivalent to a Gaussian observation model y = p + noise for a softmax output p, which is not the categorical likelihood that generated the token data. The Gaussian treatment of a one-hot vector is an approximation whose effect on the posterior is not analyzed. The paper should either specify the likelihood and justify the Gaussian approximation, or present the method as an assumed-density filter rather than exact Bayesian inference.
- [Section IV-A and Algorithm 1] Algorithm 1 preprocesses each autoregressive sequence into a batched representation and the paper states that the augmented data can be treated as i.i.d. However, the individual token targets within a sequence are conditionally dependent, and the correct sequence likelihood is a product of per-token conditional probabilities. A single batched Gaussian measurement update with a common Σ_data does not equal that product likelihood. The preprocessing relabels batches but does not make the tokens within a batch independent. This is a modeling approximation that should be stated explicitly and its consequences for the posterior update discussed.
minor comments (5)
- [Section IV-C, Eqs. (34)-(36)] The block-matrix notation for the weight initialization is difficult to parse; please specify the exact dimensions of each block and the shape of W_i.
- [Section V-B, Figures 3 and 4] The results are averages over 10 trials, but no error bars, standard deviations, or significance tests are reported; the claim that the proposed method "does not experience drops" is based on visual inspection.
- [Section V-A] The experimental section does not state the value of the covariance scale ε or the values of Σ_data used in Figures 3 and 4; please provide these in the setup.
- [Section V-B, Figure 4] The comparison with warm-started retraining uses 100 training epochs per memory buffer, while the proposed method uses a single pass; the per-sample computation time comparison in Figure 4 may conflate epoch counts, so the timing comparison should be clarified.
- [Section IV-B, Eqs. (24)-(27)] The first-order Taylor expansion for the softmax mean is correct, but the paper should state explicitly that higher-order terms are neglected and that the approximation error depends on the variance of u_L, which is not reported.
Circularity Check
No significant circularity: the derivation is a self-contained Bayesian/Kalman update with independent moment-propagation theorems, and no fitted quantity is relabeled as a prediction.
full rationale
The paper's recursive posterior update (Eqs. (6)-(13)) is a standard Bayesian factorization under an i.i.d. assumption; it does not define its conclusion into its premises. The forward-pass ReLU moment formulas are attributed to [18] and [11]; [18] is a peer-reviewed, parameter-free analytic theorem with stated Gaussian and independence assumptions, and it does not itself assert the transformer fine-tuning claim, so it is independent support rather than a circular self-citation. The RTS smoothing equations (37)-(45) and the softmax Taylor approximation are stated in the paper and are standard; the cross-covariance terms are computed from the linear structure, not fitted to the outcome. The measurement covariance Sigma_data in Algorithm 2 is a user-provided input, not a fitted parameter later relabeled as a prediction; the convergence of predicted uncertainty toward Sigma_data is the expected behavior of a Kalman filter. Section VI-A's limitation about missing closed-form backward propagation of off-diagonal covariance is an acknowledged approximation gap, not a circular derivation. Any concern that Eq. (43) yields a zero gain under the forward-pass independence assumption is a potential correctness or internal-consistency issue, not a circularity, and does not change the score.
Assumptions & free parameters
free parameters (3)
- epsilon (initial covariance scale)
- Sigma_data (measurement covariance) =
0, 10, 20, 50 in experiments
- BNN head architecture (L, layer widths n_i)
assumptions (5)
- domain assumption Training sequences are i.i.d.
- domain assumption Pre-activations u_i are Gaussian
- domain assumption Layer weights and previous-layer activations are independent
- ad hoc to paper Softmax logits have small variance for the first-order Taylor expansion to be accurate
- ad hoc to paper One-hot targets can be modeled as Gaussian with known covariance Sigma_data
Cite this review
Pith. "Pith review of Kalman Bayesian Transformer." pith.science (2026). https://pith.science/paper/NO5TXVFF
@misc{pith2026250910695,
author = {Pith},
title = {Pith review of: Kalman Bayesian Transformer},
year = {2026},
howpublished = {\url{https://pith.science/paper/NO5TXVFF}},
note = {Machine review of arXiv:2509.10695}
}
read the original abstract
Sequential fine-tuning of transformers is useful when new data arrive sequentially, especially with shifting distributions. Unlike batch learning, sequential learning demands that training be stabilized despite a small amount of data by balancing new information and previously learned knowledge in the pre-trained models. This challenge is further complicated when training is to be completed in latency-critical environments and learning must additionally quantify and be mediated by uncertainty. Motivated by these challenges, we propose a novel method that frames sequential fine-tuning as a posterior inference problem within a Bayesian framework. Our approach integrates closed-form moment propagation of random variables, Kalman Bayesian Neural Networks, and Taylor approximations of the moments of softmax functions. By explicitly accounting for pre-trained models as priors and adaptively balancing them against new information based on quantified uncertainty, our method achieves robust and data-efficient sequential learning. The effectiveness of our method is demonstrated through numerical simulations involving sequential adaptation of a decision transformer to tasks characterized by distribution shifts and limited memory resources.
Figures
Reference graph
Works this paper leans on
-
[1]
Parameter-efficient fine-tuning for large models: A comprehensive survey,
Z. Han, C. Gao, J. Liu, J. Zhang, and S. Q. Zhang, “Parameter-efficient fine-tuning for large models: A comprehensive survey,”arXiv preprint arXiv:2403.14608, 2024
arXiv 2024
-
[2]
Bayestune: Bayesian sparse deep model fine-tuning,
M. Kim and T. Hospedales, “Bayestune: Bayesian sparse deep model fine-tuning,” inAdvances in Neural Information Processing Systems (A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, eds.), vol. 36, pp. 65317–65365, Curran Associates, Inc., 2023
work page 2023
-
[3]
Bayesian active learning with pretrained language models,
K. Margatina, L. Barrault, and N. Aletras, “Bayesian active learning with pretrained language models,”CoRR, vol. abs/2104.08320, 2021
arXiv 2021
-
[4]
Bayesian low- rank adaptation for large language models,
A. X. Yang, M. Robeyns, X. Wang, and L. Aitchison, “Bayesian low- rank adaptation for large language models,” 2024
work page 2024
-
[5]
Practical variational inference for neural networks,
A. Graves, “Practical variational inference for neural networks,” in Advances in Neural Information Processing Systems(J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K. Weinberger, eds.), vol. 24, p. 2348–2356, Curran Associates, Inc., 2011
work page 2011
-
[6]
Stochastic variational inference,
M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley, “Stochastic variational inference,”the Journal of machine Learning research, vol. 14, no. 1, pp. 1303–1347, 2013
2013
-
[7]
Equation of state calculations by fast computing machines,
N. Metropolis, A. W. Rosenbluth, M. N. Rosenbluth, A. H. Teller, and E. Teller, “Equation of state calculations by fast computing machines,” The journal of chemical physics, vol. 21, no. 6, pp. 1087–1092, 1953
1953
-
[8]
Stochastic relaxation, gibbs distributions, and the bayesian restoration of images,
S. Geman and D. Geman, “Stochastic relaxation, gibbs distributions, and the bayesian restoration of images,”IEEE Transactions on pattern analysis and machine intelligence, no. 6, pp. 721–741, 1984
work page 1984
Show all 31 references
-
[9]
R. M. Neal,Bayesian learning for neural networks. Springer Science & Business Media, 1992
1992
-
[10]
Stochastic gradient markov chain monte carlo,
C. Nemeth and P. Fearnhead, “Stochastic gradient markov chain monte carlo,”Journal of the American Statistical Association, vol. 116, no. 533, pp. 433–450, 2021
2021
-
[11]
Kalman Bayesian neural networks for closed-form online learning,
P. Wagner, X. Wu, and M. F. Huber, “Kalman Bayesian neural networks for closed-form online learning,” inProceedings of the AAAI Conference on Artificial Intelligence, vol. 37, pp. 10069–10077, 2023
2023
-
[12]
Bayesian transformer language models for speech recognition,
B. Xue, J. Yu, J. Xu, S. Liu, S. Hu, Z. Ye, M. Geng, X. Liu, and H. Meng, “Bayesian transformer language models for speech recognition,” inICASSP 2021 - 2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 7378–7382, 2021
2021
-
[13]
Bayesformer: Trans- former with uncertainty estimation,
K. A. Sankararaman, S. Wang, and H. Fang, “Bayesformer: Trans- former with uncertainty estimation,” 2022
2022
-
[14]
Variational learning in nonlinear Gaus- sian belief networks,
B. J. Frey and G. E. Hinton, “Variational learning in nonlinear Gaus- sian belief networks,”Neural Computation, vol. 11, no. 1, pp. 193– 213, 1999
1999
-
[15]
Tractable inference for complex stochastic processes,
X. Boyen and D. Koller, “Tractable inference for complex stochastic processes,” inProceedings of the Fourteenth Conference on Uncer- tainty in Artificial Intelligence, UAI’98, (San Francisco, CA, USA), p. 33–42, Morgan Kaufmann Publishers Inc., 1998
1998
-
[16]
T. P. Minka,A Family of Algorithms for Approximate Bayesian Inference. PhD thesis, Massachusetts Institute of Technology, USA,
-
[17]
Deterministic variational inference for robust Bayesian neural networks,
A. Wu, S. Nowozin, E. Meeds, R. E. Turner, J. M. Hern ´andez- Lobato, and A. L. Gaunt, “Deterministic variational inference for robust Bayesian neural networks,” inInternational Conference on Learning Representations, 2019
2019
-
[18]
An analytic solution to covariance propagation in neural networks,
O. Wright, Y . Nakahira, and J. M. F. Moura, “An analytic solution to covariance propagation in neural networks,” inInternational Confer- ence on Artificial Intelligence and Statistics, pp. 4087–4095, PMLR, 2024
2024
-
[19]
Lightweight probabilistic deep networks,
J. Gast and S. Roth, “Lightweight probabilistic deep networks,” in Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3369–3378, 2018
2018
-
[20]
Probabilistic backpropaga- tion for scalable learning of Bayesian neural networks,
J. M. Hern ´andez-Lobato and R. P. Adams, “Probabilistic backpropaga- tion for scalable learning of Bayesian neural networks,” inProceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37, ICML’15, p. 1861–1869, JMLR.org, 2015
2015
-
[21]
Assumed density filtering methods for learning bayesian neural networks,
S. Ghosh, F. Delle Fave, and J. Yedidia, “Assumed density filtering methods for learning bayesian neural networks,” inProceedings of the AAAI Conference on Artificial Intelligence, vol. 30, 2016
2016
-
[22]
A practical Bayesian framework for backpropa- gation networks,
D. J. C. MacKay, “A practical Bayesian framework for backpropa- gation networks,”Neural Computation, vol. 4, no. 3, pp. 448–472, 1992
1992
-
[23]
What uncertainties do we need in bayesian deep learning for computer vision?,
A. Kendall and Y . Gal, “What uncertainties do we need in bayesian deep learning for computer vision?,”Advances in neural information processing systems, vol. 30, 2017
2017
-
[24]
Aleatoric and epistemic uncer- tainty in machine learning: An introduction to concepts and methods,
E. H ¨ullermeier and W. Waegeman, “Aleatoric and epistemic uncer- tainty in machine learning: An introduction to concepts and methods,” Machine learning, vol. 110, no. 3, pp. 457–506, 2021
2021
-
[25]
What are bayesian neural network posteriors really like?,
P. Izmailov, S. Vikram, M. D. Hoffman, and A. G. G. Wilson, “What are bayesian neural network posteriors really like?,” inInternational conference on machine learning, pp. 4629–4640, PMLR, 2021
2021
-
[26]
Learnable uncertainty un- der laplace approximations,
A. Kristiadi, M. Hein, and P. Hennig, “Learnable uncertainty un- der laplace approximations,” inUncertainty in Artificial Intelligence, pp. 344–353, PMLR, 2021
2021
-
[27]
Dropout as a Bayesian approximation: Representing model uncertainty in deep learning,
Y . Gal and Z. Ghahramani, “Dropout as a Bayesian approximation: Representing model uncertainty in deep learning,” inProceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48, ICML’16, p. 1050–1059, JMLR.org, 2016
2016
-
[28]
A simple baseline for Bayesian uncertainty in deep learning,
W. J. Maddox, T. Garipov, P. Izmailov, D. Vetrov, and A. G. Wilson, “A simple baseline for Bayesian uncertainty in deep learning,” inPro- ceedings of the 33rd International Conference on Neural Information Processing Systems, (Red Hook, NY , USA), pp. 13153–13164, Curran Assoc...
2019
-
[29]
Attention is all you need,
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” Advances in neural information processing systems, vol. 30, 2017
2017
-
[30]
Connectionist learning of belief networks,
R. M. Neal, “Connectionist learning of belief networks,”Artificial Intelligence, vol. 56, no. 1, pp. 71–113, 1992
1992
-
[31]
Continual learning us- ing bayesian neural networks,
H. Li, P. Barnaghi, S. Enshaeifar, and F. Ganz, “Continual learning us- ing bayesian neural networks,”IEEE transactions on neural networks and learning systems, vol. 32, no. 9, pp. 4243–4252, 2020
2020
Reviewed August 15, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.