Permutation Invariant Learning with High-Dimensional Particle Filters

Particle filters, or sequential Monte Carlo methods, are widely used for state estimation in non-linear and non-Gaussian settings. They represent probability distributions through a set of samples (particles), providing flexibility in capturing complex dynamics (Doucet et al., 2001a). In fields such as robotics, particle filters have been applied successfully to localization and mapping problems, where they handle uncertainty and non-linearities effectively (Thrun, 2002). However, a key limitation is their scalability: as the dimensionality of the problem increases, the number of particles needed grows exponentially, making them less practical in high-dimensional spaces like those in machine learning (Bengtsson et al., 2008). Recent efforts have focused on improving particle filter scalability through adaptive resampling and dimensionality reduction techniques (Li et al., 2015), but these approaches have not fully bridged the gap for large-scale machine learning applications. Our work addresses this gap by proposing a high-dimensional particle filter that is computationally efficient and well-suited for machine learning tasks.

Bayesian model averaging (BMA) is a powerful technique for integrating uncertainty into model predictions by averaging across multiple models (Hoeting et al., 1999; Wasserman, 2000). By weighting model predictions based on their posterior probabilities, BMA can provide more robust predictions and better capture model uncertainty compared to single-model approaches (Raftery et al., 2005). In modern machine learning, BMA has been employed to enhance performance and uncertainty estimation, notably in ensemble techniques (Lakshminarayanan et al., 2017; Wortsman et al., 2022). However, BMA has not been extensively explored in the context of continual or permutation-invariant learning, where uncertainty over tasks and sequential data plays a crucial role.

In this work, we demonstrate the benefits of particle filters in high-dimensional machine learning settings. We then describe a particular particle filter that functions as a BMA technique, and demonstrate its advantages empirically.

Consider a learning problem that provides a sequence of loss functions of model parameters, and the goal of learning is to minimize the sum of the loss functions. For instance, the loss function at each time step might correspond to the cross-entropy loss on a minibatch of points for a classification problem. We denote the model parameters at time $t$ as $x_{t}\in\mathbb{R}^{d}$ and the loss function at time $t$ as $L_{t}\in\mathbb{R}^{d}\to\mathbb{R}$. The goal is to find an $x$ minimizing $\sum_{t=1}^{T}L_{t}(x)$, where $T$ is the total number of updates.

How can we apply particle filters to this learning problem? To do this, we suppose that instead of learning a single model, we learn a distribution of models following a Bayesian approach. Specifically, suppose that at time $t=0$, we initially start with a prior distribution of candidate models; each $L_{t}$ corresponds to an observation that updates the likelihood of each model. Specifically, we suppose that the likelihood of the model $x$ is set as:

This likelihood function increases the likelihood of models that achieve lower loss values. We denote the prior distribution of models as $p_{0}$ and the posterior distribution after having observed $L_{1}$ through $L_{t}$ as $p_{t}$. Then, $p_{T}$ is given by:

where $Z_{T}$ is a normalization factor that ensures $p_{T}$ integrates to $1$. Observe that this posterior places high density in regions where the summed loss is low.

Particle filters enable the computation of $p_{T}(x)$ by incrementally computing estimates of $p_{t}(x)$. Specifically, given $p_{t}(x)$, we may compute $p_{t+1}(x)$ as:

This Bayesian update equation may often be intractable to compute exactly, particularly when $p_{t}$ does not have a known parametric form. Instead of tracking $p_{t}$ exactly, particle filters track an estimate $\hat{p}_{t}$ instead:

where $\delta$ is a delta function, $x^{(i)}_{t}$ represents the $i$th particle at time $t$ and $w_{t}^{(i)}$ represents the weight of the particle at time $t$. Each particle filter then has a different method of estimating the Bayesian update of Equation 3. After all updates are complete, an ensemble of particles is available, each of which is an estimate of the global minimizer of $\sum_{t=1}^{T}L_{t}(x)$. We denote the output distribution of a particle filter initialized at $\hat{p}_{0}$ trained on a sequence of loss functions $L_{1},...L_{T}$ as $\hat{p}_{0}[L_{1},...L_{T}]$.

Since particle filters aim to approximate Bayesian updates, we suppose that each update outputs a set of particles close to the true posterior. To formalize this, suppose that there exists a symmetric, non-negative discrepancy measure $D(p,q)$ that satisfies the triangle inequality:

for all $p,q,r$. Furthermore, suppose $D(p,p)=0$ for all $p$.

Now, suppose that the particle filter satisfies the following two conditions:

and

for some constants $C$ and $\epsilon$. $C$ can be interpreted as the estimation error propagating from the previous iteration, while $\epsilon$ can be viewed as error between each particle filter update and the true Bayesian update. This allows the discrepancy at time $T$ to be bounded as:

