Permutation Invariant Learning with High-Dimensional Particle Filters

Thank you for your detailed feedback and for acknowledging the solid contribution of our work. We address your concerns point by point.

Clarification on Notation in Section 3.1

We apologize for any confusion regarding the notation. We have revised our text to explain our notation more carefully now.

Notation $L_t: \mathbb{R}^d \rightarrow \mathbb{R}$: This means that the loss function $L_t$​ takes a parameter vector $x \in \mathbb{R}^d$ and outputs a scalar loss value in $\mathbb{R}$.

Definition of $p_t$​: We define $p_t(x)$ as the posterior distribution over model parameters after observing losses up to time $t$. We are happy to clarify any further particular points of confusion.

Significance of Permutation Invariance and Theorem 1

Practical Relevance of Theorem 1: We are unsure what you are referring to when you ask if the distance metric "converges"; we do not make any claims of convergence with respect to any variable. Instead, Theorem 1 provides a bound on how the order of data affects the particle filter's output. With $C$ close to 1 and small $\epsilon$, the discrepancy remains small, indicating approximate permutation invariance. In fact, when $\epsilon=0$, the algorithm is exactly permutation invariant.

Relation to Variance Reduction: We assess permutation invariance by measuring the variance in model performance across different data permutations. Lower variance suggests that the model's output is less sensitive to data ordering, which aligns with permutation invariance.

To further validate the connection between permutation invariance and variance reduction, we are conducting additional experiments with the baseline of SVRG as suggested. We will update with empirical results as soon as they are available.

Validation of Theorem 2 in Experiments

Theorem 2 provides a performance guarantee against both catastrophic forgetting and loss of plasticity. We explicitly demonstrate the avoidance of catastrophic forgetting and loss of plasticity in our experiments in Section 4.1. We are happy to make any modifications to our text to make this connection clear.

As we mention in Section 3.3, the assumptions in Theorem 2 are conditions on the task: the loss function must be sufficiently smooth (equation 10) and easy to optimize with SGD (equation 11). In practice, these are highly reasonable assumptions: the tasks we test on have non-pathological loss landscapes and are amenable to optimization as demonstrated by our experiments in all settings.

Difference from Mixture-of-Experts (MoE)

Our method differs from MoE in two key aspects:

Ensemble vs. Expert Selection: In MoE, only one expert (or a subset) is active for a given input, whereas our method combines all particles' outputs, representing the full posterior distribution.

Bayesian Foundation: Our particle filter is grounded in Bayesian inference, providing theoretical guarantees like permutation invariance and mitigation of forgetting, which are not inherent in standard MoE models.

Given the significant differences between MoE and our approach, we believe standard ensemble methods are a more appropriate baseline than MoE, and have used this in our experimental evaluation.

Dimensionality of Particles

Each particle encompasses the full set of model parameters. For the neural networks used in our experiments, this means each particle has the same dimensionality as the network's parameter vector. Model architecture details are discussed in Appendix D and in our attached code.

Computational Cost and Runtime

While maintaining multiple particles increases memory usage linearly with the number of particles, the updates for each particle are independent and can be parallelized efficiently, resulting in no runtime overhead. On hardware like GPUs, given sufficient memory, this practically results in runtimes comparable to single-model training.

Justification of Theorem 3 in Non-Linear Settings

Theorem 3 illustrates that our method aligns with Bayesian updates in a simplified setting. Although Theorem 3 considers linear loss functions, it demonstrates our particle filter exactly maintains the correct posterior ratios between particles, preserving Bayesian properties. We may expect that this exact equivalence is relaxed when the loss functions become nonlinear, but it still holds approximately to the extent that loss functions are locally linear.

Relation to Gradient Descent

With $N = 1$, our algorithm reduces to standard gradient descent with a fixed step size. We use SGD with a fixed learning rate of 0.01 for the gradient descent baselines in our experiments.

Minor Points

Loss Function Notation: We clarify that $L_t: \mathbb{R}^d \rightarrow \mathbb{R}$ means the loss function maps parameter vectors to scalar losses.

Abbreviation "BC": Yes, "BC" stands for Behavior Cloning.

Figure Captions and Language: Thank you for these suggestions. We have modified our text to clarify these points.

