SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Importance Weighted Sample Gradient Estimator. We propose a novel alternative method for computing gradients of samples from a continuous mixture distribution. Our importance weighted sample gradient (IWSG) estimator employs importance sampling for unbiased gradient estimation. IWSG is related to the DIS-PF gradient estimator for discrete resampling [6], but we instead consider continuous mixture distributions with arbitrary component distributions. This resolves the ancestry problem present in Scibior et al. [6], since particle locations could be sampled from multiple overlapping mixture components, whose parameters will all have non-zero gradients.
Formally, we wish to draw samples  $z^{(i)} \sim m(z \mid \phi)$  from a mixture with parameters  $\phi$ . Instead of sampling from  $m(z \mid \phi)$  directly, IWSG samples from some proposal distribution  $z^{(i)} \sim q(z)$ . Importance weights  $w^{(i)}$ , and associated gradients, for these samples then equal
$$
w ^ {(i)} = \frac {m \left(z ^ {(i)} \mid \phi\right)}{q \left(z ^ {(i)}\right)}, \quad \nabla_ {\phi} w ^ {(i)} = \frac {\nabla_ {\phi} m \left(z ^ {(i)} \mid \phi\right)}{q \left(z ^ {(i)}\right)}. \tag {14}
$$
Setting the proposal distribution  $q(z) = m(z \mid \phi_0)$  to be a mixture model with the "current" parameters  $\phi_0$  at which gradients are being evaluated, the IWSG gradient estimator (14) becomes
$$
w ^ {(i)} = \frac {m (z ^ {(i)} \mid \phi) \big | _ {\phi = \phi_ {0}}}{m (z ^ {(i)} \mid \phi_ {0})} = 1, \quad \nabla_ {\phi} w ^ {(i)} = \frac {\left. \nabla_ {\phi} m (z ^ {(i)} \mid \phi) \right| _ {\phi = \phi_ {0}}}{m (z ^ {(i)} \mid \phi_ {0})}. \tag {15}
$$
While current samples are given weight one because they are exactly sampled from a mixture with the current parameters  $\phi_0$ , gradients account for how importance weights will change as mixture parameters  $\phi$  deviate from  $\phi_0$ . Note that gradients are not taken with respect to the proposal  $q(z) = m(z \mid \phi_0)$  in the denominator of  $w^{(i)}$ , since sample locations  $z^{(i)}$  are not altered by gradient updates; changes in the associated importance weights are sufficient for unbiased gradient estimation.
Instability of Implicit Reparameterization Gradients. Reparameterization methods capture changes to a given distribution  $p(z|\phi)$ , and thus to its parameters  $\phi$ , by shifting the samples drawn from that distribution. When  $p(z|\phi)$  is unimodal, this induces a smooth shift in the samples. When  $p(z|\phi)$  contains multiple non-overlapping modes, these shifts are no longer smooth as samples
Figure 4: A simple temporal prediction problem where IRG has highly unstable gradient estimates. We use  $N = 25$  particles when training IRG-PF, IWSG-PF, and truncated gradients (TG-PF). We show the learned values for parameters  $B$ ,  $w_{2}$ ,  $C_{1}$  and their gradients during training. IRG produces unstable gradients which prevent parameters from converging at all, but IWSG allows for smooth convergence of all parameters. IWSG is faster than biased TG, and nearly as effective as (expensive, and for general models intractable) mixture KF.
jump from one mode to another; see Fig. 3. For mixtures with finitely supported kernels (like the Epanechnikov) and non-overlapping modes, discontinuities exist in  $S_{\phi}^{-1}(\epsilon)$ , and IRG estimates may have infinite variance. These discontinuities correspond to samples shifting from one mode to another, and the resulting non-smooth sample perturbations will destabilize backpropagation algorithms for gradient estimation. While the inverse CDF is always continuous for Gaussian mixtures, it may still have near-infinite slope when modes are widely separated, and induce similar IRG instabilities.
In contrast, our IWsG estimator reweights particles instead of shifting them. This reweighting is smooth and does not suffer from any discontinuities. To illustrate this, we create a simple discriminative PF with linear-Gaussian dynamics (16) and a bimodal observation likelihood (17):
$$
p \left(x _ {t} \mid x _ {t - 1}, a _ {t}\right) = \operatorname {N o r m} \left(x _ {t} \mid A x _ {t - 1} + B a _ {t}, \sigma^ {2}\right), \tag {16}
$$
$$
p \left(y _ {t} \mid x _ {t}\right) = w _ {1} \operatorname {N o r m} \left(y _ {t} \mid C _ {1} x _ {t} + c _ {1}, \gamma^ {2}\right) + w _ {2} \operatorname {N o r m} \left(y _ {t} \mid C _ {2} x _ {t} + c _ {2}, \gamma^ {2}\right). \tag {17}
$$
Exact posterior inference for this model is possible via a mixture KF [11, 40] which represents the state posterior as a Gaussian mixture. The number of mixture components grows exponentially with time, motivating PFs for long-term tracking, but remains tractable over a few time steps.
Using a fixed dataset and gradient descent, we train discriminative PFs using the IRG estimator, our IWSG estimator, as well as with biased gradients that are truncated at each resampling step [5]. We initialize with perturbed parameter values and aim to learn the values of all dynamics and likelihood parameters, except for  $\sigma$  and  $\gamma$  which are fixed to their true values. See Appendix for details.
Fig. 4 shows the estimates and gradients for a subset of the learned parameters, and several gradient estimators. IRG's unstable gradients prevent convergence of all parameters. For training (generative) marginal PFs [41] which approximate marginals with mixtures, Lai et al. [42] found that IRG gradient variance was so large that biased estimators converged faster, but provided no detailed analysis. In contrast, IWSG converges smoothly for all parameters, and more rapidly than truncated gradients.
5 Mixture Density Particle Filters
Our mixture density particle filter (MDPF) is a discriminative model inspired by regularized PFs, where resampling uses KDEs centered on the current particles. We apply IWSG to this mixture resampling, achieving unbiased and low-variance gradient estimates for our fully-differentiable PF.
MDPF does not incorporate human-specified dynamics or measurement models; some prior work assumed known dynamics to simplify learning [5, 7]. These models are parameterized as deep neural networks (NNs), and trained via stochastic gradient descent to minimize the negative-log-likelihood of (potentially sparsely labeled) states  $x_{t}$ , given observations  $y_{t}$  and (optionally) actions  $a_{t}$ . As shown in Fig. 2, dynamics and measurement models are flexible composed from smaller NNs. Our MDPF places no constraints on the functional form of NNs, allowing for modern architectures such as CNNs [43] and Spatial Transformers [44] to be used. See Appendix for implementation details.
The MDPF uses KDE mixture distributions as in Eq. (4) to define state posteriors used to evaluate the training loss, as well as resampling. A bandwidth parameter  $\beta$  is thus required for each state dimension. Rather than setting  $\beta$  via the classic heuristics discussed in Sec. 2.2, we make  $\beta$  a learnable parameter, optimizing it using end-to-end learning along with the dynamics and measurement models. This contrasts with prior work including SR-PF [7] and C-PF and OT-PF [8], which all include relaxation hyperparameters that must be tuned via multiple (potentially expensive) training trials.
We also extend the MDPF by decoupling the mixtures used for particle resampling and posterior state estimation. Using two separate measurement models, we compute two sets of particle weights  $\tilde{w}_t^{(:)}$
Figure 5: Box plots showing median (red line), inter-quartile (colored box) and range (whiskers) over several training runs on the Bearings-Only (11 runs) and Deepmind-Maze (5 runs) tracking tasks. MDPF and A-MDPF consistently perform well on both the NLL and RMSE metrics. Our IWSG estimator is critical: IRG-MDPF performs very poorly due to unstable gradients, TG-MDPF is inconsistent and sometimes becomes trapped in local optima, and baselines with biased gradient estimators typically have inferior performance. LSTM achieves low RMSE in the Deepmind-Maze task however we find it simply blind-propagates using the noisy actions. DIS-PF performs well for Maze 3 but with larger variability which can sometimes lead to worse performance than our more stable A-MDPF. Note that IRG-MDPF does not support Epanechnikov kernels, which induce discontinuous CDFs.
and  $w_{t}^{(:)}$ , one for resampling and the other for posterior state estimation; see Fig. 2. Each distribution is also given a separate bandwidth  $\tilde{\beta}, \beta$ . This adaptive mixture density PF (A-MDPF) allows the uncertainty in the estimation of the current latent state, and the degree of exploration in the particle resampling step, to be decoupled and separately optimized during training. Note that training of separate posterior and resampling distributions is impossible for PFs that truncate temporal gradients; it is only feasible due to our unbiased and low-variance IWSG estimator of resampling gradients.
We evaluate our MDPF and A-MDPF on a variety state estimation tasks in complex, but simulated, environments. We compare to TG-PF [5], SR-PF [7], OT-PF [8], DIS-PF [6], as well as C-PF which uses the Concrete distribution [32, 31] to relax discrete resampling. We also compare to variants of MDPF: Truncated-Gradient-MDPF (TG-MDPF) is similar but stops gradients at each particle resampling, and IRG-MDPF replaces our IWSG estimator with the unstable IRG estimator. Finally, we compare to a LSTM [45] which produces point predictions via uninterpretable latent states.
For all tasks, we estimate posterior distributions of a 3D (translation and angle) state of a robot,  $x_{t} = (x,y,\theta)$ . We use Gaussian kernels for position components, and von Mises kernels for angular components, of all KDEs. As is common in recurrent neural network (RNN) training, we employ conservative gradient clipping [46], as well as truncated-BPTT [47] where gradients are propagated for 4 time-steps. All methods are initialized as the true state with added noise.
We train and evaluate MDPFs to minimize the following negative-log-likelihood (NLL) loss:
$$
\mathcal {L} _ {\mathrm {N L L}} = \frac {1}{| \mathcal {T} |} \sum_ {t \in \mathcal {T}} - \log \left(m \left(x _ {t} \mid x _ {t} ^ {(:)}, w _ {t} ^ {(:)}, \beta\right)\right). \tag {18}
$$
We use sparsely labeled training data to highlight the need for effective multi-time-step gradients during training. Dense labeling of data is often expensive, so many real-world datasets have similarly sparse labels. During training, we label true states every 4th time-step  $(\mathcal{T} = \{4,8,12,\ldots \})$ , and use densely labeled datasets  $(\mathcal{T} = \{1,2,3,\dots \})$  for evaluation. See the Appendix for details.
After training to minimize NLL, we optionally refine our MDPFs to minimize the mean-squared-error (MSE) of the weighted particle mean, and report the induced Root-MSE (RMSE). Baseline methods do not estimate continuous posteriors, and thus Eq. (18) cannot be used for training. Instead we train and evaluate these methods with MSE and RMSE, before freezing the dynamics and measurement models. Then we construct KDEs, and optimize a bandwidth to report NLL values for comparison.
**Bearings-Only Tracking.** In the bearings only tracking task, we track the state of a variable-velocity car using noisy bearing observations from a radar station to the car generated as
$$
y _ {t} \sim \alpha \cdot \operatorname {U n i f o r m} (- \pi , \pi) + (1 - \alpha) \cdot \operatorname {V o n M i s e s} (\psi (x _ {t}), \kappa),
$$
where  $\psi(x_{t})$  is the true bearing and  $\alpha = 0.15$ . Actions are not available. We use 5000 training and 1000 validation trajectories of length  $T = 17$ , and 5000 evaluation trajectories of length  $T = 150$ .
Figure 6: Example trajectories from the Bearings-Only tracking task. The radar station (green circle) produces bearing observations (green line) from the radar station to the car (red arrow). Note that observations are noisy and occasionally totally incorrect (see  $t = 18$ ). We plot the posterior distribution of the current state (blue cloud), the mean particle (yellow arrow), and the  $N = 25$  particles (black dots). A-MDPF is able to capture multiple modes in the state posterior (see  $t = 15, 17, 18, 19$ ) while successfully tracking the true state of the car. Other models (DIS-PF, OT-PF, C-PF) spread their posterior distribution due to poor tracking of the true state.
Figure 7: Example trajectories (rows), with observations in corner, from the House3D task using A-MDPF. Top: Initialized with  $N = 50$  particles set to the true state with moderate noise, the posterior (blue cloud) closely tracks the robot. Middle: Multi-modal initialization with  $N = 150$  particles (blue dots). Bottom: Naive initialization with  $N = 1000$  particles drawn uniformly. A-MDPF maintains multiple modes until enough observations have been seen to disambiguate the true state.
Figure 8: Example posteriors for three time steps (rows) of House3D tracking, for several discriminative PFs given the same low-noise initialization. We show the current true state and state history (red arrow and dashes), the posterior distribution of the current state (blue cloud, with darker being higher probability), and the estimated mean state and history (yellow arrow and dashes). Several baselines completely fail to track the robot.
Each method is evaluated and trained 11 times using  $N = 25$  particles. To further highlight the stability of our IWSG estimator, we also train each method (except for IRG-MDPF) with the Epanechnikov kernel replacing the Gaussian in the KDE. Fig. 5 shows statistics of performance across 11 training runs. The very poor performance of IRG-MDPF is due to unstable gradients; without aggressive gradient clipping, IRG-MDPF fails to optimize at all. The benefits of IWSG are highlighted by the inferior performance of TG-MDPF. Stable multi-time-step gradients clearly benefit learning and are enabled by IWSG. Qualitative results are shown in Fig. 6.
Deepmind-Maze Tracking. In the Deepmind-Maze tracking task, adapted from [5], we wish to track the 3D state of a robot as it moves through modified versions of the maze environments from Deepmind Lab [48] using noisy odometry as actions, and images (from the robot's perspective) as observations. We alter the experimental setup of [5] by increasing the noise of the actions by five times, and using  $N = 25$  particles for training and evaluation. We train and evaluate on trajectories from the 3 unique mazes separately, using 5000 trajectories of length  $T = 17$  from each maze for training, and 1000 trajectories of length  $T = 99$  for evaluation. We show results over multiple training and evaluation runs in Fig. 5, and include qualitative results in the Appendix.
House3D Tracking. This 3D state tracking task is adapted from [7], where a robot navigates single-level apartment environments from the SUNCG dataset [49] using the House3D simulator [50], with noisy odometry as action and RGB images as observations (see Fig. 7). Floor plans are input into measurement model via a modern Spatial Transformation Network architecture [44]. Training data consists of 74800 trajectories (length  $T = 24$ ) from 199 unique environments, with evaluation data being 820 trajectories (length  $T = 100$ ) from 47 previously unseen floorplans. We train and evaluate with  $N = 50$  particles.
Table 1: Mean  $\pm$  inter-quartile of evaluation metrics for all evaluation trajectories from the House3D task.
Method	NLL	RMSE
LSTM	12.54 ± 3.43	48.51 ± 27.16
TG-PF	15.50 ± 2.21	206.32 ± 109.44
OT-PF	14.87 ± 2.31	159.57 ± 109.65
SR-PF	14.92 ± 2.57	191.32 ± 129.28
DIS-PF	15.06 ± 2.02	212.58 ± 133.69
TG-MDPF	8.38 ± 2.53	35.77 ± 16.67
MDPF	8.18 ± 2.72	30.59 ± 12.98
A-MDPF	8.09 ± 3.49	30.51 ± 12.32
Due to the larger computational demands of House3D,
we train each method once, reporting results in Table 1. Interestingly we find that TG-PF, OT-PF, SR-PF, and DIS-PF are unable to learn useful dynamics or measurement models. The LSTM model achieves better performance, but looking deeper we discover that it ignores observations completely, and simply blindly propagates the estimated state using actions with good dynamics (i.e., dead reckoning). This strategy is effective for noise-free actions, but in reality the actions are noisy and thus the LSTMs estimated state diverges from the true state after a moderate number of time-steps. In contrast, our MDPF methods are able to learn good dynamics and measurement models that generalize well to the unseen environments.
We also investigate MDPFs capacity for multimodal tracking in Fig. 7, by initializing A-MDPF with particles spread throughout the state space. When initialized with random particles or with particles clustered around regions with similar observations, A-MDPF does not have sufficient information to collapse the posterior and thus maintains multiple posterior modes. As the robot moves, more evidence is collected, allowing A-MDPF to collapse the state posterior to fewer distinct modes, before eventually collapsing its estimate into one mode containing the true state. Note that such estimation of multiple posterior modes is impossible for standard RNNs like the LSTM.
Limitations. A known limitation of all PFs is there inability to scale to very high-dimensional states. Sparsity increases as the dimension grows, and thus more particles (and consequently computation) are required to maintain the expressiveness and accuracy of the state posterior. Our MDPF is not immune to this issue, but we conjecture that end-to-end training of discriminative PFs will allow particles to be more effectively allocated within large state spaces, scaling better than classic PFs.
We have developed a novel importance-sampling estimator for gradients of samples drawn from a continuous mixture model. Our results highlight fundamental flaws in the application of (implicit) reparameterization gradients to any mixture model with multiple modes. Applying our gradient estimator to the resampling step of a regularized discriminative PF, we obtain a fully end-to-end differentiable MDPF that robustly learns accurate models across multiple training runs, and has much greater accuracy than the biased PFs proposed in prior work. While our experiments have focused on synthetic tracking problems in complex virtual environments, future applications of our MDPF to real-world vision and robotics data are very promising.
We thank Alexander Ihler for helpful discussions, and Zain Farhat for his assistance with the LSTM baseline. This research supported in part by NSF Robust Intelligence Award No. IIS-1816365 and ONR Award No. N00014-23-1-2712.
[1] Neil J Gordon, David J Salmond, and Adrian FM Smith. Novel approach to nonlinear/non-gaussian bayesian state estimation. *IEE proceedings F-Radar and Signal Processing*, 1993.
[2] K. Kanazawa, D. Koller, and S. Russell. Stochastic simulation algorithms for dynamic probabilistic networks. In UAI 11, pages 346-351. Morgan Kaufmann, 1995.
[3] A. Doucet, N. de Freitas, and N. Gordon, editors. Sequential Monte Carlo Methods in Practice. Springer-Verlag, New York, 2001.
[4] M. S. Arulampalam, S. Maskell, N. Gordon, and T. Clapp. A tutorial on particle filters for online nonlinear/non-Gaussian Bayesian tracking. IEEE Trans. Signal Proc., 50(2):174-188, February 2002.
[5] Rico Jonschkowski, Divyam Rastogi, and Oliver Brock. Differentiable particle filters: End-to-end learning with algorithmic priors. Proceedings of Robotics: Science and Systems (RSS), 2018.
[6] Adam Scibior, Vaden Masrani, and Frank Wood. Differentiable particle filtering without modifying the forward pass. International Conference on Probabilistic Programming (PROBPROG), 2021.
[7] Peter Karkus, David Hsu, and Wee Sun Lee. Particle filter networks with application to visual localization. Conference on Robot Learning (CORL), 2018.
[8] Adrien Corenflos, James Thornton, George Deligiannidis, and Arnaud Doucet. Differentiable particle filtering via entropy-regularized optimal transport. International Conference on Machine Learning (ICML), 2021.
[9] Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. Transactions of the ASME-Journal of Basic Engineering, 1960.
[10] T. Kailath, A. H. Sayed, and B. Hassibi. Linear Estimation. Prentice Hall, New Jersey, 2000.
[11] Sebastian Thrun, Wolfram Burgard, and Dieter Fox. Probabilistic Robotics (Intelligent Robotics and Autonomous Agents). The MIT Press, 2005. ISBN 0262201623.
[12] S. J. Julier and J. K. Uhlmann. Unscented filtering and nonlinear estimation. Proc. IEEE, 92(3): 401-422, March 2004.
[13] Víctor Elvira, Luca Martino, and Christian Robert. Rethinking the effective sample size. International Statistical Review, 2022.
[14] Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In ICML, page II-1791-1799, 2014.
[15] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. International Conference on Learning Representations (ICLR), 2014.
[16] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In ICML, volume 32, pages 1278-1286, 2014.
[17] N. Oudjane and C. Musso. Progressive correction for regularized particle filters. International Conference on Information Fusion, 2000.
[38] Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. Advances in Neural Information Processing Systems (NeurIPS), 31, 2018.
[39] Ludger Ruschendorf. Copulas, Sklar's Theorem, and Distributional Transform. Springer Berlin Heidelberg, 2013.
[40] D. Alspach and H. Sorenson. Nonlinear bayesian estimation using Gaussian sum approximations. IEEE Transactions on Automatic Control, 1972.
[41] Mike Klaas, Nando de Freitas, and Arnaud Doucet. Toward practical  $N^2$  Monte Carlo: The marginal particle filter. In Proceedings of the Twenty-First Conference on Uncertainty in Artificial Intelligence, pages 308–315, 2005.
[42] Jinlin Lai, Justin Domke, and Daniel Sheldon. Variational marginal particle filters. International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
[43] Keiron O'Shea and Ryan Nash. An introduction to convolutional neural networks. arXiv preprint arXiv:1511.08458, 2015.
[44] Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. Advances in Neural Information Processing Systems (NeurIPS), 28, 2015.
[45] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 1997.
[46] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. International Conference on Machine Learning (ICML), 2013.
[47] Herbert Jaeger. Tutorial on training recurrent neural networks, covering BPPT, RTRL, EKF and the" echo state network" approach. GMD-Forschungszentrum Informationstechnik Bonn, 2002.
[48] Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Kuttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, et al. Deepmind lab. arXiv preprint arXiv:1612.03801, 2016.
[49] Shuran Song, Fisher Yu, Andy Zeng, Angel X Chang, Manolis Savva, and Thomas Funkhouser. Semantic scene completion from a single depth image. Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
[50] Yi Wu, Yuxin Wu, Georgia Gkioxari, and Yuandong Tian. Building generalizable agents with a realistic and rich 3d environment. arXiv preprint arXiv:1801.02209, 2018.
[51] Marco Frei and Hans R. Kündsch. Mixture ensemble kalman filters. Computational Statistics and Data Analysis, 2013.
[52] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2014.
A Parameterization of the Neural Dynamics and Measurement Models