Online Learning via Learned Latent Bayesian Tracking¶

Conference: NeurIPS2026 (Accepted)
arXiv: 2609.31559
Code: https://github.com/aura-online-adaptation
Area: Optimization & Theory
Keywords: online learning, latent Bayesian tracking, extended Kalman filtering, meta-learning, distribution drift

TL;DR¶

The paper learns a low-dimensional parameter-generating space and dynamical prior offline, performs one latent extended Kalman filtering update per labeled sample online, and reconstructs prediction parameters, improving adaptation under limited supervision and computational constraints in image classification and time-varying wireless reception.

Background & Motivation¶

Learning in a non-stationary environment requires more than retraining a model for a new domain: the predictor must track current conditions as observations arrive. Wireless channels change over time, as does the composition of image corruptions, so previously accurate weights can become mismatched. Online gradient descent can correct parameters using new labels, but multiple optimization steps and learning-rate sensitivity become practical obstacles when each sample has a short processing window. The objective here is prediction quality on the current distribution, not retention of performance on every historical task.

Kalman-style methods offer a different update logic. They treat adaptation parameters as a time-evolving state, predict using a dynamical model, correct using the discrepancy between a new label and the prediction, and propagate state uncertainty. However, deep networks have too many parameters: a full covariance requires at least quadratic storage and computation. Existing methods often use diagonal or low-rank covariance, or update only a manually selected final layer. These choices reduce cost but may fail to preserve the coordinated parameter directions actually needed under distribution shift.

Rather than seeking another inexpensive filtering approximation in high-dimensional parameter space, the paper learns a representation in which filtering itself is manageable. Offline learning has access to multiple labeled non-stationary trajectories, whereas deployment provides only sparse reliable supervision. A latent space and parameter lifting map trained specifically for online adaptation connect these regimes. Core idea: jointly learn adaptation degrees of freedom, temporal evolution, and parameter generation offline, so that online Bayesian updates track a low-dimensional state rather than the entire network weight vector.

Method¶

Overall Architecture¶

The inputs are the current labeled sample and the previous latent-state mean and covariance; the output is an adapted parameter vector for subsequent inference under the current distribution. Offline trajectory meta-learning determines the parameter lifting map, state transition, and optionally noise covariances. These meta-parameters are not retrained through an entire trajectory at deployment. Online processing consists of latent dynamics prediction, label-driven correction, and parameter reconstruction, with one filtering update per supervised sample.

“Reconstructing the full parameters” means generating the selected prediction parameter vector, not allowing every network weight to vary independently in arbitrary directions. The theoretical framework considers a full parameter vector, but the image implementation adapts only the final fully connected classifier and selected batch-normalization parameters, assuming independent groups with block-diagonal covariance. The communication implementation adapts individual DeepSIC detector blocks. Unselected backbone weights do not acquire online degrees of freedom through this procedure.

Solid arrows below represent deployment-time information flow; dashed arrows show how fixed components learned offline configure online processing. Correction requires online labels, whereas subsequent prediction does not. Offline query labels supervise meta-parameters and must not enter deployment predictions.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    D["Offline labeled trajectories<br/>support and query samples"] --> A["Trajectory meta-learning"]
    A --> G["Parameter lifting map"]
    G -.->|Fixed map for observation and reconstruction| K["Label-driven correction"]
    A -.->|Learned transition and noise parameters| P["Latent dynamics prediction"]
    S["Previous latent mean and covariance"] --> P
    P --> K
    X["Current supervised input and label"] --> K
    K -->|Lift corrected latent mean| O["Adapted parameters<br/>subsequent current-distribution predictions"]
    K -->|Store latent mean and covariance| S

The v2 cache renders the method abbreviation as name in the body, while some tables and figure captions contain AURA. This note uses the descriptive term “latent Bayesian tracking” rather than inventing a formal method brand from the incomplete abbreviation. “Proposed method” below refers to these source entries.

Key Designs¶

1. Trajectory meta-learning: supervise representations through predictions after online updates