Thanks for the reply of the concern. I briefly checked the revision and have some more comments/suggestions/questions.

• First of all, notation in line 120 didn't change, it's still $L_t\in \mathbb{R}^n\rightarrow \mathbb{R}$. Please revise it. There is an obvious typo in line 226 as well, "NOw" shall be now. Please carefully check your revision.

• I appreciate author's clarification regarding their theorems. Nevertheless, several claims regarding their theorems might need further clarification. To be more specific, I would raise the score if the following questions can be answered appropriately:

  • The "converges" question regarding in theorem 1 is answered by the claim that "With $C$ close to 1 and $\epsilon$ small", but I don't see why it is true that you can choose $C$ close to 1 and $\epsilon$ can be small enough, without defining $D$. This claim needs to be checked by defining a metric $D$ and justify by some more evidence, it can be either theoretical analysis or numerical validation, but not statement like "we believe ......". Why Eq (5) (6) holds in general is also a myth to me.
  • I think the similar question arises when checking theorem 2. Namely, by using either CE loss or RL task loss as an example, explain what is the estimated value of $k$ and $\beta$ stated in Eq. (10) (11) that make the condition holds. It will help all readers understand the significance of the bound stated in Eq. (12). Another way is to explain the claim that PL condition + GD scheme automatically satisfy the condition in Eq. (10)(11), by extending the claim as a solid example to enlighten readers that these two conditions are easily met in general learning task.
  • The claim "We may expect that this exact equivalence is relaxed when the loss functions become nonlinear, but it still holds approximately to the extent that loss functions are locally linear." needs to be checked under experiment setups in Section 4. Namely, why loss in experiments are locally linear? Another clarification is that why Eq. (26) lead to " approximating Bayes optimal solutions". There is no Bayes optimization problem even being defined.

Thank you for your timely and thoughtful follow-up and for bringing these important points to our attention. We appreciate your willingness to reconsider your evaluation of our paper. We have carefully addressed each of your concerns below and made corresponding revisions to the manuscript to enhance clarity and rigor.

Typos and Notation Corrections

Thank you for pointing out these oversights.

Line 120: We apologize for the confusion. The notation indicates that the loss function maps from the parameter space $\mathbb{R}^d$ to the real numbers $\mathbb{R}$, which we have explained in our latest revision. We are happy to update the manuscript if the reviewer suggests another notation for this.

Line 226: Thank you for catching the typo; we have fixed this in our latest revision.

Clarification on Theorem 1 and Equations (5) and (6)

We appreciate the need for a concrete definition of $D$ and justification for the constants $C$ and $\epsilon$. Here's a detailed clarification:

Defining the Discrepancy Measure $D$: We may define $D(p, q)$ as the Wasserstein distance (specifically, the 2-Wasserstein distance) between two probability distributions $p$ and $q$:

$D(p, q) = \left( \inf_{\gamma \in \Gamma(p, q)} \int_{\mathbb{R}^d \times \mathbb{R}^d} \|x - y\|^2 d\gamma(x, y) \right)^{1/2}$ where $\Gamma(p, q)$ is the set of all couplings of $p$ and $q$. The Wasserstein distance is appropriate for continuous distributions and provides a meaningful measure of discrepancy that satisfies non-negativity, symmetry, the triangle inequality, and $D(p, p) = 0$.

Justifying $C \approx 1$ and Small $\epsilon$:

• Value of $C$: In Equation (6), $C$ represents how the discrepancy between two distributions changes after an update. For our gradient-based particle filter with small step sizes $\sigma^2$, the updates are gentle, and the movement of particles is limited. Specifically, for smooth loss functions $L$ with Lipschitz continuous gradients (i.e., $|\nabla L(x) - \nabla L(y)\| \leq L_{\nabla} \|x - y\|$), the gradient descent update is a Lipschitz continuous mapping with Lipschitz constant $L_\nabla$. This supports $C \approx 1 + \sigma^2 L_{\nabla}$​, which approaches $1$ as $\sigma^2$ goes to $0$.

• Value of $\epsilon$: The term $\epsilon$ captures the approximation error between the particle filter update and the true Bayesian update. In our gradient-based particle filter, the error due to linear approximation of $L$ can be bounded using Taylor's theorem. For twice-differentiable functions, the second-order remainder term involves the Hessian, and with small $\sigma^2$, this term becomes negligible.

