summaries: The authors make three contributions. 1) They apply the constrained optimization framework for VAE type models to the deep state space model case. 2) They add in the hierarchical prior approach to model the initial distribution for the first state. 3) They propose a model with linear transition matrices given by a neural network weighting learned basis matrices. Auxiliary observation variables are also introduced to allow for a linear observation model along with an extra encoder/decoder. This means filtering and smoothing distributions can be computed analytically.

summaries: This paper has two main contributions.  1. The authors adapt REWO for helping with disentanglement of learned representations and improving predictive accuracy 2. The authors replace the RNN encoding of the nonlinearity in [6] by the locally linear model of [11-12] while keeping closed form inference  The authors then illustrate how combining these together leads to better accuracy in learned latent dynamics for a set of models, in particular they show that their approach can recover the unobserved velocity as a latent variable in motion models.

summaries: This paper proposes a framework for learning neural network-based state-space models from raw data. In comparison to the standard RNN-based approach [1], there are several modifications: (i) a modified ELBO formulated in the context of a constrained optimization (CO) framework [2]. CO prioritizes reconstruction term at the beginning of the training; (ii) hierarchical prior [3] instead of a simple Gaussian; (iii) a modified Kalman-VAE framework [4] instead of RNN-based transition model. Experimental evidence on pendulum data and the reacher environment indicates that the resulting framework outperforms baselines in predictive modeling and ELBO values.   [1] A Recurrent Latent Variable Model for Sequential Data  [2] Taming VAEs  [3] Learning Hierarchical Priors in VAEs  [4]  A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning

summaries: This paper proposes various improvements to deep sequential latent variable models (or deep state space models). Namely, 1) the constrained optimization method from the Taming VAEs paper [Rezende & Viola, 2018], 2) a method for improving empirical priors over the first time step’s latent variable (referred to as VHP), and 3) a new model that combines the locally linear mappings of Karl et al., 2017 and the KVAE variable separation from Fraccaro et al., 2017. These improvements are demonstrated on pendulum and reacher datasets (both with images and joint angles), showing improvements in system identification, prediction accuracy, and the ability to use the state representations from MBRL. Overall, the paper provides several helpful tips for constructing and training deep state space models.

abstractText: Deep state-space models (DSSMs) enable temporal predictions by learning the underlying dynamics of observed sequence data. They are often trained by maximising the evidence lower bound. However, as we show, this does not ensure the model actually learns the underlying dynamics. We therefore propose a constrained optimisation framework as a general approach for training DSSMs. Building upon this, we introduce the extended Kalman VAE (EKVAE), which combines amortised variational inference with classic Bayesian filtering/smoothing to model dynamics more accurately than RNN-based DSSMs. Our results show that the constrained optimisation framework significantly improves system identification and prediction accuracy on the example of established state-of-the-art DSSMs. The EKVAE outperforms previous models w.r.t. prediction accuracy, achieves remarkable results in identifying dynamical systems, and can furthermore successfully learn state-space representations where static and dynamic features are disentangled.

authors:
Alexander Alemi
Ben Poole
Ian Fischer
Joshua Dillon
Rif A. Saurous
Kevin Murphy

title: Fixing a broken ELBO

venue: In International Conference on Machine Learning,

authors:
Justin Bayer
Maximilian Soelch
Atanas Mirchev
Baris Kayalibay
Patrick van der Smagt

title: Mind the Gap when Conditioning Amortised Inference in Sequential Latent-Variable Models

venue: In International Conference on Learning Representations. OpenReview.net,

authors:
Philip Becker-Ehmck
Maximilian Karl
Jan Peters
Patrick van der Smagt

title: Learning to Fly via Deep Model-Based Reinforcement Learning

venue: arXiv preprint arXiv::

authors:
Samuel R. Bowman
Luke Vilnis
Oriol Vinyals
Andrew Dai
Rafal Jozefowicz
Samy Bengio

title: Generating Sentences from a Continuous Space

venue: In Conference on Computational Natural Language Learning,

authors:
Junyoung Chung
Kyle Kastner
Laurent Dinh
Kratarth Goel
Aaron C Courville
Yoshua Bengio

title: A recurrent latent variable model for sequential data

venue: In Advances in Neural Information Processing Systems,

authors:
Andreas Doerr
Christian Daniel
Martin Schiegg
Nguyen-Tuong Duy
Stefan Schaal
Marc Toussaint
Trimpe Sebastian

title: Probabilistic Recurrent State-Space Models

authors:
Marco Fraccaro
Søren Kaae Sønderby
Ulrich Paquet
Ole Winther

title: Sequential Neural Models with Stochastic Layers

authors:
Marco Fraccaro
Simon Kamronn
Ulrich Paquet
Ole Winther

authors:
Danijar Hafner
Timothy Lillicrap
Ian Fischer
Ruben Villegas
David Ha
Honglak Lee
James Davidson

title: Learning Latent Dynamics for Planning from Pixels

authors:
Irina Higgins
Loïc Matthey
Arka Pal
Christopher Burgess
Xavier Glorot
Matthew Botvinick
Shakir Mohamed
Alexander Lerchner

title: beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework

title: Stochastic Processes and Filtering Theory

venue: Courier Corporation,

title: A New Approach to Linear Filtering and Prediction Problems

venue: Journal of Basic Engineering,

authors:
Maximilian Karl
Maximilian Soelch
Justin Bayer
Patrick van der Smagt

title: Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data

authors:
Maximilian Karl
Maximilian Soelch
Philip Becker-Ehmck
Djalel Benbouzid
Patrick van der Smagt
Justin Bayer

title: Unsupervised Real-Time Control through Variational Empowerment

venue: arXiv preprint arXiv:1710.05101,

authors:
Diederik P. Kingma
Max Welling

title: Auto-Encoding Variational Bayes

authors:
Alexej Klushyn
Nutan Chen
Richard Kurle
Botond Cseke
Patrick van der Smagt

title: Learning Hierarchical Priors in VAEs

authors:
Rahul G. Krishnan
Uri Shalit
David Sontag

title: Deep Kalman Filters

venue: arXiv preprint arXiv:1511.05121,

authors:
Richard Kurle
Syama Sundar Rangapuram
Emmanuel de Bezenac
Stephan Günnemann
Jan Gasthaus

title: Deep Rao-Blackwellised Particle Filters for Time Series Forecasting

authors:
Sergey Levine
Vladlen Koltun

title: Guided Policy Search

authors:
Radford M. Neal
Geoffrey E. Hinton

title: A View of the EM Algorithm that Justifies Incremental, Sparse, and other Variants

venue: In Learning in graphical models,

authors:
Syama Sundar Rangapuram
Matthias Seeger
Jan Gasthaus
Lorenzo Stella
Yuyang Wang
Tim Januschowski

title: Deep State Space Models for Time Series Forecasting

authors:
Herbert E. Rauch
Fang-Wu Tung
Charlotte T. Striebel

title: Maximum Likelihood Estimates of Linear Dynamic Systems

venue: AIAA Journal,

authors:
Danilo Jimenez Rezende
Shakir Mohamed
Daan Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

authors:
David Salinas
Valentin Flunkert
Jan Gasthaus
Tim Januschowski

title: Deepar: Probabilistic forecasting with autoregressive recurrent networks

venue: International Journal of Forecasting,

authors:
Jakub Tomczak
Max Welling

title: VAE with a VampPrior

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Manuel Watter
Jost Springenberg
Joschka Boedecker
Martin Riedmiller

title: Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images

authors:
Li Yingzhen
Stephan Mandt

title: Disentangled Sequential Autoencoder

heading: 1 Introduction

text:
Many dynamical systems can only be (partially) observed, with the exact dynamics unknown. Yet, precise models are needed for prediction and control [e.g. 19, 21, 25, 3]. Learning such accurate models is subject of current research, especially in image-based domains [e.g. 9].
Deep state-space models (DSSMs) [e.g. 28, 17, 13] describe sequence data by a (typically Markovian) nonlinear transition model and a nonlinear observation model. The transition model is assumed to capture the dynamics underlying the observed data, and the observation model maps the latent variables to the domain of observable data and accounts for the measurement noise. In this paper, we show that DSSMs, however, often do not learn the correct system dynamics, which is suboptimal for accurate predictions or performing downstream tasks, such as model-based reinforcement learning.
We identify in our experiments three main reasons causing this problem: (i) DSSMs are often trained by maximising the sequential evidence lower bound (ELBO). High ELBO values, however, do not imply the model has learned the correct system dynamics. (ii) The prior/initial distribution is usually just a Gaussian. This often leads to an over-regularisation of the approximate posterior or even to a broken generative model, where the transition model is not optimised to process samples from the prior. (iii) Most DSSMs use RNNs to approximate or support Bayesian filtering/smoothing. Yet, RNNs often prove to be a limiting factor for learning accurate models of the system dynamics; moreover, RNN-based transition models as in [8, 9] can lead to a non-Markovian state-space, where the latent variables do not capture the entire information about the system’s state.
To address these problems, we propose the following solutions: (i) we introduce a constrained optimisation (CO) framework as a general approach for learning DSSMs. It ensures a good reconstruction ∗Correspondence to a.klushyn@tum.de. †Work done while at Machine Learning Research Lab, Volkswagen Group, Munich.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
quality and thus provides a necessary basis for learning the underlying system dynamics. To this end, we extend a recent method [16] presented in the context of variational autoencoders (VAEs) to DSSMs. We do this by formulating the sequential ELBO as the Lagrangian of a CO problem and introducing the associated optimisation algorithm. (ii) We complement the proposed CO framework with a powerful empirical Bayes prior. (iii) To obtain more accurate predictions of observed dynamical systems, we introduce the extended Kalman VAE (EKVAE), where we dispense with RNNs by combining extended Kalman filtering/smoothing with amortised variational inference and a neural linearisation approach. Furthermore, we show that the EKVAE is capable of learning state-space representations where static and dynamic features are disentangled. We use this to validate the learned model in the context of model-based reinforcement learning.
Our evaluation includes experiments on the image data of a moving pendulum [13] and on the reacher environment [26], where we use angle as well as high-dimensional RGB image data as observations. We show that each of our proposed approaches significantly helps in learning accurate models of observed dynamical systems—and that applying our CO framework to established DSSMs leads to a substantial increase in their prediction accuracy.