Conventional parameter compression asks whether generated weights approximate the original weights, but does not guarantee that a new label produces an effective update in compressed space. The paper treats each non-stationary trajectory as an adaptation task. At each time step, the first labeled sample is the support sample and triggers one filtering update; the remaining samples form the query set and evaluate the updated predictor. Accumulated query loss assesses whether the latent space facilitates adaptation, rather than merely reconstructing a static network.

Meta-learning back-propagates through the differentiable filtering recursion and lifting map, jointly optimizing the map parameters and diagonal transition matrix; process and observation noise covariances can optionally be learned as well. Non-stationary training trajectories therefore determine both which variation directions to preserve and how the state should move over time. This shifts training cost offline without removing the need for representative trajectories and reliable labels. The cold-start ablation tests whether dimensionality restriction alone is sufficient without this learning.

2. Parameter lifting map: generate coordinated parameter changes from a low-dimensional state

The latent state is much smaller than the adapted parameter vector, and the lifting map turns it into parameters used by the predictor. A change in one latent coordinate can therefore move many weights together, avoiding separate update estimation in every parameter direction. Expressiveness comes from the learned parameter-generating geometry, not from lifting the state and then performing unrestricted full-weight optimization online.

The main image results correspond to an affine map. An appendix variant adds a residual nonlinear map: a single-hidden-layer MLP with width 500 and GELU activation supplements the affine component. Nonlinearity can represent curved adaptation trajectories, but the online observation Jacobian must pass through this map, and parameter reconstruction requires its forward computation. Accuracy improvements therefore have a cost. Even when the map generates many weights, reachable online parameters remain restricted to its image; this does not restore arbitrary full-parameter adaptation.

3. Latent dynamics prediction: provide a temporal prior before the next supervised observation

Before correction, the filter predicts the latent state using linear dynamics, propagates its covariance, and adds process noise. The main method restricts the transition matrix to be diagonal; OU-type dynamics with a scaled identity are a simpler special case. This is a tractable approximate prior in a learned space, not a claim that optimal network weights follow exactly linear dynamics.

The structure matters when labels are scarce: prediction carries forward previous adaptation directions, while process noise permits deviations from the old dynamics. Intervening unlabeled frames perform prediction only, without observation correction, and errors can accumulate. Such short-term extrapolation is not unsupervised adaptation. A dense transition matrix can model coupling between latent coordinates but makes full-covariance propagation cubic in cost.

The core state representation in the paper is:

\[ \theta_t^\star=\mathcal{G}(z_t),\qquad z_{t+1}=Fz_t+v_t. \]

Here \(\mathcal{G}\) is the parameter lifting map, \(F\) is the state transition matrix, and \(v_t\) has covariance \(Q\). Online processing maintains a Gaussian approximation to the latent-state posterior, not a posterior over all network weights. The predictor uses one deterministic parameter instance generated from the corrected latent mean.

4. Label-driven correction: use innovation and uncertainty to determine update direction and magnitude

The observation model treats the one-hot label as predicted class probabilities plus observation error. At the predicted latent mean, the filter computes the Jacobian through “latent state → lifted parameters → class probabilities.” The difference between the label and predicted probabilities is the innovation. The Kalman gain combines latent covariance, the Jacobian, and observation noise to translate this innovation into a latent-state correction, using state uncertainty more explicitly than multiplying every sample update by a fixed learning rate.

The following preserves the source update mechanism, abbreviating predicted covariance and mean as \(P_t=\Sigma_{t|t-1}\) and \(\bar z_t=\hat z_{t|t-1}\); \(H_t\) is the Jacobian of class probabilities with respect to the latent state, and \(R\) is observation noise covariance.

\[ \begin{aligned} K_t&=P_tH_t^\top(H_tP_tH_t^\top+R)^{-1},\\ \hat z_t&=\bar z_t+K_t\left(\operatorname{OneHot}(y_t)-f_{\mathcal{G}(\bar z_t)}(x_t)\right),\\ \Sigma_t&=P_t-K_tH_tP_t. \end{aligned} \]

Because class probabilities lie on a simplex, observation covariance needs positive-definite regularization for numerical stability. The arguments for zero-mean label residuals uncorrelated with ideal probability outputs require conditions such as sufficient model expressiveness to represent the true conditional class distribution. They do not make actual residuals exactly Gaussian. Extended Kalman filtering (EKF) still uses local first-order linearization and a Gaussian approximation; strong nonlinearity or an incorrect prior can undermine correction.