When Equations (5) and (6) Hold:

• Equation (5): This is a property of the discrepancy measure. If we choose it to be Wasserstein distance, for example, then it holds automatically.
• Equation (6): This inequality reflects the stability of the particle filter update with respect to the initial discrepancy between $\hat{p}$​ and $\hat{q}$​. It holds under the assumption that the update operator is Lipschitz continuous in the Wasserstein distance. For our particle filter, this is justified by the small gradient steps and the smoothness of the loss function as we explain above.

Clarification on Theorem 2 and Conditions in Equations (10) and (11)

We appreciate the suggestion to use a solid example to illustrate these assumptions. We will use Wasserstein distance as our discrepancy metric $D$ for this example. Consider a loss function $L$ satisfying the PL condition and Lipschitz continuity of $L$ and its gradient:

$||\nabla L(x)||^2 \geq \mu L(x)$

$|L(x) - L(y)| \leq \ell ||x -y||$

$||\nabla L(x) - \nabla L(y)|| \leq L_{\nabla} ||x - y||$

where $\mu$, $\ell$, and $L_\nabla$ are constants, and the minimum value of the loss is $0$. Recall that the PL condition and Lipschitz continuity of the gradient of $L$ guarantees that gradient descent steps of learning rate $\eta$ reduce the loss by a factor of $1 - (\eta - \frac{L \eta^2}{2}) \mu$. Assuming our update rule is gradient descent, we then have:

Equation (10):

$\mathbb{E}_{x \sim p}[L(x)]$

$-\mathbb{E}_{x \sim q}[L(x)]$

$\leq \ell D(p, q)$

Equation (11):

$\mathbb{E}_{x \sim \hat p[L]}[L(x)]$

$\leq (1 - (\eta - \frac{L \eta^2}{2}) \mu) \mathbb{E}_{x \sim \hat p}[L(x)]$

On Loss Functions Being Locally Linear

We note that loss functions are locally linear in terms of the model parameters under two conditions:

• The loss as a function of the model output is differentiable

• The model output as a function of the model parameters is differentiable

In our experiments, we use losses and models that are differentiable everywhere (or almost everywhere in the case of ReLU-activated models). Thus, our loss functions are locally linear.

Thank you for the reply. However, these justifications does not fully resolve the concern.

From the setup you provide, $C\approx 1 +\sigma^2 L_{\nabla}$ is a claim one needs to prove as a lemma or theorem. On top of my head it is vague and I have no idea why it is an obvious conclusion, especially with Wasserstein distance setup. If possible, can you provide some evidence in previous works to support your claim? Justification of $\epsilon$ is fine, but it requires a clean statement in manuscript to reflect that this is a valid assumption. This assumption supports reviewer eGfq's concern that, if your learning time is long, say $T\gg 1$, then $(1+\sigma^2)^{T-2}$, with fixed sigma, is a diverging constant, it might give you a bound which fixed gradient descent scheme can achieve easily as well (think fixed scheme as dirac-delta density). How to reconcile the concern?

Moreover, I noticed authors might mix up the probability density and a fixed scheme. Explaining why Eq.(10) (11) holds is non-trivial: by stacking up the constant, the author implicitly assumed fixed gradient descent scheme, whereas the energy model $\exp(-L(x))$, which reweigh particles, disappears when computing expectation. From my point of view, this is an inconsistency of presentation, I wonder if authors can explain further, whether or not $E_{\hat{p}[L]}[L(x)] = \int_x \frac{L(x)e^{-L(x)}}{Z} dx$, with $Z$ partition function (same as normalization constant), and if this is the correct estimation as I understand, how eq (10)(11) hold in this case.

For locally linear assumption, I presume authors are claiming "locally differentiable". This part is fine but I wonder if there are some ways of validating the theorem 3 even if first order Taylor expansion is an approximation.

I appreciate the justification from authors, and I would recommend authors state the fact in a more convinced manner by citing previous works. I wonder if Hoeting et. al. has justification of these statements or consider citation of some textbook like [1] which partially support your claim.

[1]: Simo Srkk. 2013. Bayesian Filtering and Smoothing. Cambridge University Press, USA.