heading: 2 Background: A Rate–Distortion Perspective on Deep State-Space Models

text:
DSSMs [e.g. 28, 17, 13] model an unknown distribution of observed sequence data x1:T = (x1,x2, . . . ,xT ) by means of typically lower-dimensional latent variables z1:T that represent the underlying state of the system. To achieve this, the Markov assumption is imposed. It states that the future state zt+1 as well as the current observation xt solely depend on zt:
p(x1:T , z1:T |u1:T ) = pθ(x1:T |z1:T ) pψ(z1:T |u1:T )
= p(z1) pθ(x1| z1) T∏ t=2 pψ(zt| zt−1,ut−1) pθ(xt| zt), (1)
where u1:T are optional control signals (actions), and the use of different parameters (θ, ψ) will become important in the course of this paper. The model parameters in Eq. (1) are often learned through amortised variational inference [e.g. 17, 13]. This requires introducing a recognition model qφ(z1:T |x1:T ,u1:T ) that learns—in combination with the transition model pψ(zt| zt−1,ut−1)—the dynamics underlying the observed data. The resulting objective function is known as sequential ELBO:
log p(x1:T |u1:T ) ≥ FELBO(θ, ψ, φ) = Eqφ [ log
pθ(x1:T | z1:T ) pψ(z1:T |u1:T ) qφ(z1:T |x1:T ,u1:T )
] . (2)
In the context of generative models, the ELBO can be divided into a reconstruction term (distortion) and a compression term (rate) [1]. We extend the theory in [1] to DSSMs, where the distortion D(θ, φ) = −Eqφ [ log pθ(x1:T | z1:T ) ] optimises the model’s ability for reconstructing observations,
whereas the rateR(φ, ψ) = KL ( qφ(z1:T |x1:T ,u1:T )‖ pψ(z1:T |u1:T ) ) enables learning the underlying dynamics. These definitions lead to the following general formulation of the ELBO:
FELBO(θ, ψ, φ) = −D(θ, φ)−R(ψ, φ). (3) Balancing the ratio between distortion and rate during optimisation can be an effective approach to improve the learning of DSSMs, as we discuss in the following section.

heading: 3 Constrained Optimisation Framework for Improved System Identification

text:
High ELBO values do not necessarily imply that the model has learned the underlying system dynamics of the observed data, as we verify in Sec. 6.2. This is because different combinations of rate and distortion can result in the same ELBO value. Previous work addresses this issue by introducing weighting schedules for either D(θ, φ) orR(ψ, φ) [e.g. 4] since a different ratio favours either better reconstruction or compression [1]. However, we demonstrate in Sec. 6 that balancing reconstruction and compression with predefined annealing schedules often does not achieve the desired result.
A more recent approach originates from the framework of VAEs: Rezende and Viola [23] and Klushyn et al. [16] define the VAE as a CO problem allowing for controlling the model’s reconstruction quality.
We transfer this approach to DSSMs to ensure a good reconstruction—i.e. a low D(θ, φ)—and thus provide a sufficient basis for learning the underlying system dynamics. To this end, we formulate the sequential ELBO as the Lagrangian of a CO problem (i) by specifying the rate R(ψ, φ) in Eq. (3) as optimisation objective; and (ii) by imposing the inequality constraint D(θ, φ) ≤ D0. Here, D0 is a hyperparameter that defines the baseline for our desired reconstruction quality—we provide a heuristic for the simple determination of D0 in App. A.1. The resulting Lagrangian is
L(θ, ψ, φ;λ) = R(ψ, φ) + λ ( D(θ, φ)−D0 ) , (4)
where the Lagrange multiplier λ can be viewed as a weighting term for the distortion.
The original EM algorithm [e.g. 20] for optimising the ELBO, minθ,ψ minφ −FELBO(θ, ψ, φ), provides the following connection to the CO problem:
M-step︷︸︸︷ min θ,ψ E-step︷ ︸︸ ︷ max λ min φ L(θ, ψ, φ;λ) s.t. λ ≥ 0, (5)
where, unlike in the original EM algorithm, we want qφ(z1:T |x1:T ,u1:T ) to additionally satisfy the inequality constraint D(θ, φ) ≤ D0 in the E-step. However, as detailed in [16], it can only be guaranteed that L(θ, ψ, φ;λ) optimises a lower bound on log p(x1:T |u1:T ) if and only if 1 ≥ λ ≥ 0.

heading: 3.1 Learning the Initial Distribution

text:
It is common practice to define the initial/prior distribution p(z1) as a standard normal [e.g. 17, 8]. However, in this case, the prior KL in the ELBO can cause an over-regularisation (cf. [16]) of the approximate posterior and thus of the transition model. Furthermore, if the discrepancy between prior and posterior is too large, we may obtain a broken generative model, where p(zt| zt−1,ut−1) is not trained to process samples from p(z1). We provide empirical evidence in Sec. 6.1.
This issue often arises in the context of neural models trained by stochastic gradient methods: in order to alleviate possible vanishing-gradient problems, time-series data is typically cut into equally-sized short-length units. For sufficiently large datasets, the initial states can therefore be assumed to cover most possible states, which results in a nontrivial marginal approximate posterior EpD [ qφ(z1|x,u) ] .
Since the optimal prior distribution is p∗(z1) = EpD [ qφ(z1|x,u) ] [cf. 27], an empirical Bayes prior pψ0(z1) must have the complexity to approximate pψ0(z1) ≈ p∗(z1). For this reason, we propose to learn a hierarchical prior pψ0(z1) = ∫ pψ0(z1| ζ) p(ζ) dζ as part of the DSSM by applying the variational approach in [16]. Klushyn et al. [16] define, by means of an approximate distribution qφ0(ζ|z1), a (VAE-like) lower bound on the optimal empirical Bayes prior:
Ep∗(z1) [ log pψ0(z1) ] ≥ EpD(x1:T ,u1:T ) Eqφ(z1:T |x1:T ,u1:T ) [ Eqφ0 (ζ| z1) [ log
pψ0(z1| ζ) p(ζ) qφ0(ζ| z1) ] ︸ ︷︷ ︸
=FVHP(ψ0,φ0; z1)
] , (6)
where p(ζ) is a standard normal distribution and pD(x1:T ,u1:T ) is the empirical distribution of our data. This method is referred to as variational hierarchical prior (VHP). Learning the VHP as part of the model is consistent with the CO problem since Eq. (6) introduces an upper bound on the rate, R(ψ, φ, ψ0) ≤ R(ψ, φ, ψ0, φ0) (see App. A.2), and thus a lower bound on the ELBO. As a result, we obtain the Lagrangian L(θ, ψ, φ, ψ0, φ0;λ) = R(ψ, φ, ψ0, φ0) + λ ( D(θ, φ)−D0 ) and arrive at the following optimisation problem:
Empirical Bayes︷︸︸︷ min ψ0,φ0 M-step︷︸︸︷ min θ,ψ E-step︷ ︸︸ ︷ max λ min φ L(θ, ψ, φ, ψ0, φ0;λ) s.t. λ ≥ 0. (7)

heading: 3.2 Optimisation Algorithm

