SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ali Younis and Erik B. Sudderth
{ayounis, sudderth}@uci.edu
Department of Computer Science, University of California, Irvine
Particle filters flexibly represent multiple posterior modes nonparametrically, via a collection of weighted samples, but have classically been applied to tracking problems with known dynamics and observation likelihoods. Such generative models may be inaccurate or unavailable for high-dimensional observations like images. We instead leverage training data to discriminatively learn particle-based representations of uncertainty in latent object states, conditioned on arbitrary observations via deep neural network encoders. While prior discriminative particle filters have used heuristic relaxations of discrete particle resampling, or biased learning by truncating gradients at resampling steps, we achieve unbiased and low-variance gradient estimates by representing posteriors as continuous mixture densities. Our theory and experiments expose dramatic failures of existing reparameterization-based estimators for mixture gradients, an issue we address via an importance-sampling gradient estimator. Unlike standard recurrent neural networks, our mixture density particle filter represents multimodal uncertainty in continuous latent states, improving accuracy and robustness. On a range of challenging tracking and robot localization problems, our approach achieves dramatic improvements in accuracy, while also showing much greater stability across multiple training runs.
A particle filter (PF) [1-4] uses weighted samples to provide a flexible, nonparametric representation of uncertainty in continuous latent states. Classical PFs are generative models that typically require human experts to specify the latent state dynamics and measurement models. With the growing popularity of deep learning and abundance of time-series data, recent work has instead sought to discriminatively learn PF variants that (unlike conventional recurrent neural networks) capture uncertainty in latent states [5-8]. These approaches are especially promising for high-dimensional observations like images, where learning accurate generative models is extremely challenging.
The key challenge in learning discriminative PFs is the non-differentiable particle resampling step, which is key to robustly maintaining diverse particle representations, but inhibits end-to-end learning by preventing gradient propagation. Prior discriminative PFs have typically used heuristic relaxations of discrete particle resampling, or biased learning by truncating gradients. These methods severely compromise the effectiveness of PF training; prior work has often assumed exact dynamical models are known, or used very large numbers of particles for test data, due to these limitations.
After reviewing prior work on generative (Sec. 2) and discriminative (Sec. 3) training of PFs, Sec. 4 demonstrates dramatic failures of popular reparameterization-based gradient estimators for mixture models. This motivates our importance weighted samples gradient estimator, which provides the foundation for the mixture density particle filter of Sec. 5. Experiments in Sec. 6 show substantial advances over prior discriminative PF for several challenging tracking and visual localization tasks.
Figure 1: Sequential state estimation may be formulated via either generative (left) or discriminative (right) graphical models. In either case, dynamics of states  $x_{t}$  are influenced by actions  $a_{t}$ , and estimated via data  $y_{t}$ .
2 Sequential State Estimation via Generative Models
Algorithms for sequential state estimation compute or approximate the posterior distribution of the system state  $x_{t}$ , given a sequence of observations  $y_{t}$  and (optionally) input actions  $a_{t}$ , at discrete times  $t = 1, \ldots, T$ . Sequential state estimation is classically formulated as inference in a generative model like the hidden Markov model (HMM) or state space model of Fig. 1, with dynamics defined by state transition probabilities  $p(x_{t} \mid x_{t-1}, a_{t})$ , and observations generated via likelihoods  $p(y_{t} \mid x_{t})$ .
Given a known Markov prior on latent state sequences, and a corresponding observation sequence, the filtered state posterior  $p(x_{t} \mid y_{t}, y_{t-1}, \ldots, y_{1})$  can in principle be computed via Bayesian inference. When the means of state dynamics and observations likelihoods are linear functions, and noise is Gaussian, exact filtered state posteriors are Gaussian and may be efficiently computed via the Kalman filter (KF) [9-11]. The non-Gaussian state posteriors of general state space models may be approximated via Gaussians [12], but estimating the posterior mean and covariance is in general challenging, and cannot faithfully capture the multimodal posteriors produced by missing or ambiguous data.
2.1 Particle Filters
To flexibly approximate general state posteriors, particle filters (PF) [1-4] represent possible states nonparametrically via a collection of weighted samples or particles. The classical PF parameterizes the state posterior at time  $t$  by a set of  $N$  particles  $x_{t}^{(:)} = \{x_{t}^{(1)},\dots,x_{t}^{(N)}\}$  with associated weights  $w_{t}^{(:)} = \{w_{t}^{[1]},\dots,w_{t}^{(N)}\}$ . PFs recursively update the state posterior given the latest observation  $y_{t}$  and action  $a_{t}$ , and may be flexibly applied to a broad range of models; they only assume that the state dynamics may be simulated, and the observation likelihood may be evaluated.
Particle Proposal. To produce particle locations  $x_{t}^{(:)}$  at time  $t$ , a model of the state transition dynamics is applied individually to each particle  $x_{t-1}^{(i)}$ , conditioned on external actions  $a_{t}$  if available:
$$
x _ {t} ^ {(i)} \sim p \left(x _ {t} \mid x _ {t - 1} = x _ {t - 1} ^ {(i)}, a _ {t}\right). \tag {1}
$$
Measurement Update. After a new particle set  $x_{t}^{(:)}$  has been proposed, particle weights must be updated to account for the latest observation  $y_{t}$  via the known likelihood function:
$$
w _ {t} ^ {(i)} \propto p \left(y _ {t} \mid x _ {t} ^ {(i)}\right) w _ {t - 1} ^ {(i)}. \tag {2}
$$
The updated weights are then normalized so that  $\sum_{i=1}^{N} w_t^{(i)} = 1$ . This weight update is motivated by importance sampling principles, because (asymptotically, as  $N \to \infty$ ) it provides an unbiased approximation of the true state posterior  $p(x_t \mid y_t, \ldots, y_1) \propto p(y_t \mid x_t) p(x_t \mid y_{t-1}, \ldots, y_1)$ .
Particle Resampling. Due to the stochastic nature of PFs, over time particles will slowly diverge to regions with low posterior probability, and will thus be given little weight in the measurement step. If all but a few particles have negligible weight, the diversity and effective representational power of the particle set is diminished, resulting in poor estimates of the state posterior.
To address this, particle resampling is used to maintain diversity of the particle set over time, and may be performed at every iteration or more selectively when the "effective" sample size becomes too small [13]. During resampling, a new uniformly-weighted particle set  $\hat{x}_t^{(\cdot)}$  is constructed by discrete resampling with replacement. Each resampled particle is a copy of some existing particle, where copies are sampled with probability proportional to the particle weights:
$$
\hat {x} _ {t} ^ {(i)} = x _ {t} ^ {(j)}, \quad j \sim \operatorname {C a t} \left(w _ {t} ^ {(1)}, \dots , w _ {t} ^ {(N)}\right). \tag {3}
$$
After resampling, assigning uniform weights  $\hat{w}_t^{(i)} = \frac{1}{N}$  maintains an unbiased approximate state posterior. This discrete resampling is non-differentiable, preventing gradient estimation via standard reparameterization [14-16], and inhibiting end-to-end learning of dynamics and measurement models.
2.2 Regularized Particle Filters
The regularized PF [17, 18] acknowledges that a small set of discrete samples will never exactly align with the true continuous state, and thus estimates a continuous state density  $m(x_{t} \mid x_{t}^{(:)}, w_{t}^{(:)}, \beta)$  by convolving particles with a continuous kernel function  $K$  with bandwidth  $\beta$ :
$$
m \left(x _ {t} \mid x _ {t} ^ {(:)}, w _ {t} ^ {(:)}, \beta\right) = \sum_ {i = 1} ^ {N} w _ {t} ^ {(i)} K \left(x _ {t} - x _ {t} ^ {(i)}; \beta\right). \tag {4}
$$
The kernel function could be a Gaussian, in which case  $\beta$  is a (dimension-specific) standard deviation. The extensive literature on kernel density estimation (KDE) [19] provides theoretical guidance on the choice of kernel. Quadratic Epanechnikov kernels take  $K(u;\beta) = \frac{3}{4}\left(1 - (u / \beta)^2\right)$  if  $|u / \beta |\leq 1$ ,  $K(u;\beta) = 0$  otherwise, and asymptotically minimize the mean-squared-error in the estimation of the underlying continuous density. A number of methods have been proposed for selecting the smoothing bandwidth  $\beta$  [19-21], but they often have asymptotic justifications, and can be unreliable for small  $N$ .
Attractively, regularized PFs resample particles  $\hat{x}_t^{(i)}\sim m(x_t\mid x_t^{(:,)},w_t^{(:,)},\beta)$  from the continuous mixture density rather than the discrete particle set. By resampling from a continuous distribution, regularized PFs ensure that no duplicates exist in the resampled particle set, increasing diversity while still concentrating particles in regions of the state space with high posterior probability. Our proposed mixture density  $PF$  generalizes regularized PFs, using bandwidths  $\beta$  that are tuned jointly with discriminative models of the state dynamics and observation likelihoods.
3 Conditional State Estimation via Discriminative Particle Filters
While classical inference algorithms like KFs and PFs are effective in many domains, they require faithful generative models. A human expert must typically design most aspects of the state dynamics and observation likelihoods, and algorithms that use PFs to learn generative models often struggle to scale beyond low-dimensional, parametric models [22]. For complex systems with high-dimensional observations like images, learning generative models of the observation likelihoods is often impractical, and arguably more challenging than estimating latent states given an observed data sequence. We instead learn discriminative models of the distribution of the state conditioned on observations (see Fig. 1), and replace manually engineered generative models with deep neural networks learned from training sequences with (potentially sparse) latent state observations.
Discriminative Kalman Filters. Discriminative modeling has been used to learn conditional variants of the KF [23, 24] where the posterior is parameterized by a Gaussian state space model, and state transition and observation emission models are produced from data by trained neural networks. Like conditional random field (CRF) [25] models for discrete data, discriminative KFs do not learn likelihoods of observations, but instead condition on them. Discriminative KFs desirably learn a posterior covariance, and thus do not simply output a state prediction like conventional recurrent neural networks. But, they are limited to unimodal, Gaussian approximations of state uncertainty.
Discriminative Particle Filters. It is attractive to integrate similar discriminative learning principles with PFs, but the technical challenges are substantially greater due to the nonparametric particle representation of posterior uncertainty, and the need for particle resampling to maintain diversity.
Like in generative PFs, discriminative PFs update the particle locations using a dynamics model as in Eq. (1). Unlike classic PFs, the particle weights are not updated using a generative likelihood function. Instead, discriminative PFs compute new weights  $w_{t}^{(:)}$  using a measurement model function  $\ell(x_{t}; y_{t})$  trained to properly account for uncertainty in the discriminative particle posterior:
$$
w _ {t} ^ {(i)} \propto \ell \left(x _ {t} ^ {(i)}; y _ {t}\right) w _ {t - 1} ^ {(i)}. \tag {5}
$$
Here  $\ell(x_{t}; y_{t})$  is a (differentiable) function optimized to improve the accuracy of the discriminative particle posterior, rather than a generative likelihood.
Creating a learnable discriminative PF requires parameterizing the dynamics and measurement models as differentiable functions, such as deep neural networks. This is straightforward for the measurement model  $\ell(x_{t}^{(i)}; y_{t})$ , which may be defined via any feed-forward neural network architecture like those typically used for classification (see Fig. 2). For the dynamics model, the neural network does not simply need to score particles; it must be used for stochastic simulation. Using reparameterization [14-16], dynamics simulation is decomposed as sampling from a standard Gaussian distribution, and then
Figure 2: Left: Our MDPF method showing the various sub-components. Middle: Our A-MDPF method with decoupled measurement models and bandwidths. Right: Dynamics and measurement model structures used in MDPF and A-MDPF. The dynamics and measurement models are composed of several neural networks as well as some fixed transforms, which convert the angular state dimensions of particles into a vector representation.
Nanotubular Block (General Item)
transforming that sample using a learned neural network, possibly conditioned on action  $a_{t}$ :
$$
x _ {t} ^ {(i)} = f \left(\eta_ {t} ^ {(i)}; x _ {t - 1} ^ {(i)}, a _ {t}\right), \quad \eta_ {t} ^ {(i)} \sim N (0, I). \tag {6}
$$
The dynamics model  $f(\eta ;x_{t - 1},a_t)$  is a feed-forward neural network that deterministically processes noise  $\eta$ , conditioned on  $x_{t - 1}$  and  $a_{t}$ , to implement the dynamical sampling of Eq. (1). The neural network  $f$  may be flexibly parameterized because discriminative particle filters only require simulation of the dynamical model, not explicit evaluation of the implied conditional state density.
3.1 Prior Work on Discriminative Particle Filters
Specialized, heuristic variants of discriminative PFs have been previously used for tasks like robot localization [26] and multi-object tracking [27]. Principled end-to-end learning of a discriminative PF requires propagating gradients through the PF algorithm, including the discrete resampling step. Zhu et al. [28] explore replacing the PF resampling step with a learned (deterministic) particle transform, but find that exploding gradient magnitudes prohibit end-to-end training.
Truncated-Gradient Particle Filter (TG-PF). The first so-called "differentiable" particle filter [5] actually treated the particle resampling step as a non-differentiable discrete resampling operation, and simply truncated all gradients to zero when backpropagating through this resampling. Though simple, this approach leads to biased gradients, and often produces ineffective models because backpropagation through time (BPTT [29]) is not possible. Perhaps due to these limitations, experiments in Jonschkowski et al. [5] assumed a simplified training scenario where the ground-truth dynamics are known, and only the measurement model must be learned.
Discrete Importance Sampling Particle Filter (DIS-PF). Scibior et al. [6] propose a gradient estimator for generative PFs that may also be adapted to discriminative PFs. They develop their estimator by invoking prior work on score-function gradient estimators [30], but we provide a simpler derivation via importance sampling principles. Instead of discretely resampling particles as in Eq. (3), a separate set of weights  $v_{t}^{(:)}$  is defined and used to generate samples:
$$
\hat {x} _ {t} ^ {(i)} = x _ {t} ^ {(j)}, \quad j \sim \operatorname {C a t} \left(v _ {t} ^ {(1)}, \dots , v _ {t} ^ {(N)}\right). \tag {7}
$$
To account for discrepancies between these resampling weights  $v^{(\cdot)}$  and true weights  $w^{(\cdot)}$ , the resampled particle weights  $\hat{w}_t^{(i)}$  are defined via importance sampling. The choice of  $v_t^{(\cdot)}$  critically impacts the effectiveness of the resampling step. To maintain the standard PF update of Eq. (3), we set  $v_t^{(i)} = w_t^{(i)}$ , yielding the following resampled particle weights and associated gradients:
$$
\hat {w} _ {t} ^ {(i)} = \frac {w _ {t} ^ {(i)}}{v _ {t} ^ {(i)} | _ {v _ {t} ^ {(i)} = w _ {t} ^ {(i)}}} = 1, \quad \nabla_ {\phi} \hat {w} _ {t} ^ {(i)} = \frac {\nabla_ {\phi} w _ {t} ^ {(i)}}{v _ {t} ^ {(i)} | _ {v _ {t} ^ {(i)} = w _ {t} ^ {(i)}}}. \tag {8}
$$
Intuitively, particle locations are resampled according to the current weights in the forward pass of DIS-PF. Then when computing gradients, perturbations of the observation and dynamics models are accounted for by changes in the associated particle weights.
A drawback of this discrete importance sampling is known as the ancestor problem: since the resampled particle set is constructed from a subset of the pre-resampled particles, no direct gradients exist for particles that were not chosen during resampling; they are only indirectly influenced by the weight normalization step. This increases the DIS gradient estimator variance, slowing training.
Soft Resampling Particle Filter (SR-PF). Karkus et al. [7] propose a soft-resampling (SR) mixing of the particle weights with a discrete uniform distribution before resampling:  $v_{t}^{(i)} = (1 - \lambda)w_{t}^{(i)} + \lambda /N$ .
Viewing particle locations as fixed, SR-PF evaluates the resampled particle weights and gradients as
$$
\hat {w} _ {t} ^ {(i)} = \frac {w _ {t} ^ {(i)}}{(1 - \lambda) w _ {t} ^ {(i)} + \lambda / N}, \qquad \nabla_ {\phi} \hat {w} _ {t} ^ {(i)} = \nabla_ {\phi} \left(\frac {w _ {t} ^ {(i)}}{(1 - \lambda) w _ {t} ^ {(i)} + \lambda / N}\right). \tag {9}
$$
While the SR-PF weight update is differentiable, it does not avoid the ancestor problem. By resampling low-weight particles more frequently, SR-PF will degrade overall PF performance when  $N$  is not large. More subtly, the gradient of Eq. (9) has (potentially substantial) bias: it assumes that perturbations of model parameters influence particle weights, but not the outcome of discrete resampling (7). Indeed in some experiments in [7], smoothing is actually disabled by setting  $\lambda = 0$ .
Concrete Particle Filter (C-PF). The Gumbel-softmax or Concrete distribution [31, 32] approximates discrete sampling by interpolation with a continuous distribution, enabling reparameterized gradient estimation. Each resampled particle is a convex combination of the weighted particle set:
$$
\hat {x} _ {t} ^ {(i)} = \sum_ {j = 1} ^ {N} \alpha_ {i j} x _ {t} ^ {(j)}, \quad \alpha_ {i j} = \frac {\exp \left(\left(\log \left(w _ {t} ^ {(j)}\right) + G _ {i j}\right) / \lambda\right)}{\sum_ {k = 1} ^ {N} \exp \left(\left(\log \left(w _ {t} ^ {(k)}\right) + G _ {i k}\right) / \lambda\right)}, \quad G _ {i j} \sim \text {G u m b e l}. \tag {10}
$$
Gradients of (10) are biased with respect to discrete resampling due to the non-learned "temperature" hyperparameter  $\lambda > 0$ . When  $\lambda$  is large, the highly-biased relaxation will interpolate between modes, and produce many low-probability particles. Bias decreases as  $\lambda \to 0$ , but in this limit the variance is huge. Even with careful tuning of  $\lambda$ , Concrete relaxations are most effective for small  $N$ . Maddison et al. [32] focus on models with binary latent variables, and find that performance degrades even for  $N = 8$ , far smaller than the  $N$  needed for practical PFs. While we provide C-PF as a baseline, we are unaware of prior work successfully incorporating Concrete relaxations in PFs.
Optimal Transport Particle Filter (OT-PF). OT-PF [8] modifies PFs by replacing stochastic resampling with the optimization-based solution of an entropy-regularized optimal transport (OT) problem. OT seeks a probabilistic correspondence between particles  $x_{t}^{(:)}$  with weights  $w_{t}^{(:)}$  as in (5), and a corresponding set of uniformly weighted particles, by minimizing a Wasserstein metric  $\mathcal{W}_{2,\lambda}^2$ :
$$
\min  _ {\tilde {\alpha} \in [ 0, 1 ] ^ {N \times N}} \sum_ {i, j = 1} ^ {N} \tilde {\alpha} _ {i j} \left(\left| \left| x _ {t} ^ {(i)} - x _ {t} ^ {(j)} \right| \right| ^ {2} + \lambda \log \frac {\tilde {\alpha} _ {i j}}{N ^ {- 1} w _ {t} ^ {(j)}}\right), \text {s . t .} \sum_ {j = 1} ^ {N} \tilde {\alpha} _ {i j} = \frac {1}{N}, \sum_ {i = 1} ^ {N} \tilde {\alpha} _ {i j} = w _ {t} ^ {(j)}. \tag {11}
$$
Entropy regularization is required for differentiability, and OT-PF accuracy is sensitive to the nonlearned hyperparameter  $\lambda > 0$ . OT-PF uses this entropy-regularized mapping to interpolate between particles:  $\hat{x}_t^{(i)} = \sum_{j=1}^N N \tilde{\alpha}_{ij} x_t^{(j)}$ . This assignment approximates the results of true stochastic resampling in the limit as  $N \to \infty$ , but lacks accuracy guarantees for moderate  $N$ .
OT-PF solves a regularized OT problem which relaxes discrete resampling, producing biased gradients that are reminiscent (but distinct) from C-PF. Minimization of Eq. (11) via the Sinkhorn algorithm [33] requires  $\mathcal{O}(N^2)$  operations. This OT problem must be solved at each step of both training and test, and in practice, OT-PF is substantially slower than all competing discriminative PFs (see Fig. 10).
4 Stable Gradient Estimation for Mixture Model Resampling
Training our discriminative mixture density particle filter (MDPF) requires unbiased and low-variance estimates of gradients of samples from continuous mixtures. Mixture sampling is classically decomposed into discrete (selecting a mixture component) and continuous (sampling from that component) steps. The Gumbel-softmax or Concrete distribution [31, 32] provides a reparameterizable relaxation of discrete resampling, but may have substantial bias. Some work has also applied OT methods for gradient estimation in restricted families of Gaussian mixtures [35, 36]. We instead develop importance-sampling estimators that are unbiased, computationally efficient, and low variance.
Implicit Reparameterization Gradients. Direct reparameterization of samples, as used for the latent Gaussian distributions in variational autoencoders [14-16], cannot be easily applied to mixture models. Reparameterized sampling requires an invertible standardization function  $S_{\phi}(z) = \epsilon$  that transforms a sample  $z$ , drawn from a distribution parameterized by  $\phi$ , to an auxiliary variable  $\epsilon$  independent of  $\phi$ . For mixture distributions, the inverse of  $S_{\phi}(z)$  cannot be easily computed.
The implicit reparameterization gradients (IRG) estimator [37, 38] avoids explicit inversion of the standardization function  $S_{\phi}(z)$ . Using implicit differentiation, reparameterization can be achieved so long as a computable (analytically or numerically) and invertible  $S_{\phi}(z)$  exists, without the need to
Figure 3: Left: Example changes in samples from a 1D mixture model with two Epanechnikov [34] components. Under IRG, changes in the mixture are shown by dramatic sample shifting, where some samples must change modes to account for the changes in the relative weights of each mixture mode. Explicitly plotting the particle transformation induced by IRG reveals a discontinuity (bottom). In contrast, IWSG smoothly reweights samples. Right: Example changes in samples from a 2D mixture of two Gaussians. IRG again induces large shifts as particles change modes, as demonstrated by the vector field (bottom).
explicitly compute its inverse. IRG gradients are expressed via the Jacobian of  $S_{\phi}(z)$ :
$$
\nabla_ {\phi} z = - \left(\nabla_ {z} S _ {\phi} (z)\right) ^ {- 1} \nabla_ {\phi} S _ {\phi} (z). \tag {12}
$$
For univariate distributions, the cumulative distribution function (CDF)  $M_{\phi}(z)$  is a valid standardization function. For higher dimensions, the multivariate distributional transform [39] is used:
$$
S _ {\phi} (z) = \left(M _ {\phi} \left(z _ {1}\right), M _ {\phi} \left(z _ {2} \mid z _ {1}\right), \dots , M _ {\phi} \left(z _ {D} \mid z _ {1}, \dots , z _ {D - 1}\right)\right). \tag {13}
$$
While IRG may be applied to any distribution with continuous CDF, and has been explicitly suggested (without experimental validation) for mixtures [37], we show that it may have enormous variance.