A low-dimensional state reduces covariance size but does not eliminate all large-matrix and network computations. Latent covariance is \(m\times m\), the Jacobian is \(C\times m\), and the gain requires an inverse of a \(C\times C\) matrix. The latter matters when there are many classes. With a small class dimension and diagonal transition, the paper gives per-step complexity \(\mathcal{O}(n+m^2)\), where \(n\) denotes the lifting network's parameter count. Actual execution still generates adapted parameters and differentiates through both the predictor and lifting map to obtain the Jacobian.

A Worked Example¶

Consider one frame in the QuaDRiGa communication evaluation: three users transmit to five receive antennas, with a channel fixed within a frame and varying between frames. After synchronization, each frame supplies six known pilots. The receiver uses their labels sequentially for six “dynamics prediction → innovation correction” updates. Each update changes the latent state and then generates detector-block weights, rather than treating the six pilots as a new training set to optimize to convergence.

The adapted receiver then detects another 1000 QPSK symbol vectors under the same channel and measures bit error rate (BER). Those payload symbols do not supply correction labels. If labels arrive only every several frames, intervening frames propagate the latent mean and covariance until the next pilots permit correction. This is intermittent feedback in supervised tracking, not unlabeled test-time adaptation.

Loss & Training¶

The offline objective averages cross-entropy across trajectories, time, and query samples. At each time step, the history-dependent online operator \(\mathcal{A}_t\) first adapts using the support sample, after which query loss is evaluated:

\[ \mathcal{L}_{\mathcal{D}}^{\rm meta}(\phi)=\frac{1}{|\mathcal{D}|}\sum_{i=1}^{|\mathcal{D}|}\frac{1}{T}\sum_{t=1}^{T}\frac{1}{B-1}\sum_{b=2}^{B}\ell\!\left(f_{\mathcal{A}_t(x_{t,1}^{(i)},y_{t,1}^{(i)};\phi)}(x_{t,b}^{(i)}),y_{t,b}^{(i)}\right). \]

Here \(B\) is the offline number of labeled samples per time step, not a required deployment buffer size. Offline trajectory back-propagation optimizes meta-parameters; online processing fixes those parameters and only recurses over latent means and covariances. Communication training and evaluation use disjoint channel trajectories. The main image evaluation instead uses different mixing trajectories from the same corruption pool, not unseen corruption-family generalization; an appendix experiment separately tests disjoint corruption families.

Key Experimental Results¶

Main Results¶

Image streams contain 1000 time steps with 100 images each. Supervised streaming methods use only 5, 10, or 15 labeled samples per step. MNIST-C, CIFAR-10-C, and CIFAR-100-C use ResNet-18, Wide ResNet-28-10, and a 29-layer ResNeXt, respectively; CIFAR corruption severity is 5. The following excerpts Table 1. Classification error is in percent, lower is better, and each cell uses 3 seeds.

Method	Adaptation samples and supervision per step	MNIST-C	CIFAR-10-C	CIFAR-100-C
Source	0	\(47.93\pm19.83\)	\(41.81\pm5.62\)	\(37.41\pm4.82\)
Online GD	5, labeled	\(7.50\pm0.50\)	\(29.46\pm5.55\)	\(37.93\pm1.83\)
EKF-FC	5, labeled	\(19.74\pm5.09\)	\(30.62\pm1.34\)	\(31.27\pm1.98\)
Proposed method, affine lifting	5, labeled	\(2.76\pm0.63\)	\(16.13\pm1.50\)	\(29.47\pm0.62\)
Online GD	15, labeled	\(4.53\pm0.39\)	\(22.37\pm1.38\)	\(33.22\pm3.35\)
EKF-FC	15, labeled	\(12.16\pm3.06\)	\(26.99\pm5.26\)	\(32.43\pm1.55\)
Proposed method, affine lifting	15, labeled	\(1.63\pm0.09\)	\(14.28\pm2.87\)	\(26.66\pm3.89\)
ROID	100, unlabeled, batch-wise	\(2.64\pm1.37\)	\(16.00\pm1.29\)	\(26.50\pm1.37\)