text:
In order to find the saddle point of the Lagrangian in Eq. (7), we propose Alg. 1, an extension of REWO [16] to DSSMs. Alg. 1 ensures that we optimise a lower bound on log p(x1:T |u1:T ) at the end of training, which is the case for 1 ≥ λ ≥ 0. For this purpose, we apply a special update scheme
for λ, introduced and explained in depth in [16]: λ(i) = λ(i−1) · exp [ −ν · fλ ( λ(i−1), D(i)(θ, φ)−D0; τ1, τ2 ) · ( D(i)(θ, φ)−D0 )] . (8)
In this context, i denotes the iteration step of the optimisation process. The function fλ is defined as fλ(λ, δ; τ1, τ2) = ( 1−H(δ) ) · tanh (τ1 · (1/λ− 1))− τ2 ·H (δ), where H is the Heaviside function, and τ1 as well as τ2 are slope parameters.
Furthermore, Alg. 1 allows us to efficiently learn the parameters of the VHP (ψ0, φ0) and the transition model (ψ) by dividing the CO process into two phases: an initial and a main phase, which is the reason why we use different parameters for the transition and observation model. In the initial phase, the model is optimised w.r.t. (θ, φ) to reduce the reconstruction error by learning the features of the individual observations xt. The main phase starts as soon as the inequality constraintD(θ, φ) ≤ D0 is satisfied. This serves as starting point for additionally optimising the parameters of the VHP (ψ0, φ0) and the transition model (ψ), i.e. for learning the system dynamics.
Algorithm 1 REWO for deep state-space models Initialise i = 1 Initialise λ(0) = 1 Initialise INITIALPHASE = TRUE while training do
Compute D̂ba (batch average) D̂(i) = (1− α) · D̂ba + α · D̂(i−1), (D̂(0) = D̂ba) λ(i) ← λ(i−1) · exp [ −ν · fλ ( λ(i−1), D̂(i) −D0; τ1, τ2 ) · ( D̂(i) −D0 )] if D̂(i) ≤ D0 then
INITIALPHASE = FALSE end if if INITIALPHASE then
Optimise L(θ, ψ, φ, ψ0, φ0;λ(i)) w.r.t. θ, φ else
Optimise L(θ, ψ, φ, ψ0, φ0;λ(i)) w.r.t. θ, ψ, φ, ψ0, φ0 end if i← i+ 1
end while

heading: 4 Extended Kalman VAE

text:
The CO framework can be applied to any DSSM whose objective function is covered by the general rate–distortion formulation of the ELBO (Eq. (3)). We provide derivations for popular baseline models [17, 13] in App. A.3 and A.4. However, to achieve high prediction accuracies, the model itself should not prove to be a limiting factor for learning a precise description of the system dynamics.
Most DSSMs use deterministic RNNs as part of the recognition and/or transition model [e.g. 17, 13, 8, 9]. In [17], for instance, the parameters of the approximate posterior are learned through a (bidirectional) RNN, which is expected to replace classic Bayesian filtering/smoothing. The RNNbased transition model in [e.g. 8], on the other hand, allows combining amortised variational inference with Kalman filtering/smoothing [12, 22]. However, the use of RNNs often leads to less accurate models of the dynamical system, as we show in Sec. 6.2 and 6.3.
In order to increase the prediction accuracy, we introduce the extended Kalman VAE (EKVAE), where we dispense with RNNs by combining amortised variational inference with Bayesian filtering/smoothing. To compute the posterior, we leverage the concept of extended Kalman filters/smoothers [e.g. 11] but avoid the computationally expensive linearisation (Taylor expansion) of the transition and observation model. We achieve this (i) by directly learning the Jacobian of the dynamic model as a function of the current state, which we refer to as neural linearisation (Sec. 4.1); and (ii) by introducing a linear auxiliary-variable model similar to [8, 18] (Sec. 4.2). The EKVAE can be used as filter or smoother. In the following, we focus on the more complex smoother version—as it allows learning a more precise model [2]—and refer to App. A.6 for the filter version.

heading: 4.1 Neural Linearisation of the Dynamic Model

text:
We model the nonlinear dynamical system by a Gaussian transition model that is locally linear w.r.t. discrete time steps [28, 13]:
pψ(zt+1| zt,ut) = N ( zt+1|Fψ(zt,ut) zt +Bψ(zt,ut)ut, Qψ(zt,ut) ) , (9)
where Fψ , Bψ , and Qψ are modelled by linear combinations of M weighted base matrices:
Fψ(zt,ut) = M∑ m=1 α (m) ψ (zt,ut)F (m) etc., where αψ(zt,ut) = softmax ( gψ(zt,ut) ) ∈ RM.
(10) The base matrices { F(m),B(m),Q(m) }M m=1
are learned parameters, and gψ(z,u) is implemented as a neural network.
Next, we make the connection to extended Kalman filtering/smoothing, where the prediction step is based on the local Jacobian (first-order Taylor expansion) of the nonlinear transition function zt+1 = f(zt,ut) + qt, which is unknown in our case. Instead of computing the local Jacobian, however, our transition model (Eq. (9)) is designed to globally find the best linearisation at each time step as a function of the current state and action, which we refer to as neural linearisation.
This approach allows us to apply the extended Kalman filter or smoother algorithm, but replace the computationally expensive Taylor expansion in the prediction step with Fψ(zt,ut) and Bψ(zt,ut), as we derive in App. A.5 and verify in our experiments.

heading: 4.2 Linear Auxiliary-Variable Model

text:
The observation model often needs to learn highly nonlinear functions, especially in case of highdimensional sensory data, such as images. In order to enable an analytic computation of the posterior but avoid an expensive linearisation of the observation model, we introduce auxiliary variables a1:T with a linear dependence on z1:T . As in [8, 18], we learn the nonlinear mapping from a1:T to the high-dimensional observations x1:T by a VAE’s encoder–decoder pair, qφ(at|xt) and pθ(xt|at). Since the dynamics are modelled by the transition function in z1:T—and at can be viewed as a low-dimensional representation of xt—we obtain the following observation model:
p(x1:T ,a1:T | z1:T ) = T∏ t=1 pθ(xt|at) pψ(at| zt). (11)
In contrast to [8, 18], we propose a time-invariant auxiliary-variable model,
pψ(at| zt) = N (at|Hzt, R), (12)
where H and R are globally learned or predefined. The time-invariant H additionally allows us to learn disentangled state-space representations, as we elaborate at the end of this section.
By using a1:T as pseudo observations, our update step corresponds to the classical Kalman filter/smoother algorithm because we do not need to linearise the observation model (see App. A.5). In combination with the neural linearisation approach, we can now analytically compute the filtered and smoothed distributions, pψ(zt|a1:t,u1:t−1) and pψ(zt|a1:T ,u1:T−1), respectively. Note, however, that these are generally not optimal because we have a nonlinear Gaussian system that is locally linearised. As a result, we obtain the recognition model (smoother version):
q(z1:T ,a1:T |x1:T ,u1:T ) = T∏ t=1 pψ(zt|a1:T ,u1:T−1) T∏ t=1 qφ(at|xt). (13)
Disentangling Static and Dynamic Features In the context of latent-variable models, disentanglement typically means that different features are represented by different dimensions in latent space [e.g. 10]. DSSMs learn a representation of the system’s state in the latent space. It can usually be split into static and dynamic features, e.g. the position and velocity of a robot arm, where the position can be inferred from a single frame, while the velocity requires a sequence of frames.
The EKVAE can be used to disentangle static and dynamic features due to its architecture: static features are represented separately by the auxiliary variables at, which are learned via the encoder– decoder pair. By defining H in pψ(at| zt) (Eq. (12)) as rectangular identity matrix (cf. Fig. 5),
H = (δij) ∈ RDa×Dz , (14)
the model learns a latent representation where the first Da dimensions of zt correspond to static features at ∈ RDa , such that z(d)t = a (d) t for d = 1, 2, . . . , Da. The remaining Dz −Da dimensions of zt represent dynamic features, as we verify in Sec. 6.4.

heading: 4.3 Integration With the CO Framework: Deriving Distortion and Rate

text:
To integrate the EKVAE with the CO framework introduced in Sec. 3, we define in the following the distortion D(θ, φ) and rateR(ψ, φ, ψ0, φ0) based on the transition, observation, and recognition model in Eqs. (9, 11, 13)—and the VHP in Eq. (6). A detailed derivation is provided in App. A.6.
The distortion is simply defined by the encoder–decoder pair:
D(θ, φ) = − T∑ t=1 Eqφ(at|xt) [ log pθ(xt|at) ] . (15)
Deriving the rate is more complicated: (i) we need to perform a sample-based optimisation of transition parameters ψ. This is especially crucial for gψ(zt,ut) in Eq. (10), where an optimisation solely via extended Kalman smoothing, i.e. via deterministic mean values (cf. App. A.5), does not cover the range of application and would therefore result in a poorly trained transition model. (ii) Our recognition model q(z1:T ,a1:T |x1:T ,u1:T ) (cf. Eq. (13)) does not contain the computationally more expensive pairwise smoothed distributions pψ(zt, zt−1|a1:T ,u1:T−1) but only smoothed distributions pψ(zt|a1:T ,u1:T−1). However, an optimisation of pψ(zt| zt−1,ut−1) based on samples (zt, zt−1) from smoothed distributions would lead to an inaccurate transition model.
To address these issues, we use the rate Eqφ(a1:T |x1:T ) [∑T t=1 log qφ(at|xt)
pψ(at| a1:t−1,u1:t−1)
] as our starting
point. But instead of computing pψ(at|a1:t−1,u1:t−1) analytically, we solve the corresponding integral (derived from the Bayesian filtering equations, see App. A.6.2) only w.r.t. zt in closed form and marginalise zt−1 via Monte Carlo integration:
log pψ(at|a1:t−1,u1:t−1) = log ∫∫ pψ(at| zt)pψ(zt| zt−1,ut−1)pψ(zt−1|a1:t−1,u1:t−2)dztdzt−1
≥ Epψ(zt−1| a1:T ,u1:T−1) [ log
pψ(at| zt−1,ut−1) pψ(zt−1|a1:t−1,u1:t−2) pψ(zt−1|a1:T ,u1:T−1)
] . (16)
In this context, the distribution pψ(at| zt−1,ut−1) plays a crucial role as it includes all transition parameters ψ and decouples zt from zt−1. It therefore allows a sample-based optimisation of the transition model on the basis of the smoothed distribution pψ(zt−1|a1:T ,u1:T−1). As a result—the complete derivation can be found in App. A.6.2—we obtain the following rate (smoother version):
R(ψ, φ, ψ0, φ0) = Eq(z1:T ,a1:T |x1:T ,u1:T ) [ Rinitial(ψ, φ, ψ0, φ0;a1:T , z1)
+ T∑ t=2 ( log qφ(at|xt) pψ(at| zt−1,ut−1) + log pψ(zt−1|a1:T ,u1:T−1) pψ(zt−1|a1:t−1,u1:t−2) )] , (17)
where the empirical Bayes prior introduced in Sec. 3.1 is learned via
Rinitial(ψ, φ, ψ0, φ0;a1:T , z1) = log qφ(a1|x1) pψ(a1| z1) + log pψ(z1|a1:T ,u1:T−1)−FVHP(ψ0, φ0; z1).
(18)
Note that the log distributions in Eq. (17) and (18) can be expressed as closed-form KL divergences (see App. A.6.2). The distortion (Eq. (15)) and rate (Eq. (17)) now allow us to define the Lagrangian of the CO problem in Eq. (7) and thus to integrate the EKVAE with our CO framework.