Thank you for your continued engagement with our paper and for your insightful comments. We appreciate your suggestions and have worked diligently to address each of your concerns. Below, we provide detailed responses and have made corresponding revisions to the manuscript to enhance clarity and rigor.

Approximating Bayes Optimal Solutions

Thank you for this suggestion. In the revised manuscript, we have more carefully explained how Theorem 3 implies that in the limit of infinitely many particles, the particle filter exactly matches the posterior distribution and have cited prior work in support of this (Doucet et al.).

Clarification on $C \approx 1 + \sigma^2 L_{\nabla}$ and the Wasserstein Distance

We appreciate the need for a formal justification. We are happy to provide a formal lemma here, and if satisfactory, can add it to our manuscript:

Lemma: Let $D$ be the 2-Wasserstein distance, and suppose the loss function $L$ has a Lipschitz continuous gradient with Lipschitz constant $L_{\nabla}$​. Then, for the gradient descent update $x_{t+1} = x_t - \sigma^2 \nabla L(x_t)$, the following holds:

$D(\hat{p}[L], \hat{q}[L]) \leq (1 + \sigma^2 L_{\nabla}) D(\hat{p}, \hat{q})$.

Proof Sketch:

• The gradient descent update map $\phi(x) = x - \sigma^2 \nabla L(x)$ is Lipschitz continuous with Lipschitz constant $\text{Lip}(\phi) = 1 + \sigma^2 L_{\nabla}$.
• Using properties of the Wasserstein distance under Lipschitz mappings (Santambroglio, 2015), we have:

$D(\hat{p}[L], \hat{q}[L]) = D(\hat{p} \circ \phi, \hat{q} \circ \phi) \leq \text{Lip}(\phi) D(\hat{p}, \hat{q})D(p^​[L],q^​[L])$

We have also indicated the error scaling of all approximations in our particle filter derivation, and indicated the error we would expect in the final result as a function of $\sigma^2$.

Exponential Dependence on $T$

You express concern that with $T \gg 1$, the term $(1 + \sigma^2 L_{\nabla})^{T-2}$ may diverge, making the bound less practically useful. We acknowledge that exponential growth with respect to $T$ can be problematic. However, with small $\sigma^2$ and bounded $L_{\nabla}$​, $C^{T-2} = (1 + \sigma^2 L_{\nabla})^{T-2} \approx 1 + (T-2) \sigma^2 L_{\nabla})$, which is only a linear dependence on $T$. Note that this approximation requires $\sigma^2$ to approach zero faster than $T$ approaches infinity.

If $\sigma^2$ is fixed and $T$ approaches $\infty$, then we don't believe it is possible to remove the exponential dependence of the bound on $T$ without stronger assumptions.

Clarification on Equations (10) and (11)

If we understand correctly, you are asking us to justify Equations (10) and (11) in the case that our update rule is not gradient descent, but rather our particle filter update. We justify this below.

Note that our particle filter has two update steps: 1) the gradient descent update, 2) the weight updates of each particle. We make the same assumptions about $L$ as before:

$||\nabla L(x)||^2 \geq \mu L(x)$

$|L(x) - L(y)| \leq \ell ||x -y||$

$||\nabla L(x) - \nabla L(y)|| \leq L_{\nabla} ||x - y||$

where $\mu$, $\ell$, and $L_\nabla$ are constants, and the minimum value of the loss is $0$. We denote time $t$ as pre-update and time $t+1$ as post-update.

Equation (10) holds as before:

$\mathbb{E}_{x \sim p}[L(x)]$

$- \mathbb{E}_{x \sim q}[L(x)]$

$\leq \ell D(p, q)$

Equation (11):

We may first write,

$\mathbb{E}_{x \sim \hat{p}[L]}[L(x)]$

$= \sum_i L(x^i_{t+1})$

Since $x^i_{t+1}$ is found by taking a gradient step on $x^i_t$, we have:

$L(x^i_{t+1}) \leq (1 - (\eta - \frac{L \eta^2}{2}) \mu) L(x^i_t)$

by the same argument from earlier. Summing over $i$ on both sides, we have:

$\mathbb{E}_{x \sim \hat{p}[L]}[L(x)]$

