Brain-like variational inference (2410.19315v2)

Abstract: Inference in both brains and machines can be formalized by optimizing a shared objective: maximizing the evidence lower bound (ELBO) in machine learning, or minimizing variational free energy (F) in neuroscience (ELBO = -F). While this equivalence suggests a unifying framework, it leaves open how inference is implemented in neural systems. Here, we show that online natural gradient descent on F, under Poisson assumptions, leads to a recurrent spiking neural network that performs variational inference via membrane potential dynamics. The resulting model -- the iterative Poisson variational autoencoder (iP-VAE) -- replaces the encoder network with local updates derived from natural gradient descent on F. Theoretically, iP-VAE yields a number of desirable features such as emergent normalization via lateral competition, and hardware-efficient integer spike count representations. Empirically, iP-VAE outperforms both standard VAEs and Gaussian-based predictive coding models in sparsity, reconstruction, and biological plausibility. iP-VAE also exhibits strong generalization to out-of-distribution inputs, exceeding hybrid iterative-amortized VAEs. These results demonstrate how deriving inference algorithms from first principles can yield concrete architectures that are simultaneously biologically plausible and empirically effective.

• The paper’s main contribution is bridging the Free Energy Principle with ELBO maximization to construct a brain-inspired model for iterative Bayesian inference.
• It employs Poisson assumptions with recurrent updates that mimic spiking neural dynamics, enhancing adaptiveness and generalization over traditional models.
• Empirical results show improved reconstruction accuracy and efficient transfer learning, highlighting its potential for neuromorphic and energy-efficient AI systems.

An Evaluation of a Brain-Inspired Framework for Iterative Inference in NeuroAI

The paper "A prescriptive theory for brain-like inference" proposes an innovative approach to constructing a brain-inspired neural network that employs iterative inference to perform Bayesian posterior updates. This theory integrates neuroscience principles with machine learning techniques, achieving enhanced adaptiveness and generalization capabilities over traditional neural models by optimizing the Evidence Lower Bound (ELBO) in a spiking neural network.

Overview and Theoretical Foundations

The paper underlines two intertwined components: the Free Energy Principle (FEP), a theoretical framework in neuroscience, and ELBO maximization, widely used in machine learning. The researchers bridge these components to derive a model reflective of spiking neural networks capable of performing robust Bayesian inference through membrane potential dynamics.

Key to this approach is the adoption of Poisson distributions over Gaussian assumptions, resulting in the iterative Poisson Variational Autoencoder (i$$). This choice aligns with empirical findings in neuroscience, where biological neurons are conditionally Poisson. This contrasts sharp Gaussian assumptions, which often create mismatches when mapping ML models onto neural circuits.</p> <h3 class='paper-heading' id='model-and-algorithmic-implementation'>Model and Algorithmic Implementation</h3> <p>The i$$ model extends the Poisson VAE framework into an iterative design, accommodating sequence data under Poisson assumptions. Methodologically, it completely abandons the encoder used in VAEs, instead using a recurrent update rule combining the ELBO objective with neural dynamics inspired by biological neurons.

Through iterative updates on log rates $\bm{u}(t)$, the framework naturally aligns with the membrane potential dynamics of spiking neurons. The updates effectively embody a predictive coding scheme enhanced with Bayesian updates for reconstructing the input sequence data. This approach elegantly avoids several issues associated with earlier predictive coding models, notably by ensuring positivity in firing rates and avoiding explicit prediction terms.

Empirical Results and Generalization

The empirical evaluations demonstrate the iterative i$ model&#39;s effectiveness across several key metrics. It achieves superior reconstruction accuracy, adaptability, and generalizes well to out-of-distribution (OOD) conditions compared to existing amortized models and iterative alternatives like ia-VAE and sa-VAE.</p> <p>Significantly, i$ closes the performance gap identified in preceding work comparing Poisson variants and traditional sparse coding algorithms. The model's learned features are compositional, enhancing its ability to generalize across different datasets and perturbations—for instance, handling rotated MNIST and transferring learning from MNIST to Omniglot efficiently without retraining.

Implications and Future Directions

The i$ framework ushers in an adaptable, energy-efficient model that resonates with the brain's well-documented neural computation methods. Its successful mapping onto biological principles could significantly impact developments in neuromorphic hardware and energy-efficient AI systems.

Future explorations could delve into hierarchical versions of this model or extend its capabilities to dynamic, non-stationary data, such as video sequences. Such enhancements could deepen our knowledge of neural-inspired computation in artificial systems, pushing the boundaries of unsupervised learning paradigms aligned with biological efficiency and complexity.

In conclusion, the paper solidifies the practical and theoretical importance of integrating Poisson-based ELBO optimization with iterative inference. This research not only introduces a compelling brain-inspired architecture but also provides a fertile ground for ongoing advancements in NeuroAI and machine learning's future role in decoding and emulating neural processing.