Paper: Implicit Maximum a Posteriori Filtering via Adaptive Optimization (ICLR-2024)

Reviewer comment (Theory):
**From linear to nonlinear problems:** For linear problems, Santos (1996) shows that a choice of a (truncated) gradient descent corresponds to an implicit choice of predictive covariances. Thus, for linear problems, estimating the filtering-mean by applying a suitably truncated version of gradient descent "makes sense" -- provided the prior model implicit to the length-scale in gradient descent "makes sense". Using the same idea for nonlinear problems, arbitrary optimisers, and possibly optimisers with momentum is, technically speaking, a leap of faith, and there is nothing to support this transition other than "it seems to work well on the test problems". Fortunately, it does seem to work well on the test problems. (Unless one chooses to use AdaDelta, that is.)

Author response:
- We agree that there is some intuition that is lost when using adaptive optimizers over nonlinear problems. I think we can remedy this issue (in part) by being more clear about how the number of optimizer steps, choice of optimizer, and choice of initialization affect the “implicit” prior. We will adjust section 4.2 to improve clarity.

Reviewer comment (Theory):
It would be good if the authors could provide a bit more detail and intuition in the second half of Section 3.1 and in Section 3.2 about the role of the matrices \( B_t \) and \( C_t \) and why we can be sure that these are well defined.

Author response:
- We will revise 3.1 and 3.2 to include your suggestions and fix minor typos.

Reviewer comment (Writing):
It would be good if the authors could provide a bit more detail and intuition in the second half of Section 3.1 and in Section 3.2 about the learning-rate matrix, perhaps by stating the plain gradient-descent update equation which includes this matrix.

Author response:
- We will revise 3.1 and 3.2 to include your suggestions and fix minor typos.

Reviewer comment (Writing):
There are a few minor typos, e.g., the notation \( \mathbb{E}[p(x_t|y_{1:t-1})] \) does not make sense and presumably \( \mathbb{E}[x_t|y_{1:t-1}] \) is meant.

Author response:
- We will revise 3.1 and 3.2 to include your suggestions and fix minor typos.

Reviewer comment (Experiments):
Some of the "true" model parameters (especially the process and observation noise) in the examples from Sections 5.1 and 5.2 seem to be slightly different than those used in the references cited for these applications. The reviewer asks if there is a particular reason for this change and whether the results in these sections would differ qualitatively if the parameters from the literature were used.

Author response:
- The Lorenz Attractor system was specified such that we could show some level of robustness to model parameters. The Lorenz Attractor follows the most classical setup for the differential equation parameters (sigma = 10, beta = 8/3, and rho = 28), which is known to induce chaotic behavior. We chose the initial starting distribution such that we can sample a different behavior with every starting seed (see Appendix A.5.2). By ensuring that every seed essentially produces a different filtering problem, we can get a good measure of efficacy across 100 starting initializations. 
- We selected the Lorenz system and the current Lorenz configuration to show that the method is robust to many different filtering problems! We found that other configurations (in literature) were less “chaotic” and varied paper to paper with no established standard.

Reviewer comment (Experiments):
The reviewer is willing to raise their score if the authors can show that the results in Sections 5.1–5.2 are robust to the choice of model parameters.

Author response:
- We are happy to delegate the current table for Experiment 1 to the Appendix and replace it with results where the true process noise Q = 1, 3, and 5 (instead of showing only Q = 3). This would show the reader performance for each baseline and the proposed method in low, medium, and high process noise cases. As stated above, the Lorenz system already naturally emits many different filtering problems (addressing Section 5.2) and this change would address robustness to model parameters in the 5.1 system. Would this revision increase your assessment of the paper?

Thank you again for your feedback, especially the point about robustness to true model parameters. We are adjusting 5.1 to address this and we believe 5.2 already accomplishes this. We look forward to discussion!