$\leq (1 - (\eta - \frac{L \eta^2}{2}) \mu) \mathbb{E}_{x \sim \hat{p}}[L(x)]$

Validating Theorem 3 with First-Order Taylor Approximation

To clarify, we believe Theorem 3 holds approximately (within some error bounds) when loss functions are differentiable (and thus locally linear).

If we understand correctly, you are asking whether Theorem 3 remains correct when we consider higher order terms: when perturbations to $x$ are large enough that $L$ changes non-linearly.

Under our analytical approach, which relies on local linearity of the loss function, we believe Theorem 3 would not remain true; in fact, the derivation of our particle filter itself relies on local linearity. On the other hand, we note that with sufficiently small step sizes (i.e. $\sigma^2$ near 0), any differentiable loss function can be treated as locally linear. Thus, we expect that the Bayesian properties of our particle filter break down with large step sizes.

We hope that these detailed explanations and revisions address your concerns. Your feedback has significantly improved the clarity and rigor of our paper. We are committed to ensuring that our work is both theoretically sound and practically relevant. Thank you again for your thoughtful consideration.

Thank you for your constructive feedback and for acknowledging the strengths of our paper. We address your comments below.

Limitations Section

We appreciate the suggestion to discuss limitations. We have now added limitations to our conclusion section to produce a unified discussion section. We highlight two limitations:

Maintaining multiple particles increases memory usage linearly with $N$. We acknowledge this as a limitation, especially with large models. One possible workaround is to train multiple models in series, although this trades the memory cost with a time cost.

Another limitation is that our method alone is often not as effective as in combination with other methods. This limits its use as a stand-alone algorithm, and we instead advocate using it as an ensembling strategy to enhance the performance of an existing method (to address continual learning or loss-of-plasticity, for example).

Accessibility of Writing

Thank you for your suggestions! We have made the specific changes you suggested and are happy to make any further modifications to improve the clarity of our work.

Permutation Invariance Claim

Thank you for this point; we have modified our abstract as you suggested.

Table Formatting

Thank you! We have now bolded the best-performing numbers in Table 1 for clarity.

Comparison with Larger Models

Thank you for this suggestion. We are conducting experiments to compare our method with larger single models and will include the results in the revised manuscript once available.

Thank you for your insightful review and for highlighting both the strengths and areas for improvement in our paper. We address your concerns point by point.

Clarification on Equations (6) and (7)

Why do Equations (6) and (7) hold for suitable constants?

Equation (6): This equation states that the discrepancy between the updated particle filter distributions $D(\hat{p}[L], \hat{q}[L])$ is bounded by $C$ times the discrepancy between the initial distributions. Intuitively, this means that if two particle filters start off similar, they remain similar after an update. This holds under the assumption that the particle filter update is Lipschitz continuous with respect to the discrepancy measure D. In practice, this is reasonable for particle filters where updates involve smooth operations like weighting and resampling.

Equation (7): This equation bounds the discrepancy between the updated particle filter distribution $\hat{p}[L]$ and the true Bayesian posterior $p(\cdot|L)$. It reflects the particle filter's approximation of the Bayesian update, where $\epsilon$ captures the error due to approximations (e.g., finite number of particles, linearization). This is typical in particle filter analyses, where each update aims to approximate the true posterior as closely as possible.

Exponential Bounds in T and Potential Vacuity

We acknowledge that the exponential dependence on $T$ may seem concerning. However, the growth rate is governed by the constant $C$, which represents the stability of the particle filter.

Interpretation of $C$: In practice, $C$ is typically close to 1. It measures how small variations in particles propagate through updates: in other words, the stability of updates. If $C \approx 1$, then small fluctuations in the initialization of particles does not significantly affect the outcome after training. We expect this to be a reasonable assumption for many practical algorithms.

Non-Vacuous Bounds: In Theorem 2, the loss is upper bounded by $\beta M$ plus a term involving $C$ and $\epsilon$. If $\epsilon$ is small and $C$ is close to 1, the additional term remains small compared to $\beta M$. Thus, the bound is not vacuous and provides a meaningful guarantee on the loss.

Comparison with Linear Bounds: While sublinear regret bounds from online optimization are tighter, our analysis focuses on the behavior of particle filters approximating Bayesian updates, which is a different setting. The exponential term arises due to the recursive nature of the discrepancy bound. Under our assumptions, these are the tightest bounds we can obtain.