heading: 5 Related Work

text:
A popular method for avoiding local optima when maximising the ELBO is referred to as annealing [e.g. 4]. Here, the rate is multiplied by a weighting term β that is increased from 0 to 1 during training. However, such predefined schedules often prove to be suboptimal, as we show in Tab. 1. For this reason, we extend the VAE-based approach in [16] to DSSMs by deriving a general Lagrangian formulation of the sequential ELBO on the basis of distortion and rate. This allows to represent the above weighting term by a Lagrange multiplier λ = 1/β, which is updated based on the reconstruction quality. Our proposed optimisation algorithm builds on [16] and includes several modifications to facilitate learning the underlying system dynamics, as we detail in Sec. 3.2.
Many DSSMs [e.g. 17, 8] use simple Gaussian prior/initial distributions, resulting in less accurate transition models (cf. Sec. 6.2 and App. A.7.2). In the empirical Bayes approach of [13], a separate recognition model learns an initial pseudo state ζ ∼ p(ζ|x1:T ,u1:T ), which is then mapped to z1 = f(ζ) through a neural network. By contrast, the VAE-like empirical Bayes method (VHP) [16] that we use in our CO framework can directly substitute a Gaussian p(z1) without further restrictions on the model architecture (cf. Sec. 3.1), which moreover leads to better results (see Tab. 1).
Since the introduction of stochastic gradient variational Bayes [15, 24], various extensions have been proposed for learning DSSMs via amortised variational inference, where, in contrast to the EKVAE (ours), classic Bayesian filtering/smoothing is approximated/replaced by deterministic RNNs [17, 13, 9, 5, 7, 29, 6]. Two popular examples—which we evaluate and integrate in our CO framework—are deep Kalman filters/smoothers (DKF/DKS) [17] and deep variational Bayes filters/smoothers (DVBF/DVBS) [13, 14]. Krishnan et al. [17] define two different recognition models based on uni-/bidirectional RNNs that parametrise the approximate filtered/smoothed distribution. In [13, 14], the approximate posterior is obtained by sharing parameters between the recognition and transition model; and an RNN is used for the initial time step, as described above. Although, DVBF/DVBS uses the same locally-linear transition model as the EKVAE, it does not take advantage of closed-form Bayesian inference, leading to a less accurate dynamic model, as we verify in Tab. 1.
Previous work has shown that hidden states of RNNs [8, 21] or probabilistic switch variables [18] can be used to predict the parameters of a (time-inhomogeneous) linear SSM in order to enable closedform Bayesian inference. The Kalman VAE (KVAE) [8], for example, uses an auxiliary-variable model p(x1:T ,a1:T , z1:T |u1:T ) = p(x1:T |a1:T ) p(a1:T | z1:T ) p(z1:T |u1:T ) and is based on linear Gaussian p(at| zt,ht−1) and p(zt| zt−1,ht−1,ut−1), whose model parameters are conditioned on a deterministic hidden state ht−1 = LSTM(a1:t−1) for modelling nonlinear dynamics. This allows analytically computing the posterior by Kalman filtering/smoothing. In the EKVAE, we use a similar auxiliary-variable model as in [8], but dispense with RNNs/switch variables by choosing a transition model with a nonlinear dependence on zt−1. To compute the posterior, we leverage the concept of extended Kalman filtering/smoothing (cf. Sec. 4). This is beneficial as, for example, the LSTM-based deterministic path in the transition of [8] leads to a less accurate dynamic model and a non-Markovian state-space, meaning that not all information about the state is encoded in zt (see Sec. 6.3).

heading: 6 Experiments

text: We validate our approach on image data of a moving pendulum and on the reacher environment of Deepmind’s control suite [26], where we use angle as well as high-dimensional image data. The pendulum dataset was originally introduced in [13] and consists of 500 sequences with 15 images each, which have a size of 16× 16 pixels. The reacher dataset consists of 2000 sequences with 30 time steps each. We use two versions in our experiments: (i) partially observed system states, i.e. the angles of the first and second joint; and (ii) RGB images of 64× 64 pixels in size. In our experiments, we use smoothing posteriors. This leads to more precise models [cf. 2] and allows inferring an accurate state-space representation of partially observed systems already in the initial time step. In Sec. 6.1, we demonstrate our CO framework on the example of deep Kalman smoothers (DKSs, see App. A.3) [17]; and show in Sec. 6.2 that it significantly improves learning the system dynamics on the example of DKSs, deep variational Bayes smoothers (DVBSs, see App. A.4) [13, 14], and EKVAEs (ours). Furthermore, we verify in Sec. 6.3 that RNN-based transition models, as in KVAEs [8] and RSSMs [9], lead to a non-Markovian state-space. In Sec. 6.4, we show the benefits of disentangled state-space representations for model-based reinforcement learning.

heading: 6.1 Demonstrating the CO Framework on the Example of DKS

text:
Karl et al. [13] have shown that DKS is not capable of learning the angular velocity of the pendulum, i.e. to accurately predict the system, when trained classically or with annealing. In the following, we show that our CO framework solves this problem. We refer to the resulting model as VHP-DKS (CO), implying that it is trained via CO with the VHP as part of the model (see App. A.3 for the derivation).
Fig. 1 shows the optimisation process of VHP-DKS (CO): the model learns the underlying system dynamics, which is indicated by the barrel shape of the state-space representation [cf. 28, 13] and verified in Sec. 6.2. Complementary to this, we demonstrate in Fig. 2 that the VHP significantly improves the quality of generated sequences (no broken generative model) by learning a prior that matches EpD [ q(z1|x1:T ,u1:T ) ] —and show that annealing [4], in contrast to our CO framework, does not enable DKS to infer the angular velocity, which confirms the experimental findings in [13].

heading: 6.2 The Influence of the Learned State-Space Representation on the Prediction Accuracy

text:
High ELBO values do not imply the model can accurately predict the observed system, as we show in Tab. 1. Our CO framework solves this problem: it improves system identification, leading to a significant increase in the prediction accuracy of the models (note the impact of the VHP). We evaluate if a system has been identified based on the correlation between inferred and ground-truth states—i.e. rotation angles and angular velocities—which is measured by R2 of an OLS regression [cf. 13]. The prediction accuracy is evaluated by the MSE of 500 predicted sequences (pendulum/reacher: 15/30 times steps), conditioned on { z (n) 1 ∼ q(z1|x (n) 1:5 ,u (n) 1:4 ) }500 n=1
. This allows us to additionally verify the quality of the learned state-space representation in the initial time step. The EKVAE outperforms DKS and DVBS w.r.t. prediction accuracy and is even capable of identifying the dynamical system of reacher on the basis of 64× 64 pixels RGB images (see Fig. 3 and Tab. 1c).
Supplementary to Tab. 1, we provide in App. A.7.2 (i) a statistic evaluation of different annealing schedules compared with CO, which is based on 25 runs each; (ii) visualisations of the state-space representations (initial time step and entire sequence) learned by the different models; and (iii) further evaluations including reconstructed, predicted, and generated sequences.

heading: 6.3 Comparison With the KVAE: Limitations of RNN-Based Transition Models

text:
RNN-based transition models can lead to a non-Markovian state space, i.e. not all information about the state is encoded in zt, but partially in the RNN. This can significantly restrict (Bayesian) filtering and smoothing, resulting in a lower prediction accuracy of the model, as we show in Tab. 2 and Fig. 4.
The KVAE uses the transition model p(zt+1| zt,ht,ut), where ht=LSTM(a1:t) (see Sec. 5). As shown in Tab. 2, this leads to a lower prediction accuracy compared to the EKVAE (cf. Tab. 1a); and predictions conditioned on the smoothed {z1 ∼ p(z1|a1:5,u1:4), h1} are significantly less accurate than predictions conditioned on the filtered {z5 ∼ p(z5|a1:5,u1:4), h5}, which we explain in Fig. 4 and App. A.7.3. Note that the same applies to the RSSM, as we detail in App. A.7.3.