This decomposes the discrepancy at time $T$ as a term depending on the initial discrepancy $D(\hat{p}_{0},p_{0})$ and a term depending on the incremental discrepancy $\epsilon$.

Next, we demonstrate that particle filters are approximately permutation-invariant: they produce an output that is nearly invariant to the ordering of loss functions $L_{t}$. We show that $\hat{p}_{0}[L_{1},...L_{T}]$ is similar to $\hat{p}_{0}[L_{\sigma_{1}},...L_{\sigma_{T}}]$ where $\sigma_{1},\sigma_{2},...\sigma_{T}$ is a permutation of $1,2,...T$.

See Appendix A for a proof. This result demonstrates that particle filters are approximately permutation invariant, especially over small sequences. Standard learning algorithms such as gradient descent are notably not permutation-invariant: they tend to be highly dependent on the ordering of data points. Permutation-invariance enables learning algorithms with less stochastic outputs: in a perfectly permutation-invariant particle filter, the only potential sources of randomness are the initial selection of particles and the randomness in the particle filter updates themselves.

Now, we demonstrate that particle filters naturally avoid catastrophic forgetting and loss of plasticity. Catastrophic forgetting can be formalized in our framework as the phenomenon where a learning algorithm trained on a sequence of losses $L_{1},...L_{T}$ performs poorly on the earlier losses it is trained on. Similarly, loss of plasticity corresponds to performing poorly on later losses. We provide an upper bound on the loss at any point in training:

See Appendix B for a proof. We make two key assumptions in this theorem: the difference in average loss under two different distributions can be bounded by $D$ and that each step in the particle filter reduces the loss on which it is trained by at least a fixed factor. We believe that the first assumption may in many cases be reasonable if the loss function is sufficiently slow-changing: small changes in the distribution over $x$ should not change the average loss value. The second assumption may also be reasonable under many settings for effective particle filters as well as other standard learning algorithms; with a fixed loss function, it corresponds to a linear convergence rate. Gradient descent, for example, satisfies this assumption on loss functions satisfying the Polyak-Łojasiewicz inequality. The resulting bound on the loss $L_{i}$ guarantees that the performance can be no worse than if there had only been a single update $L_{i}$ (yielding loss at most $\beta M$) plus an additional error term.

Given the desirable properties of particle filters in learning problems, here, we describe a particular particle filter well suited to the high-dimensional spaces found in most machine learning settings.

Suppose that our particle filter’s particle distribution $\hat{p}_{t}(x)=\sum_{i}w_{t}^{(i)}\delta(x-x^{(i)}_{t})$ represents another distribution $\tilde{p}_{t}(x)$ constructed as:

where $\sigma$ is a variance parameter and $Z$ is a normalizing constant. Essentially, $\tilde{p}_{t}(x)$ replaces each delta function in $\hat{p}_{t}(x)$ with a isotropic Gaussian. We derive the particle filter’s update on $\hat{p}_{t}(x)$ as an approximation of the optimal Bayesian update applied to $\tilde{p}_{t}(x)$. Observe that given a prior of $\tilde{p}_{t}(x)$ and the loss function $L_{t+1}$, the posterior is proportional to:

We manipulate this expression until it can be expressed in the form of Equation 13. First, we make a linear approximation of $L_{t+1}$ centered at each particle $x^{(i)}_{t}$:

Pulling $e^{-L_{t+1}(x)}$ into the summation and applying this approximation:

Grouping terms:

Completing the square in the exponent:

where $x^{(i)}_{t+1}=x^{(i)}_{t}-\sigma^{2}\nabla L_{t+1}(x^{(i)}_{t})$. Simplifying the constant terms:

Observe that under our linear approximation, $L_{t+1}(x_{t+1}^{(i)})=L_{t+1}(x^{(i)}_{t})-\sigma^{2}||\nabla L_{t+1}(x^{(i)}_{t})||^{2}$. Thus, we may write the expression as:

Finally, we define $w_{t+1}^{(i)}=w_{t}^{(i)}e^{-\frac{L_{t+1}(x_{t+1}^{(i)})+L_{t+1}(x^{(i)}_{t})}{2}}$ to arrive at our final approximation of the posterior:

We represent this posterior with particles $x^{(i)}_{t+1}$ and respective weights $w^{(i)}_{t+1}$.

We summarize the update equations of this particle filter below:

Algorithm 1 shows the full pseudocode of the filter. For simplicity, we do not normalize the weights of the particles at each iteration; this can be done once at the end of training. Intuitively, this filter updates the positions of the particles with gradient descent, but reweights the particles based on their performance at the old and new points, with lower-loss particles weighted higher. Like gradient descent and other gradient-based optimization procedures, this particle filter is well suited for optimization on high-dimensional spaces, while retaining the properties of a particle filters outlined in the prior sections such as permutation-invariance and avoidance of catastrophic forgetting. Figure 1 illustrates how our method converges to well-performing regions of the parameter space over time.