With the same budget of 5 labels, the proposed method reduces CIFAR-10-C error by 13.33 percentage points relative to Online GD. With 15 labels, its CIFAR-100-C error remains slightly above ROID's 26.50%. ROID uses statistics from an unlabeled batch, whereas the proposed method updates sample by sample using reliable labels. Neither supervision nor buffering budgets match, so this table does not establish comprehensive superiority in unsupervised test-time adaptation.

Main-paper Table 1 specifies latent dimensions \(m=256,500,500\), but Appendix C.2 states \(m=128\). The main paper also calls the CIFAR-100 backbone ResNet-29, whereas the appendix specifies ResNeXt. This note retains main-table dimensions and uses the detailed implementation description for the backbone while explicitly preserving these source discrepancies rather than harmonizing the configurations.

The communication main experiment uses QuaDRiGa trajectories under 3GPP Indoor Office LOS, with both linear channels and nonlinear channels incorporating tanh distortion. Of 150 frames, the first four provide 64 labels per frame for synchronization, and the remaining 146 provide six labels per frame for tracking. Main figures report lower BER and more stable tracking, but the text cache lacks curve data; exact SNR–BER values and FLOPs are not inferred from captions.

Ablation Study¶

Table 2 compares learning versus no learning offline under the same dynamics, latent dimension, and online budget. The following selects two dynamics models and two dimensions. BER is in percent; the reported error terms are not reinterpreted as confidence intervals unspecified by the source.

Dynamics	Latent dimension	Offline learning	BER, 8 dB	BER, 10 dB
OU	100	Yes	\(0.832\pm0.485\)	\(0.512\pm0.379\)
OU	100	No	\(14.399\pm3.024\)	\(11.542\pm1.054\)
OU	700	Yes	\(0.311\pm0.367\)	\(0.210\pm0.152\)
OU	700	No	\(14.870\pm1.626\)	\(12.980\pm1.487\)
F-Diag	100	Yes	\(0.0757\pm0.0079\)	\(0.0608\pm0.0045\)
F-Diag	100	No	\(18.471\pm3.024\)	\(17.464\pm1.622\)
F-Diag	700	Yes	\(0.0199\pm0.0009\)	\(0.0176\pm0.0004\)
F-Diag	700	No	\(18.225\pm1.626\)	\(17.027\pm1.815\)

Simple parameter compression does not reproduce the benefit of offline learning. Learned diagonal dynamics also outperform these OU configurations, but this does not imply that every environment needs a more complex transition matrix. The dense-transition ablation has covariance prediction cost \(\mathcal{O}(m^3)\) and is attractive in runtime only around \(m\leq250\) in the reported implementation.

Table 4 compares lifting maps under the same latent dimension and online protocol. Classification error is in percent, reported as mean and standard deviation over 3 seeds.

Dataset	Labeled samples	Affine lifting	Nonlinear lifting
CIFAR-10-C	5	\(16.13\pm1.50\)	\(10.21\pm0.40\)
CIFAR-10-C	10	\(16.84\pm0.51\)	\(10.15\pm0.81\)
CIFAR-10-C	15	\(14.28\pm2.87\)	\(9.75\pm0.54\)
CIFAR-100-C	5	\(29.47\pm0.62\)	\(26.34\pm1.77\)
CIFAR-100-C	10	\(29.65\pm1.97\)	\(25.21\pm1.04\)
CIFAR-100-C	15	\(26.66\pm3.89\)	\(23.00\pm0.79\)

Key Findings¶

