Learn to Adapt Doubly-sequential Bayesian Neural Networks with Complex Architectures

Neural networks used for control must adapt from a few sequentially arriving samples, often drawn from suboptimal data, and do so efficiently. Bayesian neural networks (BNNs) provide a principled framework for sample-efficient sequential adaptation with quantified uncertainty, but few techniques can use data to optimize that adaptation at scale. In this paper, we first show that adaptation can be learned from data by optimizing its backward update one layer at a time. This stands in stark contrast to existing meta-learning techniques that learn only from input–output matching and become prohibitive for large networks. The key enabler is to view training as a sequence of layerwise Bayesian posterior updates, which decomposes the problem into optimizing the update rule of each layer. Layerwise moment propagation of this kind was previously confined to feedforward networks; we extend it to computation graphs composed of attention, normalization, feed-forward, residual, and branching operations, so that transformers and other architectures can use data-driven techniques to optimize adaptation/fine tuning at scale. In parallel, we generalize BNNs to a hybrid of imitation-based and reinforcement-based learning so that BNNs learn a more optimal mapping than the data distribution. This builds on a new perspective: the uncertainty a BNN maintains can be repurposed to distinguish optimal from suboptimal samples and to weight the update toward the former. Finally, we integrate the two parts into a data-driven framework to optimize both the update rules of individual layers and the balance between imitation and reinforcement. The proposed framework allows the adaptation algorithm to use past adaptation histories to self-evolve toward improved sample efficiency, optimality, and robustness.