Approximate Bayesian Properties of Our Solution

It is unclear why the approximate solution in (15) satisfies the Bayesian properties of particle filters.

Our gradient-based particle filter is designed to approximate the Bayesian update efficiently in high-dimensional spaces. Main Approximation in Equation (15): We use a Gaussian approximation around each particle and linearize the loss function. This is exact in the limit as the Gaussian variance $\sigma^2$ approaches zero, effectively capturing the local behavior of the loss function. Theoretical Justification (Theorem 3): Although Theorem 3 considers linear loss functions, it demonstrates that in this setting, our particle filter exactly maintains the correct posterior ratios between particles, preserving Bayesian properties. We may expect that this exact equivalence is relaxed when the loss functions become nonlinear, but still hold approximately to the extent that loss functions are locally linear.

Empirical Evidence: Our experiments show that our method achieves permutation invariance and mitigates catastrophic forgetting, consistent with Bayesian approaches. Additionally, our method differs from independent gradient descent models because particle weights are updated based on their likelihoods, reflecting the posterior distribution.

Clarification on Notation

Can you explain the notation $\hat{p}[L]$?

Yes, as we explain in Section 3.1, $\hat{p}[L]$ denotes the updated particle filter after processing the loss function $L$. Starting from the particle distribution $\hat{p}$, the update incorporates the information from $L$ to produce a new distribution $\hat{p}[L]$.

Objective of the Particle Filter

Our particle filter aims to approximate the Bayesian posterior distribution $p_T(x)$, which is proportional to $p_0(x) e^{-\sum_{t=1}^T L_t(x)}$. While the posterior places higher probability density near the global minimizer of the total loss, it represents the full distribution over model parameters. The particles collectively capture this distribution, including uncertainty and multiple modes, rather than just estimating the global minimizer.

Thank you for your positive feedback and encouraging comments. We are glad that you found our method theoretically principled and practically effective. We address your questions below.

Relation to Ensemble Methods in Continual Learning

Thank you for pointing out relevant works on ensembles in continual learning [1-3]. We have added a discussion in the related work section.

While these works employ ensembles to mitigate forgetting, our method differs by providing a Bayesian framework with theoretical guarantees like permutation invariance. Additionally, we focus on approximating the posterior distribution over model parameters rather than just maintaining multiple models.

Connection Between Sections 3.2-3.3 and 3.4

The particle filter proposed in Section 3.4 does satisfy the conditions of Section 3.2-3.3 given three key additional assumptions:

1. The gradient descent step of the algorithm is sufficiently stable such that equation 6 holds; small changes to the initialization of particles make only small changes to the trained particles

2. The step size $\sigma^2$ is small enough; this makes the assumption of the local linearity of the loss valid in Section 3.4

3. The task loss function $L$ satisfies the conditions of Theorem 2; these are conditions specifying how smooth and optimizable the loss function is

The specific constants to make these equivalences are highly dependent on the task loss function; thus, it is difficult to define a general set of constants that would hold these assumptions hold. Nevertheless, we expect that for any particular set of tasks, it is possible to define constants such that our proposed particle filter fits the assumptions in Sections 3.2-3.3.

Future Work Directions

We believe the one exciting direction for future work is applying our approach to machine forgetting. Machine forgetting is a challenging problem in which the goal is to modify a trained model so as to remove the effect of a particular training input (or set of inputs) on the trained model. With typical gradient-based learning, this is challenging because a single training input can completely change the course of a model's optimization trajectory, making it difficult to undo the effect of that input.

Permutation-invariant algorithms on the other hand can easily forget arbitrary past inputs because any training input can be considered the "last" training datapoint by permuting the training points. Forgetting then simply involves undoing the effect of the last learning step.

Equality in Line 256

Thank you for pointing this out. The equality on line 256 is approximate under linear approximation of the loss. In defining $w_{t+1}^{(i)}$​, we use the exact values computed from the loss function on the new point; thus, equation 23 is exact. We have updated our notation in the manuscript to reflect this.

Table Formatting

Thank you for the suggestion. We have reformatted Table 1 to clearly distinguish between continual learning baselines and those combined with our particle filter.