heading: 6.4 Encoding Rewards: Policy Learning With Disentangled State-Space Representations

text:
The EKVAE can learn state-space representations where static and dynamic features are disentangled (see Sec. 4.2). In the context of model-based reinforcement learning, these are often position and velocity, as we demonstrate on the example of pendulum (Fig. 5) and reacher (Fig. 3).
Such a disentangled representation allows us to use observations for encoding a goal position pg = a through q(a|x) or a goal velocity vg through p(z|a1:T ,u1:T−1) q(a1:T |x1:T ), where vg ∈ RDv is represented by the last Dv = Dz −Da dimensions of z (cf. Sec. 4.2). Therefore, if rewards are not available, the EKVAE can be used for defining reward functions based on an encoded pg or vg , that target dimensions in zt either representing the position or the velocity: rpost (zt,pg) = − ∑Da d=1 ( z (d) t − p (d) g )2 or rvelt (zt,vg) = − ∑Dv d=1 ( z (Da+d) t − v (d) g )2 , where the negative mean squared error is a natural choice motivated by the Euclidean distance metric [cf. 16]. Depending on the task, we can use either rpost (zt,pg) or r vel t (zt,vg) to learn a policy πω(ut| zt) by
maximising J(ω) = ∑H−1 t=1 Eπω(ut| zt) Ep(zt+1| zt,ut) [ rt+1 ] , where H is the planning horizon.
In Fig. 6 and 7, we use the above method to validate the EKVAE in the context of model-based reinforcement learning. Our results show that the EKVAE allows learning accurate policies without having access to (external) rewards. We provide experimental details and further results in App. A.7.4.

heading: 7 Conclusion

text:
In this paper, we have dealt with the question of how to learn DSSMs to obtain accurate predictions of observed dynamical systems. We have addressed the learning problem by proposing a CO framework for generic DSSMs. To this end, we have derived a general Lagrangian formulation of the sequential ELBO on the basis of distortion and rate—and extended the empirical Bayes prior (VHP) and the associated optimisation algorithm introduced in the context of VAEs to DSSMs. Building upon the CO framework, we have introduced the EKVAE, which combines extended Kalman filtering/smoothing with amortised variational inference and a neural linearisation approach.
Our experimental evaluations have demonstrated that applying the proposed CO framework to established DSSMs (e.g. DKF/DKS and DVBF/DVBS) facilitates system identification, with the VHP avoiding over-regularisation and broken generative models. The result is a substantial increase in prediction accuracy. In this context, we have shown that the EKVAE achieves a significantly higher prediction accuracy than state-of-the-art (RNN-based) models. Furthermore, we have shown that the EKVAE can learn disentangled position–velocity representations and demonstrated how these can be used for model-based reinforcement learning to define/encode reward functions and learn policies.

heading: Acknowledgements

text:
We would like to thank Maximilian Karl and Djalel Benbouzid for valuable feedback and discussions.
Funding Transparency Statement
None of the authors received related third party funding or third party support during the 36 months prior to the submission of this work. None of the authors had financial relationships with entities that could potentially be perceived to influence the submitted work in the 36 months prior to submission.

title: Latent Matters: Learning Deep State-Space Models


summaries: **(motivation)** Many approaches exist to generate counterfactual images, but what should be done when no training data is available? Current approaches lead to adversarial-looking approaches, which are not satisfying enough as an explanation.  **(approach)** The authors use Deep Inversion and a novel objective that ensures the data stays in distribution and propose a method to evaluate the quality of counterfactuals.  **(experiments)** The authors run experiments on Celebi and ISIC.

summaries: Summary: The paper proposed a novel approach to achieve deep model inversion for a pre-trained deep classifier using the classifier alone, without having access to the training dataset or generative models. Given a query image, the proposed deep model inversion generates a counterfactual explanation that is realistic looking and produces the desired classification outcome.    Earlier work on image synthesis by inverting deep classifiers either creates unrealistic images or creates perturbation of query image that successfully changes the classification decision but resulting pixel-level manipulations are not interpretable/meaningful (similar to an adversarial attack).   The proposed approach has four essential components: 1.	Counterfactual image is synthesized using an image prior. The authors experimented with an un-trained UNET-based model as a deep image prior (DIP) and Fourier mapping with SIREN activation as implicit neural representation (INR) for mapping noise to the counterfactual image. 2.	To ensure high similarity between counterfactual and query images, the authors used ISO to minimize discrepancies over the entire image and LSO to minimize differences between embeddings extracted from different layers of the deep classifier. 3.	To ensure that the generated counterfactual image is realistic and lies on the data manifold, the proposed model uses a manifold consistency (MC) loss. The MC-loss is quantified using an auxiliary loss predictor function in DEP or the radial basis kernel similarity in DUQ. Both methods provide ways to capture epistemic uncertainty due to out-of-distribution (OOD) samples. 4.	The model uses a functional consistency loss to ensure the counterfactual image produces a desired outcome from the deep classifier. Finally, the model uses progressive optimization to learn the image prior layer-by-layer.

summaries: The paper proposes a deep inversion approach that uses a pre-trained classifier to generate the counterfactual explanations of a given image. In other words, the method introduces small, discernible perturbations in an image that changes its class while keeping it realistic. The method makes use of certain inductive biases such as strong image priors, manifold consistency objective, and progressive optimization strategy to generate those counterfactual explanations.

summaries: In this paper authors propose to generate counterfactuals from a query image and a given trained deep classifier. One is interested to find out a semantically reasonable alternative explanation, where image would be classified from class y to y'. This is all performed without the help of an explicit generator (such as GAN).

abstractText: Explanation techniques that synthesize small, interpretable changes to a given image while producing desired changes in the model prediction have become popular for introspecting black-box models. Commonly referred to as counterfactuals, the synthesized explanations are required to contain discernible changes (for easy interpretability) while also being realistic (consistency to the data manifold). In this paper, we focus on the case where we have access only to the trained deep classifier and not the actual training data. While the problem of inverting deep models to synthesize images from the training distribution has been explored, our goal is to develop a deep inversion approach to generate counterfactual explanations for a given query image. Despite their effectiveness in conditional image synthesis, we show that existing deep inversion methods are insufficient for producing meaningful counterfactuals. We propose DISC (Deep Inversion for Synthesizing Counterfactuals) that improves upon deep inversion by utilizing (a) stronger image priors, (b) incorporating a novel manifold consistency objective and (c) adopting a progressive optimization strategy. We find that, in addition to producing visually meaningful explanations, the counterfactuals from DISC are effective at learning classifier decision boundaries and are robust to unknown test-time corruptions.

authors:
S. Verma
J. Dickerson
K. Hines

title: Counterfactual explanations for machine learning: A review

venue: arXiv preprint arXiv:2010.10596

authors:
RR. Selvaraju
M. Cogswell
A. Das
R. Vedantam
D. Parikh
D. Batra

title: Grad-cam: Visual explanations from deep networks via gradient-based localization

venue: Proceedings of the IEEE international conference on computer vision, pages 618–626

authors:
H. Lakkaraju
N. Arsov
O. Bastani

title: Robust black box explanations under distribution shift

venue: International Conference on Machine Learning (ICML)

authors:
A. Shrikumar
P. Greenside
A. Kundaje

title: Learning important features through propagating activation differences

venue: International Conference on Machine Learning, pages 3145–3153. PMLR

authors:
MT Ribeiro
S. Singh
C. Guestrin

title: Why should i trust you?" explaining the predictions of any classifier

venue: Proc. of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 1135–1144

authors:
MT. Ribeiro
S. Singh
C. Guestrin

title: Anchors: High-precision model-agnostic explanations

venue: Proceedings of the AAAI Conference on Artificial Intelligence, volume 32

authors:
AV. Looveren
J. Klaise

title: Interpretable counterfactual explanations guided by prototypes

venue: Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 650–665. Springer

authors:
A. Dhurandhar
P. Chen
R. Luss
C. Tu
P. Ting
K. Shanmugam
P. Das

title: Explanations based on the missing: Towards contrastive explanations with pertinent negatives

venue: Advances in Neural Information Processing Systems, volume 31

authors:
S. Sumedha
P. Brian
C. Junxiang
B. Kayhan

title: Explanation by progressive exaggeration

venue: International Conference on Learning Representations

authors:
Y. Goyal
Z. Wu
J. Ernst
D. Batra
D. Parikh
S. Lee

title: Counterfactual visual explanations

venue: International Conference on Machine Learning, pages 2376–2384. PMLR

authors:
A. Mordvintsev
C. Olah

title: and M

venue: Tyka. Inceptionism: Going deeper into neural networks

authors:
A. Mahendran
A. Vedaldi

title: Visualizing deep convolutional neural networks using natural pre-images

venue: International Journal of Computer Vision, 120(3):233–255

authors:
H. Yin
P. Molchanov
JM. Alvarez
Z. Li
A. Mallya
D. Hoiem
NK. Jha
J. Kautz

title: Dreaming to distill: Data-free knowledge transfer via deepinversion

venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8715–8724

authors:
S. Ioffe
C. Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: International conference on machine learning, pages 448–456. PMLR

authors:
Geoffrey H
V. Oriol
D. Jeffrey

title: Distilling the knowledge in a neural network

venue: In NIPS Deep Learning and Representation Learning Workshop,

authors:
D. Ulyanov
A. Vedaldi
V. Lempitsky

title: Deep image prior

venue: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 9446–9454

authors:
V. Sitzmann
J. Martel
A. Bergman
D. Lindell
G. Wetzstein

title: Implicit neural representations with periodic activation functions

venue: Advances in Neural Information Processing Systems, 33

authors:
D. Yoo
I.S Kweon

title: Learning loss for active learning

venue: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 93–102

authors:
JJ. Thiagarajan
V. Narayanaswamy
R. Anirudh
PT. Bremer
A. Spanias

title: Accurate and robust feature importance estimation under distribution shifts

venue: Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 7891–7898

authors:
J. Van Amersfoort
L. Smith
YW. Teh
Y. Gal

title: Uncertainty estimation using a single deep deterministic neural network

venue: International Conference on Machine Learning, pages 9690–9700. PMLR

authors:
A. Dosovitskiy
T. Brox

title: Inverting visual representations with convolutional networks

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4829–4837

title: Understanding deep image representations by inverting them

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5188–5196

authors:
D. Nemirovsky
N. Thiebaut
Y. Xu
A. Gupta

title: Countergan: Generating realistic counterfactuals with residual generative adversarial nets

venue: arXiv preprint arXiv:2009.05199

authors:
A. Sauer
A. Geiger

title: Counterfactual generative networks

venue: ICLR

authors:
O. Ronneberger
P. Fischer
T. Brox

title: U-net: Convolutional networks for biomedical image segmentation

venue: International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer

authors:
Y. Chen
S. Liu
X. Wang

title: Learning continuous image representation with local implicit image function

venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8628–8638

authors:
M. Tancik
PP. Srinivasan
B. Mildenhall
S. Fridovich-Keil
N. Raghavan
U. Singhal
R. Ramamoorthi
JT. Barron
R. Ng

title: Fourier features let networks learn high frequency functions in low dimensional domains

venue: NeurIPS

authors:
Y. LeCun
L. Bottou
Y. Bengio
P. Haffner

title: Gradient-based learning applied to document recognition

venue: Proceedings of the IEEE, 86(11):2278–2324

authors:
G. Daras
J. Dean
A. Jalal
AG. Dimakis

title: Intermediate layer optimization for inverse problems using deep generative models

venue: arXiv preprint arXiv:2102.07364

authors:
T. Karras
S. Laine
M. Aittala
J. Hellsten
J. Lehtinen
T. Aila

title: Analyzing and improving the image quality of stylegan

venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8110–8119

authors:
Z. Liu
P. Luo
X. Wang
X. Tang

title: Deep learning face attributes in the wild

venue: In Proceedings of International Conference on Computer Vision (ICCV),

authors:
N. Codella
V. Rotemberg
P. Tschandl
M.E. Celebi
S. Dusza
D. Gutman
B. Helba
A. Kalloo
K. Liopyris
M. Marchetti

title: et al

venue: Skin lesion analysis toward melanoma detection 2018: A challenge hosted by the international skin imaging collaboration (isic). arXiv preprint arXiv:1902.03368

authors:
P. Tschandl
C. Rosendahl
H. Kittler

title: The ham10000 dataset

venue: a large collection of multisource dermatoscopic images of common pigmented skin lesions. Scientific data, 5:180161

authors:
K. He
X. Zhang
S. Ren
J. Sun

title: Deep residual learning for image recognition

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770– 778

authors:
DP. Kingma
J. Ba

title: Adam: A method for stochastic optimization

venue: International Conference on Learning Representations, San Diego, California

authors:
C. Chang
E. Creager
A. Goldenberg
D. Duvenaud

title: Explaining image classifiers by counterfactual generation

heading: 1 Introduction

text:
With the growing need for deploying deep black-box models into critical decision-making, there is an increased emphasis on explainability methods that can reveal intricate relationships between data signatures (e.g., image features) and predictions. In this context, the so-called counterfactual (CF) explanations [1] that synthesize small, interpretable changes to a given image while producing desired changes in model predictions to support user-specified hypotheses (e.g., progressive change in predictions) have become popular. Though counterfactual explanations provide more flexibility over conventional techniques, such as feature importance estimation [2, 3, 4, 5, 6], by exploring the vicinity of a query image, an important requirement to produce meaningful counterfactuals is to produce discernible local perturbations (for easy interpretability) while being realistic (close to the underlying data manifold). Consequently, existing approaches rely extensively on pre-trained generative models to synthesize plausible counterfactuals [1, 7, 8, 9, 10]. By design, this ultimately restricts their utility to scenarios where one cannot access the training data or pre-trained generative models, for example, due to privacy requirements commonly encountered in many practical applications.
In this paper, we focus on the problem where we have access only to trained deep classifiers and not the actual training data or generative models. Synthesizing images from the underlying data distribution
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
by inverting a deep model, while not requiring access to training data, is a well investigated topic of research. For e.g., Deep Dream [11] synthesizes class-conditioned images by manipulating a noisy image directly in the space of pixels (or more formally Image Space Optimization (ISO)) constrained by image priors such as total variation [12] to regularize this ill-posed inversion. However, Deep Dream is known to produce images that look unrealistic, often very different from the training images, thus limiting their use in practice. Consequently, Yin et al., proposed DeepInversion [13] that performs image synthesis in the latent space of a pre-trained classifier (Latent Space Optimization (LSO)) and leverages layer-specific statistics (from batchnorm [14]) to constrain the images to be consistent with the training data distribution. This was showed to produce higher-quality images, particularly in the context of performing knowledge distillation [15] using the synthesized images.
Proposed Work. In contrast, this work aims to develop a deep model inversion approach that generates counterfactual explanations by exploring the vicinity of a given query image, instead of synthesizing an arbitrary realization from the entire image distribution. As illustrated in the example in Figure 1, existing deep inversion methods are ineffective when natively adopted for counterfactual generation. Due to use of weak priors, and the severely ill-posed nature of the problem, it introduces irrelevant pixel manipulations that easily satisfy the desired change in prediction. Hence, we propose DISC (Deep Inversion for Synthesizing Counterfactuals) that improves upon conventional deep model inversion by utilizing: (i) stronger image priors through the use of deep image priors [16]
(DIP) and implicit neural representations [17] (INR); (ii) a novel manifold consistency objective that ensures the counterfactual remains close to the underlying manifold; and (iii) a progressive optimization strategy to effectively introduce discernible, yet meaningful, changes to the query image.
From Figure 1, we find that our approach produces meaningful image manipulations, in order to change the prediction to the smiling class, while other deep inversion strategies cannot. Using empirical studies, we show that DISC consistently produces visually meaningful explanations, and that the counterfactuals from DISC are effective at learning model decision boundaries and are robust to unknown test-time corruptions.
Our Contributions.
1. A general framework to produce counterfactuals on-the-fly using deep model inversion;
2. Novel objectives to ensure consistency to the data manifold. We explore two different strategies based on direct error prediction [18, 19] and deterministic uncertainty estimation [20];
3. A progressive optimization strategy to introduce discernible changes to a given query image, while satisfying the manifold consistency requirement;
4. A classifier discrepancy metric to evaluate the quality of counterfactuals;
5. Empirical studies using natural image and medical image classifiers to demonstrate the effectiveness of DISC over a variety of baselines and ablations.

heading: 2 Related Work

text:
Image Synthesis from Classifiers. Inverting a pre-trained deep model is a popular strategy for image synthesis in scenarios where there is no access to the underlying training data. While deep model inversion-based methods such as Deep Dream [11] and DeepInversion [13] have been successful in generating class-conditioned images, there have been other extensions to such approaches. For example, Dosovitskiy et al. [21] proposed to invert representations of a pre-trained CNN to obtain insights about what a deep classifier has learned. On similar lines, Mahendran et al. [12, 22] addressed the problem of pre-image recovery, which in essence attempts to recover an arbitrarily
encoded representation (in the latent space of a classifier) to a realization on the (unknown) image manifold and to enable model diagnosis. Ulyanov et al. [16] improved upon this ill-posed inversion by utilizing a strong image prior in the form of deep image priors (DIP). They also explored the related problem of activation-maximization [16], where the goal is to generate an image that maximizes the activation of a given output neuron in a pre-trained classifier, and demonstrated the effectiveness of DIP. Despite the effectiveness of the deep model inversion methods for image synthesis, we find that such methods cannot be natively adopted for CF generation and are insufficient for producing meaningful pixel manipulations to a given query image.
Counterfactual Generation. Existing methods extensively rely on generative models to provide counterfactuals that explain the decisions of a black-box model. Examples including CounteRGAN [23], Counterfactual Generative Networks (CGNs) [24] and the methods reported in [7, 8], have clearly demonstrated the use of generative models in synthesizing high-quality CFs for any user-specified hypothesis on the predictions. However, this requirement of access to training data or generative models can be infeasible in practical scenarios, for e.g., restrictions arising due to privacy requirements. In contrast, our approach formulates the problem of counterfactual generation using deep model inversion, and can produce meaningful counterfactuals using only the trained classifier.