Larger latent dimensions are not universally better. Image dimension sweeps improve initially, then saturate or deteriorate slightly. Larger dimensions usually improve the communication table results, but covariance and lifting costs still matter.
Label intervals reveal a clear boundary. In Appendix Table 3, six pilots every frame yield BER \(0.07\pm0.01\)%, whereas six pilots every five frames yield \(4.54\pm0.96\)%. Prediction alone does not guarantee long-term stability.
Cross-simulator transfer does not establish uniform superiority. When trained on QuaDRiGa and tested on COST2100, the paper describes BER as broadly comparable to leading full-parameter Bayesian baselines, not consistently lower.
Abrupt changes permit gradual, not immediate, recovery. Across five CIFAR-10-C shifts in Appendix Table 8, affine lifting reduces mean error from 45.96% to 19.44% after 40 updates, and nonlinear lifting reduces it from 42.44% to 17.66%. Each update step supplies 15 labels.
Hardware results establish a millisecond-order budget for a particular prototype. On CPU, a 1.8K-parameter WiFi CNN requires 0.084 ms per packet for inference and 10.573 ms for adaptation. The 6.6K- and 8.2K-parameter models require 40.414 and 84.393 ms for adaptation, not universally sub-1-ms operation.

Highlights & Insights¶

The learned object is a parameter geometry that facilitates updates, not a static low-dimensional encoding. Support-sample adaptation followed by query-sample evaluation binds the representation to the actual filtering rule.
Bayesian machinery serves online optimization here, rather than uncertainty quantification over every network weight. Distinguishing the latent posterior from deterministic parameter instances clarifies both efficiency and expressiveness boundaries.
Lifting maps and dynamical priors provide separate cost controls. In another streaming classifier, one can independently test whether richer parameter generation justifies its Jacobian cost and whether the temporal prior matches deployment drift.

Limitations & Future Work¶

The current method requires online labels, naturally available as pilots but absent in fully unsupervised deployment. Pseudo-labeling and self-supervised extensions are future directions, not validated by the main experiments.
Low-dimensional compressibility, representative offline trajectories, and moderately smooth dynamics are important premises. Domain mismatch, abrupt shifts, or very large adapted parameter blocks can make the generating space inadequate or the prior misleading.
EKF uses first-order linearization and an approximate noise model for simplex-valued label observations. Richer nonlinearity may improve representation while weakening local linear approximations and increasing cost.
Appendix theory requires conditions including Lipschitz compressibility, strongly quasi-convex risk, and bounded loss. It is not a convergence or arbitrary-drift guarantee for general deep classifiers. The proof also applies a general compression map to a difference vector and introduces an insufficiently matched operator symbol; these details require further verification against the source.
The PlutoSDR+ IEEE 802.11a prototype does not automatically meet real-time deadlines for arbitrary wireless standards or packet sizes. Further work could study change detection, covariance inflation, latent-state resets, and optimized Jacobian and parameter-generation execution.

vs Online GD: GD optimizes directly in parameter space, while this method generates weights through learned latent filtering. The experimental GD baseline performs five gradient updates per label, so the single-correction advantage must be interpreted alongside different update costs and offline training overhead.
vs EKF-FC / BONG: EKF-FC restricts filtering to the classifier, while BONG uses full, diagonal, or diagonal-plus-low-rank covariance in parameter space. The proposed method changes the filtering state space. Image implementation still assumes independent parameter groups; it does not eliminate all structural approximations.
vs ROID / EcoTTA / ELaTTA: These methods do not use ground-truth labels, whereas the proposed method uses sparse supervision. In the efficient-TTA comparison, CIFAR-100-C adaptation latency is 2.01 ms/sample for the proposed method, 1.15 for EcoTTA, and 0.41 for ELaTTA. Accuracy gains do not imply uniformly better latency.
vs LoRA and differentiable filtering: LoRA provides compact adaptation parameterization, and differentiable filters learn state estimators. This work uses related ideas to learn an online parameter-tracking rule at trajectory level. The appendix LoRA comparison supports added value from learning adaptation space, but adapted parameter locations differ, preventing attribution to a single mechanism.

Rating¶

Novelty: 4/5, jointly trains latent parameter generation, dynamical priors, and differentiable online filtering beyond simple covariance compression.
Experimental Thoroughness: 4/5, covers images, wireless channels, cold starts, nonlinear lifting, sparse feedback, abrupt shifts, and hardware, with some inconsistent configuration descriptions.
Writing Quality: 3/5, the main argument is clear, but cached abbreviations, backbone and dimension descriptions, and proof details need checking.
Value: 4/5, useful for low-latency systems with reliable sparse supervision and structured drift, but not yet a general unsupervised adaptation solution.