Pith. sign in

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 →

arxiv 2509.10695 v1 pith:NO5TXVFF submitted 2025-09-12 cs.LG cs.AI

classification cs.LGcs.AI
keywords KalmanBayesianTransformersequentialfine-tuningmomentpropagationuncertaintyquantificationcatastrophicforgettingdecisionRauch-Tung-Striebelsmootheronlineinference
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper proposes a way to fine-tune a pre-trained transformer sequentially by treating each new sample as a Bayesian measurement that updates a distribution over the weights of the output head. The method, called the Kalman Bayesian Transformer, freezes the transformer body, replaces the linear head with a Bayesian neural network, and propagates means and covariances in closed form through linear, ReLU, and softmax layers using Kalman smoothing. In an inverted-pendulum task, the method needs to store only the latest training sample and achieves higher stabilization success rates than warm-started retraining with replay buffers of 10, 20, 25, and 50 samples, while also returning explicit predictive uncertainty. If this result holds, online fine-tuning in memory-limited, latency-critical settings becomes substantially cheaper and more stable than replay-based retraining.

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.

Watch

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

Editorial extensions of the paper, not claims the author makes directly.

  • 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.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 5 minor

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)
  1. [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.
  2. [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.
  3. [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)
  1. [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.
  2. [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.
  3. [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.
  4. [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.
  5. [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

0 steps flagged · score 0.0 of 10

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 3 free parameters · 5 assumptions · 0 invented entities

The central mechanism rests on standard Bayesian recursion plus a set of approximations: Gaussianity of intermediate activations, independence between weights and activations in the forward pass, first-order Taylor expansion of softmax, and treating one-hot targets as Gaussian measurements with user-supplied covariance. The only numerical quantities fitted or chosen for the experiments are epsilon, Sigma_data, and the BNN head width/depth; none of these is reported fully.

free parameters (3)
  • epsilon (initial covariance scale)
    Sets Sigma_{w,w} = epsilon I and Sigma_{z0,z0} = epsilon I (Algorithm 2 lines 2-8, 13). It must be small to keep the prior near-deterministic but nonzero to avoid singular inversions; its value is not reported, yet it controls the strength of the prior in every Kalman update.
  • Sigma_data (measurement covariance) = 0, 10, 20, 50 in experiments
    Used as the covariance of the target one-hot vector in the backward pass (Algorithm 2 line 24). The paper sweeps values in the uncertainty experiment but does not provide a principled way to choose it in practice; it directly sets the update gain.
  • BNN head architecture (L, layer widths n_i)
    The paper specifies only that the decision transformer has hidden dimension 16; the number of layers L and widths n_i of the replacing Bayesian network are not given, though they determine the representational capacity and the validity of the positive/negative copy initialization.
assumptions (5)
  • domain assumption Training sequences are i.i.d.
    Stated in Section III: 'We assume that all sets of sequences in the training data are independent and identically distributed.' This underpins the recursive factorization of the posterior in Eq. (7)-(12).
  • domain assumption Pre-activations u_i are Gaussian
    The ReLU moment propagation formulas (Eq. 22-23) from [18] are exact only for Gaussian inputs; the paper tracks only mean and covariance, thereby assuming near-Gaussianity of u_i without stating it explicitly.
  • domain assumption Layer weights and previous-layer activations are independent
    Section IV-B: 'Assuming w_i and z_{i-1} are independent, the forward pass through the linear layer i is given by...' This is used to compute mean and covariance of u_i in Eqs. (19)-(21). In sequential updating, the weights depend on past data that also influenced activations, so the assumption is an approximation.
  • ad hoc to paper Softmax logits have small variance for the first-order Taylor expansion to be accurate
    Eq. (24) expands softmax around the mean and keeps only the Jacobian term. The paper notes 'When the variance is small' but does not verify this condition for the trained network.
  • ad hoc to paper One-hot targets can be modeled as Gaussian with known covariance Sigma_data
    Algorithm 2 lines 23-24 set the smoothed output mean to the one-hot target vector and covariance to Sigma_data. This treats a discrete target as a continuous Gaussian observation, which is not the true categorical likelihood.

how reviews work

0 comments
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

Figures reproduced from arXiv: 2509.10695 by the authors.

Figure 1
Figure 1. Proposed architecture. The original transformer archi [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Illustration for one layer of the Bayesian neural [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗
Figure 5
Figure 5. The uncertainty predicted by the proposed method [PITH_FULL_IMAGE:figures/full_fig_p007_5.png] view at source ↗
Figures from the paper (1 more)
Figure 4
Figure 4. Figure 4: The average computation time. The blue line shows [PITH_FULL_IMAGE:figures/full_fig_p007_4.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

31 extracted references · 25 canonical work pages

  1. [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

  2. [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

  3. [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

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

Show all 31 references
  1. [9]

    R. M. Neal,Bayesian learning for neural networks. Springer Science & Business Media, 1992

  2. [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

  3. [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

  4. [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

  5. [13]

    Bayesformer: Trans- former with uncertainty estimation,

    K. A. Sankararaman, S. Wang, and H. Fang, “Bayesformer: Trans- former with uncertainty estimation,” 2022

  6. [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

  7. [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

  8. [16]

    T. P. Minka,A Family of Algorithms for Approximate Bayesian Inference. PhD thesis, Massachusetts Institute of Technology, USA,

  9. [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

  10. [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

  11. [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

  12. [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

  13. [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

  14. [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

  15. [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

  16. [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

  17. [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

  18. [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

  19. [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

  20. [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...

  21. [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

  22. [30]

    Connectionist learning of belief networks,

    R. M. Neal, “Connectionist learning of belief networks,”Artificial Intelligence, vol. 56, no. 1, pp. 71–113, 1992

  23. [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

Pith tools

Reviewed August 15, 2026 · model on record in the stance chip above.