- We address all of your concerns in the new versions. Most importantly, the newest version tests 9 systems for the toy nonlinear system instead of 1 and 4 systems for the stochastic lorenz attractor instead of 1. These experiments cover the common range used in literature for running these systems and should address your robustness concern. Results remain consistent.

Reviewer comment (Novelty):
The authors should note some of the modern KF based research in this space (Chang et al. 2023 (https://arxiv.org/abs/2305.19535), Titsias et al. 2023 (https://arxiv.org/abs/2306.08448)) that successfully works on larger-scale problems than the ones considered in this paper.

Author response:
- Thank you for referencing two literature sources that should be mentioned in the related works. Similar to the Auvinen et al. (2010) and Bardsley et al. (2013) papers (mentioned in related works of current version), Chang et al. (2023) constructs a low-rank EKF. However, these low-rank filtering methods are not straightforward to implement (relative to proposed method) and are sensitive to the convexity of the loss landscape (most deep learning libraries have all the optimizers we tested and an implicit MAP filter can be implemented in minutes). We believe that Chang et al. (2023) is outside the scope of extremely straightforward large-scale filtering.
- In the yearbook example, our objective was to test an explicit approach (VKF) and our implicit approach on equal standing of simplicity. 
- Titsias et al. (2023) focuses on dealing with the non-stationarity of the convex readout layer in a continual learning setting, using a simple formulation of a Kalman filter, and focusing more on the detection of transitions. This puts Titsias et al. (2023) slightly outside the scope of our paper because it is not explicitly designed for evolving every weight of a pre-trained network.

Thank you again for the feedback, especially the references to Chang et al. (2023) and Titsias et al. (2023). We are happy to discuss any of these points!

Reviewer comment (Experiments):
The experiments chosen are somewhat small-scale given that the authors' argument for not having access to estimates of uncertainty is having a scalable method.

Author response:
- Since we are using optimizers that are deployed at massive scale for DNNs, we focused on verifying that the proposed method can operate as a Bayesian filter.
- Since we are comparing to classical Bayesian filters that do not scale like the proposed method, the choice of neural network for the yearbook dataset cannot be too large.
- We hope these two points clarify the choice of scale.

Reviewer comment (Experiments):
Since the authors explicitly connect the choice of optimizer to defining a prior, they should for example compare different optimizers to one another on the same problem. Have you tried doing this?

Author response:
- We address your concerns about modifying the number of steps K in Tables 5, 7 in the Appendix. The optimizers are not robust to large changes in K, since this partially encodes the balance between process noise and measurement noise! We also compare different optimizers in Tables 4, 6 in the Appendix. 
- We are hoping that Tables 4-7 include all of the desired information and fully address your concerns!

Reviewer comment (Evaluation):
Table 1 states that the MAP methods are not dependent on the process noise covariance Q_t-1, which is true, but they are dependent on the number of steps K, which as stated earlier, in e.g., 3.3, implicitly defines a prior. The authors however state that there is no sensitivity to missspecification (last paragraph of Page 6). This should be clarified.

Author response:
- We agree that we can clarify “misspecification” for sections 5.1 and 5.2. We will make 5.1 more sound and add results to 5.2 that clarify our meaning of “misspecification”.

Reviewer comment (Experiments):
In practical contexts, we often don't have access to a pretrained model. How would the method perform when the model is trained completely from scratch?

Author response:
- You are right that pre-trained models are not always available. However, without a pre-trained model, the adaptation would be relatively simple. In an explicit Bayesian filter (ie EKF), one would assume a very large prior covariance if there is no starting information. This is equivalent to, in our implicit formulation, optimization to convergence (because the measurement noise << prior, so you just fit the likelihood). We opted against this style of running this experiment because there would be difficulties in running the particle filter baseline if the transition distribution is not already centered on a relatively good parameterization. 
- We will clarify in Section 5.3 that we can run this implicit MAP filter easily without a pre-trained model by fitting the likelihood exactly at the first timestep.