heading: 3 Proposed Approach

text:
In this section, we describe our approach for counterfactual generation that improves upon deep model inversion, and introduce a new classifier discrepancy metric for evaluating CFs. There are four key components that are critical to designing classifier-based counterfactual generators: (i) choice of metric for semantics preservation; (ii) choice of image priors to regularize image synthesis; (iii) manifold consistency to ensure that the synthesized counterfactual lies close to the true data manifold; and (iv) progressive optimization strategy to introduce gradual meaningful changes to a query image. Figures 2a and 2b provide overview of DISC implemented with deep image prior and INR-based for regularizing the counterfactual optimization process.
In its simplest form, for a given query x, a counterfactual explanation can be obtained as follows: arg min
x̄ d(x̄, x) s.t. F(x̄) = ȳ, x̄ ∈M(x), (1)
where x̄ = C(x) is a counterfactual explanation for x, F is a pre-trained classifier model,M denotes the data manifold and ȳ = y + δ is the desired change in the prediction. The metric d(., .) measures the discrepancy between the query image and the counterfactual (i.e., semantics preservation).

heading: 3.1 Choice of Metric for Semantics Preservation

text:
We can measure the discrepancy between a query and its CF explanation, d(x̄, x), in the pixel space or in the latent space of the classifier. We now describe the high-level formulation of deep model inversion-based CF generation.
a. Image Space Optimization (ISO). ISO for counterfactual generation involves the ill-posed optimization of an input x̄ directly in the pixel space to generate an image semantically similar to the query x, while being consistent with a user-hypothesis on the prediction, ȳ. Strategies such as Deep Dream [11] perform ISO to synthesize artistic variations of images. Mathematically,
arg min x̄ d(x̄, x) +R(x̄) s.t. F(x̄) = ȳ. (2)
whereR(x̄) is a suitable image prior to regularize the optimization.
b. Latent Space Optimization (LSO). LSO refers to the ill-posed problem of inverting an arbitrarily encoded representation from the latent space of a deep classifier to a realization on the (unknown) image manifold. Let Ψl(.) denote the lth differentiable layer of the deep classifier. Then, counterfactual generation using LSO can be mathematically formulated as
arg min x̄ ∑ l d(Ψl(x̄),Ψl(x)) +R(x̄) s.t. F(x̄) = ȳ. (3)
Approaches such as DeepInversion [13] perform LSO for conditional image synthesis (though with distribution-level comparison instead of our sample-level comparison) and achieve visually superior images when compared to ISO approaches.

heading: 3.2 Choice of Image Priors

text:
As observed from (2) and (3), the choice of the regularizer or image prior R(.) is central towards regularizing and tractably solve this challenging inverse problem. A variety of image priors have been proposed in the literature, and we investigate the following in this work.
a. Total Variation + `2. Total variation [22] (TV) is a popular regularizer that encourages images to contain piece-wise constant patches while the `2 norm regularizes the range and energy of the image to remain within a given interval. The TV norm and the `2 norm are given by:
RTV (x̄) = ∑ i,j √ (x̄i,j+1 − x̄i,j)2 + (x̄i+1,j − x̄i,j)2; R`2(x̄) = √∑ i,j ‖x̄i,j‖2 (4)
b. Deep Image Priors (DIP). A Deep Image Prior (DIP) [16] leverages the structure of an untrained, carefully tailored convolutional neural network (e.g., U-Net [25]) to generate images and solve a variety of ill-posed restoration tasks in computer vision. DIP has been found to produce high-quality reconstructions, based on the key insight that the structure of the network itself can act as a regularizer. Consequently, the synthesized image is re-parameterized in terms of the weights θ of the prior model fθ i.e x̄ = fθ(z).
c. Implicit Neural Representations (INR). We also considered an alternative approach based on INR, which provide a cheap and convenient way to learn a continuous mapping from the image coordinates to the pixel values (RGB). While they have been found to be effective for image/volume rendering and designing generative models [26], we explore their use in deep model inversion (details in appendix). We build upon two key results to design our INR-based counterfactual generators:
(i) Fourier mapping: Based on NTK (neural tangent kernel) theory, [27] showed that using Fourier mapping can recover high-frequency features in low-dimensional coordinate-based image reconstruction. Hence, we use a Fourier feature mapping z to featurize 2−D input coordinates v ∈ [0, 1]2 before passing them through a coordinate-based MLP:
z(v) = [a1 cos(2πb T 1 v), a1 sin(2πb T 1 v), · · · ].
Using a set of randomly chosen sinusoids, this maps the input points to the surface of a highdimensional hypersphere. Training the MLP network on these embedded points corresponds to kernel regression with the stationary composed NTK hNTK ◦ hz, where hNTK denotes the neural tangent kernel corresponding to the MLP;
(ii) SIREN activation: In [17], Sitzmann et al. showed that periodic activation functions are better suited for recovering natural images and their derivatives, when compared to standard activation functions. More specifically, SIREN uses a sinusoid activation Φ(x) = sin(Wx + b). We find that using both a Fourier mapping coupled with SIREN activation leads to a very strong image prior.

heading: 3.3 Manifold Consistency

text:
A key constraint in (1) that is not included in the formulations in (2), (3) is the manifold consistency, and interestingly, this is not inherently satisfied in deep model inversion. Consequently, even with a strong image prior, it can synthesize images that do not belong to the original data distribution. This
can be particularly challenging when producing CFs that represent change in class labels, wherein one expects patterns specific to a target class to be emphasized. To address this challenge, we extend the formulation in (3) (and equivalently (2)) to include a manifold consistency constraint, that is defined directly based on the classifier, without assuming access to class-specific statistics in the latent space. More specifically,:
arg min x̄ λ1 ∑ l d(Ψl(x̄),Ψl(x)) + λ2Lmc(x̄;F) + λ3Lfc(F(x̄), ȳ). (5)
The first term for semantics preservation is same as that of (3) and is used to ensure that the inherent semantics of the query image is retained in the generated counterfactual (implemented as the `2 error). The second term Lmc (manifold consistency) penalizes solutions that do not lie close to the data manifold and is designed by assuming access only to the classifier F . The final loss term Lfc (functional consistency) ensures that the prediction for the counterfactual matches the desired target ȳ, e.g., categorical cross entropy. As illustrated in Figure 3, the manifold consistency objective plays a central role in deep inversion-based CF generation. Even with a strong image prior (DIP in this example) and LSO, the generator produces out-of-distribution (OOD) CFs (with missing pixels) as guided by the semantic preservation term. Though the synthesized CF produces the desired class label with the classifier F , the explanation is not interpretable. In contrast, including the Lmc term (using DEP explained next) leads to a meaningful counterfactual that automatically fills in the missing pixels. Note that, this example is different from prior-based image inpainting [16], where a known
mask is used to alter the loss function to recover the missing pixels. In this work, we explore two different strategies to implement Lmc: a. Direct Error Prediction (DEP). Recently, in [19], it was found that a loss predictor trained jointly with the classifier can be used to effectively detect distribution shifts and obtain accurate uncertainty estimates for a given sample. Assuming that, a classifier model F is trained to optimize the primary loss Lpri = LCE(F(x), y), we construct an auxiliary loss predictor G to estimate the loss s = G(x) = Lpri. Similar to [18, 19], we utilize an auxiliary loss function Laux(s, ŝ) to train the parameters of G. In particular, we adopt the contrastive loss which aims to preserve the ordering of samples based on their corresponding losses from F . Let si and sj denote the
losses of samples xi and xj , while the corresponding estimates from G are ŝi and ŝj respectively. Now,
Laux = ∑ (i,j) max ( 0,−I(si, sj).(ŝi − ŝj) + γ ) , (6)
where I(si, sj) = {
1, if si > sj , −1, otherwise.
Here γ is an optional margin parameter (γ = 1 for our implementation). The overall objective for the joint optimization of the classifier and loss predictor is given by
Ltotal = β1Lpri + β2Laux. (7)
After the models are trained, we implement the manifold consistency term Lmc directly based on the loss predictor, since the losses indicate regimes where the model fails to make an accurate prediction:
Lmc = ‖G(x̄)− s∗‖1, (8) where s∗ denotes the target loss to be achieved. In our experiments, we set s∗ as the median error estimate from DEP on a held-out validation set.
b. Deterministic Uncertainty Quantification (DUQ). The recently proposed DUQ [20] method is based on Radial Basis Function (RBF) [28] networks and has been showed to be highly effective at OOD detection. In its formulation, a model is comprised of a deep feature extractor F , an exponential
kernel function along with a set of prototypical feature vectors (or centroids) for each class. DUQ is trained by optimizing the kernel distance between the features from F and the class-specific centroids and using a moving average process to update the centroids. Once DUQ is trained, the uncertainty can be measured as the distance between the model output and the closest centroid. Details regarding DUQ training can be found in the appendix. In this work, we implement Lmc using a margin-based loss that maximizes the kernel similarity of the synthesized CF with the centroid of the target class, relative to the source class. Denoting the kernel similarity for a CF x̄ with the centroid for class y as K(F(x̄), φ(y)), where φ(y) corresponds the pre-computed centroid for class y, we define:
Lmc = max ( K(F(x̄), φ(y))−K(F(x̄), φ(ȳ)) + τ, 0 ) , (9)
which indicates that kernel similarity w.r.t. the target class ȳ should be greater than that with the source class y at least by the margin τ (set to 0.5 in our experiments).

heading: 3.4 Progressive Optimization

text: CF generation is a highly under-constrained problem, that even with a strong image prior and the proposed manifold consistency constraint, it can easily converge to trivial solutions, i.e., irrelevant image manipulations. For example, one might expect to introduce large discernible changes by reducing the penalty λ1 for semantics preservation. However, given the large solution space (defined by the number of parameters in the DIP/INR generator fθ), this often leads to unrealistic images. To circumvent this, we propose to adopt a progressive optimization strategy that gradually increases the number of layers in fθ to be optimized and steadily relaxing the penalty λ1 (by factor κ) to allow for larger, yet interpretable, changes. More specifically, denoting the number of layers in fθ by L, in each iteration we train the parameters of the first i layers (i is incremented by 1 in the subsequent iteration) while keeping the parameters of the remaining L− i layers at their initial state (details on how the layers are chosen for DIP and INR based generators can be found in the appendix). A similar strategy has been shown to be effective for ill-posed restoration tasks using large-scale generative models such as Style-GAN [29, 30]. An outline of this progressive optimization process is provided in the appendix. We find that such a progressive optimization leads to significantly better quality solutions allowing meaningful traversal from one class to another.

heading: 3.5 Evaluating Quality of CF Explanations using Classifier Discrepancy

text: A desired property in query-based explainers is that the synthesized changes are both interpretable and representative of the target class (e.g., smiling). To systematically evaluate the latter property, we propose the following synthetic experiment: Given a binary classifier F and training images, i.e., X0, belonging to Class 0, we use our CF generator to synthesize examples for Class 1, i.e., X̄1 (using Class 0 images as input), and finally build a secondary classifier Fc using [X0, X̄1]. The quality of the counterfactuals can thus be measured using the gap between the performance of F and Fc on a common test set. We refer to this score as classifier discrepancy (CD).

heading: 4 Experiment Setup

text:
Datasets. (i) CelebA Faces [31]: This dataset contains 202,599 images along with a wide-range of attributes. For our experiments, we consider 3 different attributes, namely smiling, bald and young. Note, we train a classifier for predicting each of the attributes independently. We report the results for bald and young attributes in the appendix; (ii) ISIC 2018 Skin Lesion Dataset [32]: This lesion diagnosis challenge dataset contains a total of 10, 015 dermoscopic lesion images from the HAM10000 database [33]. Each image is associated with one out of 7 disease states: Melanoma (MEL), Melanocytic nevus (MN), Basal cell carcinoma (BCC), Actinic keratosis (AK), Benign keratosis (BK), Dermatofibroma (DF) and Vascular lesion (VASC). Note, in all cases, we used a stratified 90− 10 data split to train the classifiers. Classifier Design. For all experiments, we resized the images to size 96× 96 and used the standard ResNet-18 architecture [34] to train the classifier model with the Adam optimizer [35], batch size 128, learning rate 1e− 4 and momentum 0.9. For the DEP implementation (Section 3.3), we performed average pooling on feature maps from each of the residual blocks in ResNet-18, and applied a linear layer of 128 units with ReLU activation. The hyper-parameters in (7) were set at β1 = 1.0 and
β2 = 0.5. For the case of DUQ, we set both the length scale parameter and the gradient penalty to 0.5.
Image Generator Design. For generator design, the deep image prior used the standard U-Net architecture and input noise images drawn from the uniform distribution U [−1, 1]. For INR, we chose 256 random sinusoids with frequencies bi drawn from a Gaussian distribution with mean 0 and variance 100 to compute the Fourier mapping for the input coordinates.

heading: 5 Findings

text:
LSO Offers Better Feature Manipulation Control over ISO with Different Priors. An important design component in DISC is how we compute the semantic discrepancy between query x and CF x̄ - either in the pixel space using ISO or in the latent space of the classifier using LSO. We perform a comparative analysis of their behavior in CF generation using CelebA faces, in particular when manipulating an image from the non-smiling class to be classified as smiling, by varying the choice of image priors. As observed in Figure 4, though LSO produces higher quality images compared to ISO when using a weak image prior (TV + `2), that quality gap vanishes with the use of a stronger prior, e.g., DIP. However, in terms of producing discernible changes that reflect the properties of the smiling class, neither approach is sufficient. In particular, ISO shows a higher risk of making minimal, irrelevant modifications (refer to difference images |x− x̄| in Figure 4) that drive the prediction to a desired label and hence, similar to [13], we recommend the use of LSO but with stronger image priors to design CF generators.
Manifold Consistency is Essential to Produce Meaningful Explanations. As showed in the previous experiment, deep model inversion does not produce discernible (and interpretable) image changes when applied for CF generation. In this context, we explore the impact of enforcing manifold consistency in DISC. In particular, we compare the following LSO-based CF generation implementations (with INR prior): (i) no manifold consistency; (ii) DUQ-based Lmc from (9); and (iii) DEP-based Lmc from (8). From the results in Figure 5, we find that the proposed DEP objective significantly improves over the standard LSO (with no Lmc) by introducing appropriate pixel manipulations near the mouth and cheeks in all examples and clearly represents the underlying semantics of the smiling class. In contrast, the RBF network-based DUQ performs very similar to the LSO + INR baseline and
this emphasizes the inability of the kernel similarity metric to detect mild distribution shifts in data, though it has been proven successful for detecting severely OOD samples.
INR based Priors Produce Highly Concentrated Changes over DIP. Although both DIP and INR are effective, a more rigorous comparison of those two image priors is required. For this purpose, we utilize two important evaluation metrics, namely: (i) ability to produce large discernible changes, measured using `2 error between x and x̄ in the pixel-space; and (ii) ability to produce concentrated image manipulations [36], measured by thresholding (< 0.05) the difference image |x − x̄| and determining the area of the largest bounding box that contains all the non-zero values (between 0 and 1). Ideally, explanations that have a larger MSE and consistently lower concentration are more likely to produce easily interpretable, localized changes. The results in Table 1 are obtained by randomly choosing 5000 images from the non-smiling class and synthesizing the corresponding counterfactuals for smiling. Similarly, for the ISIC2018 dataset, we used 800 randomly chosen images from the MEL class and generated CFs for NEV. The naïve LSO + TV + `2 baseline produces counterfactuals with a significantly large MSE as well as concentration score, indicating that the CFs are uninterpretable and contain irrelevant perturbations all over the image. As expected, incorporating a stronger prior improves the concentration significantly, while also being highly conservative in terms of MSE, i.e., non-discernible changes. Finally, enforcing manifold consistency using DEP achieves an optimal trade-off between the two metrics and produces meaningful CFs (Figures 6 and 7). In particular, INR based generators produce highly concentrated image manipulations.
LSO + INR + DEP Produces CFs with Low Classifier Discrepancy Score. We now evaluate the quality of the counterfactuals using the proposed classifier discrepancy (CD) score. For this purpose, we consider a random subset of 5000 images each from non-smiling (X0) and smiling (X1) classes respectively. Following the strategy described in Section 3.5, we train the classifiers F and Fc, and measure the CD score as the difference in test accuracies, Acc.(F , Xtest)− Acc.(Fc, Xtest). Similarly, we repeat a simple evaluation for ISIC data by using 800 images from class MEL asX0 and images from class NEV as X1. From Table 1, we find that, using manifold consistency along with a strong image prior produces significantly lower CD scores (0.08 with LSO + INR + DEP on CelebA), when compared with LSO without manifold consistency (0.26 with LSO + INR). In particular, using INR + DEP fairs the best and consistently produces highly meaningful counterfactuals.
DISC Explanations are Robust Under Test-Time Distribution Shifts. In several practical applications, we often encounter shifts between the train and test distributions which makes model deployment challenging. Hence, we evaluate the behavior of the proposed approach under unknown distribution shifts at test time. Note, we train the classifier on the clean CelebA faces dataset without introducing any corruptions. However, when we introduce corruptions at test-time, we find that (see Figure 8), our approach (LSO + INR + DEP) is able to robustly manipulate the appropriate regions in the query image. This can be attributed both to the ability of DEP to reflect challenging distribution shifts with higher error estimates [19] and our progressive optimization strategy to induce larger yet semantically meaningful changes (i.e., inherent noise clean up).

heading: 6 Conclusions

text: In this paper, we presented DISC, a general approach to design counterfactual generators for any deep classifier, without requiring access to the training data or generative models. We draw connections to the problem of deep model inversion and extend it to support counterfactual generation. DISC can learn a CF generator on-the-fly by leveraging different image priors and manifold consistency constraints, along with a progressive optimization strategy, to synthesize highly-plausible explanations. Future extensions to this work include exploring the use of multiple target attributes simultaneously in our optimization and applying this method to time-varying data.

heading: 7 Acknowledgements

text: This work was performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344.

title: Designing Counterfactual Generators using Deep Model Inversion