Original: Variational Inference for Soil Biogeochemical Models
Revised: Variational Inference for Soil Biogeochemical Models

Original: Anonymous Authors

Original: Abstract
Revised: Abstract

Original: Soil biogeochemical models (SBMs) are an important tool used by Earth scientists to quantify the impact of rising global surface temperatures.
Revised: Soil biogeochemical models (SBMs) are an important tool used by Earth scientists to quantify the impact of rising global surface temperatures.

Original: SBMs represent the soil carbon and microbial dynamics across time as differential equations, and inference on model parameters is conducted to project changes in parameter values under warming climate conditions.
Revised: SBMs represent the soil carbon and microbial dynamics across time as differential equations, and inference on model parameters is conducted to project changes in parameter values under warming climate conditions.

Original: Traditionally, the field has relied on MCMC algorithms for posterior inference, often implemented via probabilistic programming languages like Stan.
Revised: Traditionally, the field has relied on MCMC algorithms for posterior inference, often implemented via probabilistic programming languages like Stan.

Original: However, computational cost makes it difficult to scale MCMC methods to more complex SBM models and largescale datasets.
Revised: However, computational cost makes it difficult to scale MCMC methods to more complex SBM models and largescale datasets.

Original: In this paper, we develop variational inference methods for time-discretized
Revised: In this paper, we develop variational inference methods for time-discretized

Original: SBMs as an alternative to MCMC.
Revised: SBMs as an alternative to MCMC.

Original: We propose an efficient family of variational approximations based on Gauss-Markov distributions that leverages the temporal structure of sequential models, scaling linearly in both time and space with respect to the sequence length.
Revised: We propose an efficient family of variational approximations based on Gauss-Markov distributions that leverages the temporal structure of sequential models, scaling linearly in both time and space with respect to the sequence length.

Original: We show in experiments with simulated data and real CO 2 response ratios that our approach converges faster, and recovers posterior that more accurately captures uncertainty than previous variational methods.
Revised: We show in experiments with simulated data and real CO 2 response ratios that our approach converges faster, and recovers posterior that more accurately captures uncertainty than previous variational methods.

Original: Our black-box inference approach is designed to integrate with probabilistic programming languages to enable future scientific applications.
Revised: Our black-box inference approach is designed to integrate with probabilistic programming languages to enable future scientific applications.

Original: 1. Introduction
Revised: 1. Introduction

Original: Soil microbes play a crucial role in global carbon cycle and
Revised: Soil microbes play a crucial role in global carbon cycle and

Original: Earth ecosystem functions.
Revised: Earth ecosystem functions.

Original: Although it is widely known that rising global surface temperatures affect soil microbes, it is hard to quantify the impact.
Revised: Although it is widely known that rising global surface temperatures affect soil microbes, it is hard to quantify the impact.

Original: Soil biogeochemical models (SBMs) are used to represent the transfer of elements, such as carbon, between types of organic molecules and quantify the response of biological soil systems to global warming (Xie et al., 2020).
Revised: Soil biogeochemical models (SBMs) are used to represent the transfer of elements, such as carbon, between types of organic molecules and quantify the response of biological soil systems to global warming (Xie et al., 2020).

Original: However, there can be vast differences between the predictions of competing models.
Revised: However, there can be vast differences between the predictions of competing models.

Original: To quickly compare different hypotheses we seek fast, accurate, and easy inference methods that also maintain biological interpretability of model parameters and latent variables.
Revised: To quickly compare different hypotheses we seek fast, accurate, and easy inference methods that also maintain biological interpretability of model parameters and latent variables.

Original: SBMs represent the soil carbon and microbial dynamics across time as differential equations, and inference on model parameters is conducted to project changes in parameter values under warming climate conditions.
Revised: SBMs represent the soil carbon and microbial dynamics across time as differential equations, and inference on model parameters is conducted to project changes in parameter values under warming climate conditions.

Original: Partly due to ease of use in probabilistic programming languages (PPLs) like Stan (Carpenter et al., 2017), the field has mainly relied on MCMC methods for posterior inference of model parameters, where latent states are approximated by deterministic ordinary differential equation (ODE) solvers (Li et al., 2019; Wang et al., 2022; Xie et al., 2020).
Revised: Partly due to ease of use in probabilistic programming languages (PPLs) like Stan (Carpenter et al., 2017), the field has mainly relied on MCMC methods for posterior inference of model parameters, where latent states are approximated by deterministic ordinary differential equation (ODE) solvers (Li et al., 2019; Wang et al., 2022; Xie et al., 2020).

Original: PPLs allow a user to specify a model in an intuitive modeling language and (mostly) automate the inference process, which is important for scientific applications that require iterating over many competing models.
Revised: PPLs allow a user to specify a model in an intuitive modeling language and (mostly) automate the inference process, which is important for scientific applications that require iterating over many competing models.

Original: However, computational cost makes it difficult to scale MCMC methods to more complex, non-linear SBMs for larger datasets spanning decades.
Revised: However, computational cost makes it difficult to scale MCMC methods to more complex, non-linear SBMs for larger datasets spanning decades.

Original: Variational inference (VI; Wainwright & Jordan (2008)) has been widely successful as a faster alternative to MCMC in other large scale applications of machine learning (Gan et al., 2015; Gopalan et al., 2016; Ji et al., 2019).
Revised: Variational inference (VI; Wainwright & Jordan (2008)) has been widely successful as a faster alternative to MCMC in other large scale applications of machine learning (Gan et al., 2015; Gopalan et al., 2016; Ji et al., 2019).

Original: VI reframes the task of posterior inference as an optimization problem by minimizing the KL divergence between the true posterior and an approximate variational distribution.
Revised: VI reframes the task of posterior inference as an optimization problem by minimizing the KL divergence between the true posterior and an approximate variational distribution.

Original: For complex non-conjugate models, variational algorithms often require tedious and model-specific manual derivations, and are thus not ideal for scientific applications.
Revised: For complex non-conjugate models, variational algorithms often require tedious and model-specific manual derivations, and are thus not ideal for scientific applications.

Original: Recently, “black box” variational inference methods have been proposed that can be readily applied to general models.
Revised: Recently, “black box” variational inference methods have been proposed that can be readily applied to general models.

Original: For differentiable models with continuous latent variables, Automatic Differentiation Variational Inference (ADVI; Kucukelbir et al.
Revised: For differentiable models with continuous latent variables, Automatic Differentiation Variational Inference (ADVI; Kucukelbir et al.

Original: (2017)) provides unbiased and low variance gradient estimates through transformation of random variables and the “reparameterization trick” (Kingma & Welling, 2014;
Revised: (2017)) provides unbiased and low variance gradient estimates through transformation of random variables and the “reparameterization trick” (Kingma & Welling, 2014;

Original: Rezende et al., 2014).
Revised: Rezende et al., 2014).

Original: ADVI has been implemented as the default variational algorithm in Stan.
Revised: ADVI has been implemented as the default variational algorithm in Stan.

Original: In this paper, we propose an extension of ADVI that leverages the temporal structure of sequential models.
Revised: In this paper, we propose an extension of ADVI that lever- ages the temporal structure of sequential models.

Original: Our pro- posed variational family based on Gauss-Markov distributions captures temporal dependence efficiently and scales linearly in the length of the sequence.
Revised: Our proposed variational family based on Gauss-Markov distributions captures temporal dependence efficiently and scales linearly in the length of the sequence.

Original: We reformulate ODE models from prior work as stochastic differential equations (SDEs), as the smoother trajectories of ODEs oversimplify the substantial noise that is inherent to biological systems (Browning et al., 2020; Abs et al., 2020; Golightly & Wilkinson, 2011).
Revised: We reformulate ODE models from prior work as stochastic differential equations (SDEs), as the smoother trajectories of ODEs oversimplify the substantial noise that is inherent to biological systems (Browning et al., 2020; Abs et al., 2020; Golightly & Wilkinson, 2011).

Original: Rather than solving the SDEs directly, we use an approximation based on time discretization and express the problem as a sequential latent variable model, where we jointly optimize model parameters and latent states.
Revised: Rather than solving the SDEs directly, we use an approximation based on time discretization and express the problem as a sequential latent variable model, where we jointly optimize model parameters and latent states.

Original: We show that previous variational methods perform poorly by systematically under or over-estimating variance of the posterior.
Revised: We show that previous variational methods perform poorly by systematically under or over-estimating variance of the posterior.

Original: In experiments with simulated data and real COresponse ratios, we show that our method leads to faster convergence to better loss and variational approximations compared to several baseline methods, including standard
Revised: In experiments with simulated data and real COresponse ratios, we show that our method leads to faster convergence to better loss and variational approximations compared to several baseline methods, including standard

Original: ADVI and an amortized variational inference approach.
Revised: ADVI and an amortized variational inference approach.

Original: Several authors have examined inference in SDEs and related sequence models.
Revised: Several authors have examined inference in SDEs and related sequence models.

Original: Ryder et al.
Revised: Ryder et al.

Original: (2018) proposed an amortized VI algorithm specifically designed for SDEs, where variational distributions are parameterized by neural networks.
Revised: (2018) proposed an amortized VI algorithm specifically designed for SDEs, where variational distributions are parameterized by neural networks.

Original: Other authors have recently proposed continuoustime methods for sequential and time series models (Chen et al., 2018; Rubanova et al., 2019; Kidger et al., 2020;
Revised: Other authors have recently proposed continuoustime methods for sequential and time series models (Chen et al., 2018; Rubanova et al., 2019; Kidger et al., 2020;

Original: Li et al., 2020; Schirmer et al., 2021).
Revised: Li et al., 2020; Schirmer et al., 2021).

Original: Here, we favor the simpler approach based on ADVI and discrete-time approximations.
Revised: Here, we favor the simpler approach based on ADVI and discrete-time approximations.

Original: Our method is broadly applicable to general sequential latent variable models and designed to enable integration with PPLs that allows scientists to easily and quickly design, test, and revise different models.
Revised: Our method is broadly applicable to general sequential latent variable models and designed to enable integration with PPLs that allows scientists to easily and quickly design, test, and revise different models.

Original: 2. Automatic Differentiation Variational Inference
Revised: 2. Automatic Differentiation Variational Inference

Original: Given any model with latent variable z and observation y , the goal of Bayesian inference is to infer posterior distribution p ( z | y )
Revised: Given any model with latent variable z and observation y , the goal of Bayesian inference is to infer posterior distribution p ( z | y )

Original: =
Revised: =

Original: p ( z ) p ( y | z ) p ( y ) .
Revised: p ( z ) p ( y | z ) p ( y ) .

Original: Except for very simple models, the normalizing constant p ( y ) =
Revised: Except for very simple models, the normalizing constant p ( y ) =

Original: (cid:82) z p ( z ) p ( y | z ) dz is intractable and thus approximation of the posterior is needed.
Revised: (cid:82) z p ( z ) p ( y | z ) dz is intractable and thus approximation of the posterior is needed.

Original: Variational inference seeks a variational distribution q ( z ; λ ) parameterized by variational parameter λ by maximizing the evidence lower bound (ELBO):
Revised: Variational inference seeks a variational distribution q ( z ; λ ) parameterized by variational parameter λ by maximizing the evidence lower bound (ELBO):

Original: [Equation]
Revised: [Equation]

Original: Maximizing the ELBO is equivalent to minimizing KL divergence of q ( z ; λ ) from the true posterior p ( z | y ) .
Revised: Maximizing the ELBO is equivalent to minimizing KL divergence of q ( z ; λ ) from the true posterior p ( z | y ) .

Original: Automatic Differentiation Variational Inference (ADVI; Kucukelbir et al.
Revised: Automatic Differentiation Variational Inference (ADVI; Kucukelbir et al.

Original: (2017)) is a black-box VI algorithm that can be applied to any differentiable probability models.
Revised: (2017)) is a black-box VI algorithm that can be applied to any differentiable probability models.

Original: It automates ELBO optimization of equation (1) via automatic differentiation.
Revised: It automates ELBO optimization of equation (1) via automatic differentiation.

Original: It works by first transforming the support of the latent variables z to the real coordinate space 
Revised: It works by first transforming the support of the latent variables z to the real coordinate space ˜ z

Original: [Equation] K .
Revised: = T ( z ) ∈ R K .

Original: The transformed joint density is given by
Revised: The transformed joint density is given by

Original: [Equation]
Revised: [Equation]

Original: where J (˜ z ) is the Jacobian of the inverse transformation
Revised: where J (˜ z ) is the Jacobian of the inverse transformation

Original: T − 1 .
Revised: T − 1 .

Original: The resulting ELBO in the transformed coordinates is
Revised: The resulting ELBO in the transformed coordinates is

Original: [Equation]
Revised: [Equation]

Revised: q

Revised: [Equation]

Original: , where H [ q (˜ z ; λ )] is the entropy of the variational distribution.
Revised: ˜ z [Equation] λ )] . H [ q (˜ z ; λ )] is the entropy of the variational distribution.

Original: ADVI employs Gaussian variational approximations for the transformed variable ˜ z , which may implicitly induce a non-Gaussian variational approximation in the original latent space z .
Revised: ADVI employs Gaussian variational approximations for the transformed variable ˜ z , which may implicitly induce a non-Gaussian variational approximation in the original latent space z .

Original: The Gaussian assumption allows the entropy H [ q (˜ z ; λ )] to be computed analytically.
Revised: The Gaussian assumption allows the entropy H [ q (˜ z ; λ )] to be computed analytically.

Original: A simple but naive approximation assumes a fully-factorized variational distribution with independent latent variables: q (˜ z ; λ ) =
Revised: A simple but naive approximation assumes a fully-factorized variational distribution with independent latent variables: q (˜ z ; λ ) =

Original: [Equation]
Revised: (cid:81) K k

Original: K k =1 Normal (˜ z k ; µ k , σ 2 k ) .
Revised: =1 Normal (˜ z k ; µ k , σ 2 k ) .

Original: Because σ must always be positive, we parameterize as ω = log( σ ) .
Revised: Because σ must always be positive, we parameterize as ω = log( σ ) .

Original: The set of variational parameters in the mean-field approximation is therefore λ k
Revised: The set of variational parameters in the mean-field approximation is therefore λ k

Original: = ( µ k , ω k ) .
Revised: = ( µ k , ω k ) .

Original: Another convenient yet more expressive option for the variational approximation is the multivariate or full-rank
Revised: Another convenient yet more expressive option for the variational approximation is the multivariate or full-rank

Original: Gaussian, q (˜ z ; λ ) =
Revised: Gaussian, q (˜ z ; λ ) =

Original: MultivariateNormal (˜ z ; µ, Σ) , where Σ is the covariance matrix.
Revised: MultivariateNormal (˜ z ; µ, Σ) , where Σ is the covariance matrix.

Original: Unlike the mean-field assumption, this allows arbitrary correlation structure between any pair of the latent variables.
Revised: Unlike the mean-field assumption, this allows arbitrary correlation structure between any pair of the latent variables.

Original: To ensure Σ is always positive semi-definite, we parameterize the covariance matrix as Σ = LL T , where L is a lower triangular matrix with K ( K +1) unconstrained real-valued nonzero entries.
Revised: To ensure Σ is always positive semi-definite, we parameterize the covariance matrix as Σ = LL T , where L is a lower triangular matrix with K ( K +1) unconstrained real-valued nonzero entries.

Original: In the full-rank case, the variational parameter becomes λ = ( µ, L ) .
Revised: In the full-rank case, the variational parameter becomes λ = ( µ, L ) .

Original: Now that we can freely optimize the ELBO in the real coordinate space, we require its gradient with respect to λ for stochastic gradient optimization.
Revised: Now that we can freely optimize the ELBO in the real coordinate space, we require its gradient with respect to λ for stochastic gradient optimization.

Original: To enable automatic differentiation, we push the gradient operation inside the expectation, applying the “reparameterization trick” (Rezende et al., 2014; Kingma & Welling, 2014).
Revised: To enable automatic differentiation, we push the gradient operation inside the expectation, applying the “reparameterization trick” (Rezende et al., 2014; Kingma & Welling, 2014).

Original: We use the inverse Gaussian standardization S − 1 to reparameterize ˜ z
Revised: We use the inverse Gaussian standardization S − 1 to reparameterize ˜ z

Original: = S − 1 λ ( ϵ ) as a deterministic function of the standard Gaussian noise ϵ ∼ N (0 , 1) given the variational parameters λ .
Revised: = S − 1 λ ( ϵ ) as a deterministic function of the standard Gaussian noise ϵ ∼ N (0 , 1) given the variational parameters λ .

Original: For example, the standardization in the mean-field case is
Revised: For example, the standardization in the mean-field case is

Original: [Equation]
Revised: [Equation]

Original: This allows us to approximate the gradient using the Monte Carlo estimator:
Revised: This allows us to approximate the gradient using the Monte Carlo estimator:

Original: [Equation]
Revised: [Equation]

Original: Note that since the entropy is evaluated analytically, we do not need to reparameterize ˜ z in the entropy term, as the gradient can also be computed analytically.
Revised: Note that since the entropy is evaluated analytically, we do not need to reparameterize ˜ z in the entropy term, as the gradient can also be computed analytically.

Original: 3. Soil Biogeochemical Models
Revised: 3. Soil Biogeochemical Models

Original: Soil biogeochemical models (SBMs) are differential equation models that represent the dynamics of soil pool densities (Xie et al., 2020).
Revised: Soil biogeochemical models (SBMs) are differential equation models that represent the dynamics of soil pool densities (Xie et al., 2020).

Original: We consider two classes of models: the linear conventional (CON) model and the nonlinear Allison-Wallenstein-Bradford (AWB) model (Allison et al., 2010).
Revised: We consider two classes of models: the linear conventional (CON) model and the nonlinear Allison-Wallenstein-Bradford (AWB) model (Allison et al., 2010).

Original: State variables x correspond to densities of elements in organic molecules that evolve over time t following an ODE.
Revised: State variables x correspond to densities of elements in organic molecules that evolve over time t following an ODE.

Original: CO 2 emissions at time t can be estimated as a function of these state variables.
Revised: CO 2 emissions at time t can be estimated as a function of these state variables.

Original: Noisy measurements y are collected at potentially irregular time intervals.
Revised: Noisy measurements y are collected at potentially irregular time intervals.

Original: Model parameters θ correspond to other biological elements that govern the system.
Revised: Model parameters θ correspond to other biological elements that govern the system.

Original: We define the joint distribution p ( θ, x, y )
Revised: We define the joint distribution p ( θ, x, y )

Original: =
Revised: =

Original: p ( θ ) p ( x | θ ) p ( y | x, θ ) .
Revised: p ( θ ) p ( x | θ ) p ( y | x, θ ) .

Original: The CON system consists of three state variables x t =
Revised: The CON system consists of three state variables x t =

Original: ( S t , D t , M t ) : soil organic carbon ( S ) indicates the carbon density of stable organic soil molecules; dissolved organic carbon ( D ) denotes the carbon density contained in less stable, more decomposed organic molecule types; and microbial biomass carbon ( M ) describes the carbon density encompassed by the population of soil microbial organisms.
Revised: ( S t , D t , M t ) : soil organic carbon ( S ) indicates the carbon density of stable organic soil molecules; dissolved organic carbon ( D ) denotes the carbon density contained in less stable, more decomposed organic molecule types; and microbial biomass carbon ( M ) describes the carbon density encompassed by the population of soil microbial organisms.

Original: The AWB system consists of four state variables x t =
Revised: The AWB system consists of four state variables x t =

Original: ( S t , D t , M t , E t ) , which additionally includes extracellular enzyme carbon ( E ) that signifies the carbon density tied up in the enzymes secreted by the microbial organisms to help break down organic material in the system for consumption.
Revised: ( S t , D t , M t , E t ) , which additionally includes extracellular enzyme carbon ( E ) that signifies the carbon density tied up in the enzymes secreted by the microbial organisms to help break down organic material in the system for consumption.

Original: The nonlinearity in AWB model comes from explicitly representing micorbial processes with nonlinear Michaelis-Menten functions (Wieder et al., 2015).
Revised: The nonlinearity in AWB model comes from explicitly representing micorbial processes with nonlinear Michaelis-Menten functions (Wieder et al., 2015).

Original: The diagrams in figure 1 show the interaction between state variables and the various model parameters in the CON and
Revised: The diagrams in figure 1 show the interaction between state variables and the various model parameters in the CON and

Original: AWB models.
Revised: AWB models.

Original: To more realistically capture the stochasticity inherent to biological systems, we formulate the SDE parameterization of the CON and AWB models by adding noise to the system, which we call the stochastic CON (SCON) and stochastic
Revised: To more realistically capture the stochasticity inherent to biological systems, we formulate the SDE parameterization of the CON and AWB models by adding noise to the system, which we call the stochastic CON (SCON) and stochastic

Original: AWB (SAWB) models.
Revised: AWB (SAWB) models.

Original: Concretely, we have the SDE
Revised: Concretely, we have the SDE

Original: [Equation]
Revised: [Equation]

Original: where x t is an M -dimensional vector of random variables, α ( x t , θ, t ) is an M -dimensional drift vector , β ( x t , θ, t ) is an
Revised: where x t is an M -dimensional vector of random variables, α ( x t , θ, t ) is an M -dimensional drift vector , β ( x t , θ, t ) is an

Original: M × M diffusion matrix , and W t is a standard Wiener process.
Revised: M × M diffusion matrix , and W t is a standard Wiener process.

Original: The drift and diffusion depend on θ , a D -dimensional vector of unknown parameters.
Revised: The drift and diffusion depend on θ , a D -dimensional vector of unknown parameters.

Original: Our drift vector α corresponds directly to the original ODE model (details in the supplement).
Revised: Our drift vector α corresponds directly to the original ODE model (details in the supplement).

Original: We consider two versions of diagonal-only diffusion matrices: constant diffusion β =
Revised: We consider two versions of diagonal-only diffusion matrices: constant diffusion β =

Original: Iσ (SCON-C and SAWB-C) and state-scaling diffusion β =
Revised: Iσ (SCON-C and SAWB-C) and state-scaling diffusion β =

Original: I ( σ ⊙ √ x ) (SCON-SS and SAWB-SS), where ⊙ denotes elementwise multiplication and σ ∈ θ is a model parameter that controls the noise of the dynamics.
Revised: I ( σ ⊙ √ x ) (SCON-SS and SAWB-SS), where ⊙ denotes elementwise multiplication and σ ∈ θ is a model parameter that controls the noise of the dynamics.

Original: In practice, all variables in the model are constrained to fall within biologically realistic intervals.
Revised: In practice, all variables in the model are constrained to fall within biologically realistic intervals.

Original: For example, all latent states must be positive and activation energy parameters must be between 0 and 100.
Revised: For example, all latent states must be positive and activation energy parameters must be between 0 and 100.

Original: The SDE trajectory x across a finite timespan can be discretized into a series of T equally spaced steps of length ∆ t .
Revised: The SDE trajectory x across a finite timespan can be discretized into a series of T equally spaced steps of length ∆ t .

Original: Transition densities between states at successive times are approximated as truncated Gaussian to enforce positivity:
Revised: Transition densities between states at successive times are approximated as truncated Gaussian to enforce positivity:

Original: [Equation]
Revised: [Equation]

Original: where TN ( x ; µ, Σ , a, b ) denotes the truncated Gaussian density with mean µ , covariance Σ , lower bound a , and upper bound b .
Revised: where TN ( x ; µ, Σ , a, b ) denotes the truncated Gaussian density with mean µ , covariance Σ , lower bound a , and upper bound b .

Original: The transition likelihood is p ( x | θ ) =
Revised: The transition likelihood is p ( x | θ ) =

Original: p ( x 0 | θ ) (cid:81) Tt =1 p ( x t | x t − 1 , θ ) .
Revised: p ( x 0 | θ ) (cid:81) Tt =1 p ( [Equation] t | x [Equation]

Original:  We place an informative logit-normal prior on model parameter θ based on expert knowledge, p ( θ ) =
Revised: − 1 , [Equation] We place an informative logit-normal prior on model parameter θ based on expert knowledge, p ( θ ) =

Original: LogitNormal ( θ ; µ θ , Σ θ , a θ , b θ ) .
Revised: LogitNormal ( θ ; µ θ , Σ θ , a θ , b θ ) .

Original: Here, lower and upper bounds ( a θ , b θ ) vary across different parameters.
Revised: Here, lower and upper bounds ( a θ , b θ ) vary across different parameters.

Original: Finally, we assume truncated Gaussian observations, p ( y t | x t , θ ) =
Revised: Finally, we assume truncated Gaussian observations, p ( y t | x t , θ ) =

Original: TN ( y t ; µ ( x t) ,Σ( θ ) ,  a y , b y ) , where y t is an
Revised: TN ( y t ; µ ( x t , θ ) , Σ y , a y , b y ) , where y t is an

Original: N -dimensional vector of observations, µ ( x t  ) is a function of latent states x t, andΣ( θ ) is a function of model parameters θ .
Revised: N -dimensional vector of observations, µ ( x t , θ ) is a function of latent states x t and model parameters θ , and Σ y can either be a model parameter or a known constant.

Original: For example, the observations could be noisy measurements of CO 2 , which is computed as a function of the states, or of the state variables themselves.
Revised: For example, the observations could be noisy measurements of CO 2 , whose means are computed as a function of the states (see equations (17) and (21) in the supplement), or of the state variables themselves.

Original: The lower and upper bounds ( a y , b y ) can also vary across different types of observations.
Revised: The lower and upper bounds ( a y , b y ) can also vary across different types of observations.

Original: Letting T be the time steps where observations are available, the observation likelihood is given by p ( y | x, θ ) =
Revised: Letting T be the time steps where observations are available, the observation likelihood is given by p ( y | x, θ ) =

Original: (cid:81) t ∈T p ( y t | x t , θ ) .
Revised: (cid:81) t ∈T p ( y t | x t , θ ) .

Original: 4. ADVI for Soil Biogeochemical Models
Revised: 4. ADVI for Soil Biogeochemical Models

Original: The goal of SBM inference is to project changes in parameter values under warming climate conditions.
Revised: The goal of SBM inference is to project changes in parameter values under warming climate conditions.

Original: We would like to infer the posterior of model parameters conditioned on post-warming treatment observations, p ( θ | y ) .
Revised: We would like to infer the posterior of model parameters conditioned on post-warming treatment observations, p ( θ | y ) .

Original: We can apply
Revised: We can apply

Original: ADVI by defining state variables x and model parameters θ as our latent variables, z
Revised: ADVI by defining state variables x and model parameters θ as our latent variables, z

Original: =
Revised: =

Original: ( θ, x ) .
Revised: ( θ, x ) .

Original: We further assume factorized variational distribution p ( θ, x | y )
Revised: We further assume factorized variational distribution p ( θ, x | y )

Original: ≈ q
Revised: ≈ q

Original: ( θ, x ) =
Revised: ( θ, x ) =

Original: q ( θ ) q ( x ) .
Revised: q ( θ ) q ( x ) .

Original: We proceed to infer posterior of model parameters θ and latent states of the SDE dynamics x jointly.
Revised: We proceed to infer posterior of model parameters θ and latent states of the SDE dynamics x jointly.

Original: To ensure biologically realistic values, θ is constrained between some lower and upper bounds ( a θ , b θ ) , while x is constrained to be positive.
Revised: To ensure biologically realistic values, θ is constrained between some lower and upper bounds ( a θ , b θ ) , while x is constrained to be positive.

Original: We apply the rescaled logit function to transform θ to the unconstrained real space: ˜ θ =
Revised: We apply the rescaled logit function to transform θ to the unconstrained real space: ˜ θ =

Original: T θ ( θ ) =
Revised: T θ ( θ ) =

Original: logit (cid:16) θ − a b − a (cid:17) .
Revised: logit (cid:16) θ − a b − a (cid:17) .

Original: This transformation implicitly induces a logit-normal variational approximation on θ .
Revised: This transformation implicitly induces a logit-normal variational approximation on θ .

Original: For x , we apply the inverse softplus transformation that maps positive values to the real space:
Revised: For x , we apply the inverse softplus transformation that maps positive values to the real space:

Original: ˜ x
Revised: ˜ x

Original: = T x ( x ) = log(exp
Revised: = T x ( x ) = log(exp

Original: ( x ) − 1) .
Revised: ( x ) − 1) .

Original: [Figure]
Revised: [Figure]

Original: The ADVI objective of Equation (3) becomes:
Revised: The ADVI objective of Equation (3) becomes:

Original: [Equation]
Revised: [Equation]

Original: In practice, users do not need to derive this equation for new models, since the process is automated.
Revised: In practice, users do not need to derive this equation for new models, since the process is automated.

Original: They only need to specify the generative model, any constraints on the support of the latent variables, and the choice of desired variational approximation: mean-field, full-rank, or Gauss-Markov.
Revised: They only need to specify the generative model, any constraints on the support of the latent variables, and the choice of desired variational approximation: mean-field, full-rank, or Gauss-Markov.

Original: 4.1. Gauss-Markov ADVI
Revised: 4.1. Gauss-Markov ADVI

Original: A natural choice for the variational approximation of θ is the full-rank distribution to allow correlations to be captured between any arbitrary pair of model parameters.
Revised: A natural choice for the variational approximation of θ is the full-rank distribution to allow correlations to be captured between any arbitrary pair of model parameters.

Original: The full covariance representation scales quadratically in the number of parameters, but this is fine as both models are limited to a relatively small number of parameters.
Revised: The full covariance representation scales quadratically in the number of parameters, but this is fine as both models are limited to a relatively small number of parameters.

Original: However, for the sequential latent states x , neither full-rank nor mean-field is an ideal choice.
Revised: However, for the sequential latent states x , neither full-rank nor mean-field is an ideal choice.

Original: Mean-field, although computationally very cheap (linear in the sequence length), is too unrealistic since it falsely assumes temporal independence.
Revised: Mean-field, although computationally very cheap (linear in the sequence length), is too unrealistic since it falsely assumes temporal independence.

Original: As we will show in our experiments, this may lead to biased parameter estimates.
Revised: As we will show in our experiments, this may lead to biased parameter estimates.

Original: The full-rank approximation allows arbitrary correlations between any states at any time step, but this requires a covariance matrix of size MT × MT that scales quadratically in the length of the sequence, which could be very long.
Revised: The full-rank approximation allows arbitrary correlations between any states at any time step, but this requires a covariance matrix of size MT × MT that scales quadratically in the length of the sequence, which could be very long.

Original: In addition to being costly, this also makes optimization sensitive to divergence or local optima due to having too many irrelevant parameters.
Revised: In addition to being costly, this also makes optimization sensitive to divergence or local optima due to having too many irrelevant parameters.

Original: In sequential models with first-order Markov priors and observations that are local to single time points, as is the case in
Revised: In sequential models with first-order Markov priors and observations that are local to single time points, as is the case in

Original: SDEs, we are guaranteed that the posterior is also first-order
Revised: SDEs, we are guaranteed that the posterior is also first-order

Original: Markov.
Revised: Markov.

Original: In this case, parameterizing the full covariance is wasteful as it has provably extraneous parameters.
Revised: In this case, parameterizing the full covariance is wasteful as it has provably extraneous parameters.

Original: In order tocapture more parsimoniously capture temporal dependencies, we propose to directly parameterize our variational approximation as a Markov chain.
Revised: In order to more parsimoniously capture temporal dependencies, we propose to directly parameterize our variational approximation as a Markov chain.

Original: This results in a sampling procedure that scales linearly with respect to sequence length in both time and space.
Revised: This results in a sampling procedure that scales linearly with respect to sequence length in both time and space.

Original: Consider the set of variational parameters λ t = ( µ t , A t , L t ) for t = 1 , . .
Revised: Consider the set of variational parameters λ t = ( µ t , A t , L t ) for t = 1 , . .

Original: , T , where µ t ∈ R M is the mean at time t ,
Revised: , T , where µ t ∈ R M is the mean at time t ,

Original: A t ∈ R M × M represents correlations between state variables at times t − 1 and t , and L t ∈ R M × M is a lower triangular covariance square-root matrix that represents correlations between different state variables at time t .
Revised: A t ∈ R M × M represents correlations between state variables at times t − 1 and t , and L t ∈ R M × M is a lower triangular covariance square-root matrix that represents correlations between different state variables at time t .

Original: Letting ϵ t ∼ N (0 , I M ) be standard Gaussian noise of dimension
Revised: Letting ϵ t ∼ N (0 , I M ) be standard Gaussian noise of dimension

Original: M and η 0 =
Revised: M and η 0 =

Original: L 0 ϵ 0 , then for subsequent t = 1 , ..., T :
Revised: L 0 ϵ 0 , then for subsequent t = 1 , ..., T :

Original: [Equation]
Revised: [Equation]

Original: The first line adds the temporal dependence, where the total stochasticity of ˜ x t is the sum of two sources: a linear transformation of the state variables from the previous time step
Revised: The first line adds the temporal dependence, where the total stochasticity of ˜ x t is the sum of two sources: a linear transformation of the state variables from the previous time step

Original: A t ˜ x t − 1 , and independent noise at the current time step L t ϵ t .
Revised: A t ˜ x t − 1 , and independent noise at the current time step L t ϵ t .

Original: The second line sets the mean µ t .
Revised: The second line sets the mean µ t .

Original: In general, the resulting sample x t is correlated with all previous time points x s<t with covariance Cov ( x s , x t ) =
Revised: In general, the resulting sample x t is correlated with all previous time points x s<t with covariance Cov ( x s , x t ) =

Original: (cid:81) ti
Revised: (cid:81) ti

Original: =
Revised: =

Original: s +

Revised: s +

Original: A i Var ( x s ) without having to directly parameterize the full covariance matrix.
Revised: A i Var ( x s ) without having to directly parameterize the full covariance matrix.

Original: Because the states are defined by a Markov chain, the resulting entropy is a sum of conditional entropies, H [ q (˜ x ; λ )] 
Revised: Because the states are defined by a Markov chain, the resulting entropy is a sum of conditional entropies, H [ q (˜ x ; λ )] =

Original: = (cid:80)
Revised: [Equation]

Original: Tt =1 H [ q (˜ x t | ˜ x t − 1 ; λ t )] .
Revised: Tt = [Equation] H [ q (˜ x t | ˜ x t − 1 ; λ t )] .

Original: Furthermore, the conditional entropy at time t is equal to that of a Gaussian with covariance
Revised: Furthermore, the conditional entropy at time t is equal to that of a Gaussian with covariance

Original: L t L Tt , so:
Revised: L t L Tt , so:

Original: [Equation]
Revised: [Equation]

Original: Similar to the inverse standardization, this procedure parameterizes ˜ x as a deterministic function of standard Gaussian noise, which allows the gradient to be computed using automatic differentiation of the Monte Carlo estimator in (3).
Revised: Similar to the inverse standardization, this procedure parameterizes ˜ x as a deterministic function of standard Gaussian noise, which allows the gradient to be computed using automatic differentiation of the Monte Carlo estimator in (3).

Original: Intuitively, the Gauss-Markov variational approximation is a middle ground between the mean-field and the full-rank approximations.
Revised: Intuitively, the Gauss-Markov variational approximation is a middle ground between the mean-field and the full-rank approximations.

Original: It is more expressive than the mean-field approximation by allowing arbitrary first-order correlations.
Revised: It is more expressive than the mean-field approximation by allowing arbitrary first-order correlations.

Original: It enables more efficient inference than the full-rank approximation by encoding the Markov structure of the true posterior, reducing the effective number of parameters.
Revised: It enables more efficient inference than the full-rank approximation by encoding the Markov structure of the true posterior, reducing the effective number of parameters.

Original: 4.2. Related Work
Revised: 4.2. Related Work

Original: Archer et al. (2015) proposed a black-box variational inference algorithm for other types of state space models with a Gaussian variational posterior that is parameterized by a tridiagonal inverse covariance matrix.
Revised: Prior work has developed black-box VI algorithms for other types of state space models with a Gaussian variational posterior that is parameterized by a tridiagonal inverse covariance matrix (Archer et al., 2015; Bamler & Mandt, 2017).

Original: It can be shown that our sampling procedure also leads to a block tridiagonal inverse covariance matrix, although we never directly parameterize this matrix.
Revised: It can be shown that our Gauss-Markov posterior also has a block tridiagonal inverse covariance, but we never explicitly construct this matrix.

Original: Unlike our method, it uses amortized inference where the mean and inverse covariance of the Gaussian are parameterized by the output of a neural network with observations y as its input.

Original: Furthermore, it only allows Gaussian posteriors, while our approach induces implicit non-Gaussian  posteriors.
Revised: While we use tranformations to implicitly allow non-Gaussian posteriors, prior work assumed Gaussian posteriors.

Original: Most importantly, the variational approximation of Archer et al. 
Revised: Archer et al. (2015) use amortized inference where the mean and inverse covariance of the

Revised: Gaussian are determined by the output of a neural network with observations y as its input.

Original:  is  parameterized in a way that still requires matrix inversion.
Revised: The variational approximation is thus parameterized in a way that still requires matrix inversion.

Original: Although matrix inversion can be done more efficiently in sparse block tridiagonal matrices, this is still more expensive than our parameterization that samples via a single forward recursion.
Revised: Although matrix inversion can be done more efficiently in sparse block-tridiagonal matrices, this is still more expensive than our parameterization that samples via a single forward recursion.

Revised: To solve the linear system induced by the sparse inverse-covariance, Bamler & Mandt (2017) apply a forward and backward pass of temporal messagepassing, while we only require a single forward recursion.

Original: Ryder et al.
Revised: Ryder et al.

Original: (2018) proposed an algorithm for black-box variational inference for SDE that similarly works by discretizing the time and reframing the problem as state space models, then jointly optimizing for parameters and latent states.
Revised: (2018) proposed an algorithm for black-box variational inference for SDE that similarly works by discretizing the time and reframing the problem as state space models, then jointly optimizing for parameters and latent states.

Original: The authors use mean-field approximation on the model parameters θ and applies amortized inference for the latent states x .
Revised: The authors use mean-field approximation on the model parameters θ and applies amortized inference for the latent states x .

Original: The variational posterior of the latent states are assumed to be Gaussian with mean and variance parameterized by the output of recurrent neural networks.
Revised: The variational posterior of the latent states are assumed to be Gaussian with mean and variance parameterized by the output of recurrent neural networks.

Original: A more recent follow-up work (Ryder et al., 2021) uses normalizing flows as the variational approximation to allow non-Gaussian posterior, and replaces recurrent neural networks with convolutional networks for faster inference.[Figure]
Revised: A more recent follow-up work (Ryder et al., 2021) uses normalizing flows as the variational approximation to allow non-Gaussian posterior, and replaces recurrent neural networks with convolutional networks for faster inference.

Original: 5. Experiments 5.1. Simulated Data
Revised: 5. Experiments 5.1. Simulated Data

Original: We consider two variants of the SCON model: SCON-C and
Revised: We consider two variants of the SCON model: SCON-C and

Original: SCON-SS.
Revised: SCON-SS.

Original: We define informative priors on model parameters p ( θ ) and the initial condition p ( x 0 ) based on expert knowledge, and simulate data from the joint distribution p ( θ, x, y )=
Revised: We define informative priors on model param- [Figure] eters p ( θ ) and the initial condition p ( x 0 ) based on expert knowledge, and simulate data from the joint distribution p ( θ, x, y )

Original:  p ( θ ) p ( x | θ ) p ( y | x, θ ) with T = 1000 time steps.
Revised: = p ( θ ) p ( x | θ ) p ( y | x, θ ) with T = 1000 time steps.

Original: We sample noisy observations of S , D , M , and CO 2 everytime steps.
Revised: We sample noisy observations of S , D , M , and CO 2 everytime steps.

Original: We use regular gap between observations in our simulated data for simplicity, but this is not a requirement for our method.
Revised: We use regular gap between observations in our simulated data for simplicity, but this is not a requirement for our method.

Original: When the system is linear with Gaussian noise that does not depend on the state variables, the Kalman smoother (Kalman, 1960) can exactly recover the true posterior p ( x | y ; θ ) .
Revised: When the system is linear with Gaussian noise that does not depend on the state variables, the Kalman smoother (Kalman, 1960) can exactly recover the true posterior p ( x | y ; θ ) .

Original: We first use the constant diffusion model SCONC to verify that our proposed method can recover the optimal
Revised: We first use the constant diffusion model SCONC to verify that our proposed method can recover the optimal

Original: Kalman smoother solution.
Revised: Kalman smoother solution.

Original: We also compare Gauss-Markov
Revised: We also compare Gauss-Markov

Original: ADVI (GM-ADVI) to three alternative variational approximations for the latent states: full-rank ADVI (FR-ADVI), mean-field ADVI (MF-ADVI), and an amortized inference method based on normalizing flows (Ryder et al., 2021).
Revised: ADVI (GM-ADVI) to three alternative variational approximations for the latent states: full-rank ADVI (FR-ADVI), mean-field ADVI (MF-ADVI), and an amortized inference method based on normalizing flows (Ryder et al., 2021).

Original: For the Kalman smoother, we fix θ using the true values used to generate the data.
Revised: For the Kalman smoother, we fix θ using the true values used to generate the data.

Original: For the rest of the methods, θ and x are inferred jointly.
Revised: For the rest of the methods, θ and x are inferred jointly.

Original: For all these methods where θ is inferred, we use the full-rank variational approximation on θ with rescaled sigmoid transformation to enforce the (fixed) lower and upper bounds.
Revised: For all these methods where θ is inferred, we use the full-rank variational approximation on θ with rescaled sigmoid transformation to enforce the (fixed) lower and upper bounds.

Original: For SCON-SS, we drop the comparison with FR-ADVI and
Revised: For SCON-SS, we drop the comparison with FR-ADVI and

Original: MF-ADVI since they perform poorly on the simpler model, and focus on the comparison between GM-ADVI and the normalizing flow.
Revised: MF-ADVI since they perform poorly on the simpler model, and focus on the comparison between GM-ADVI and the normalizing flow.

Original: Note that although the SCON-SS model is linear, Kalman smoother cannot be applied here, since the state-scaling diffusion noise depends on state variables.[Figure] [Figure] [Figure]
Revised: Note that although the SCON-SS model is linear, Kalman smoother cannot be applied here, since the state-scaling diffusion noise depends on state variables.

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: hour

Original: [Figure] hour [Figure] In all experiments, we fix a budget of 50 samples and 50, iterations for all methods.
Revised: In all experiments, we fix a budget of 50 samples and 50, iterations for all methods.

Original: We use AdaGrad with learning rate 0.1 for ADVI and Adam with learning rate 0.01 for the normalizing flows.
Revised: We use AdaGrad with learning rate 0.1 for ADVI and Adam with learning rate 0.01 for the normalizing flows.

Original: GM-ADVI converges more quickly to better variational bounds.
Revised: GM-ADVI converges more quickly to better variational Gauss-Markov ADVI Full-rank ADVI Mean-field ADVI

Revised: Normalizing Flow

Revised: Kalman Smoother

Original: Figure 2 compares the negative ELBO losses of different methods across inference iterations.
Revised: S 1. S K a l m a n S 2.

Revised: D 0.

Revised: D

Revised: K a l m a n

Revised: D      hour 4.

Revised: M

Revised: [Figure]

Revised: estimates      hour 0.

Revised: M

Revised: [Figure]

Revised: bounds.

Revised: Figure 2 compares the negative ELBO losses of different methods across inference iterations.

Original: The ELBOs are estimated via Monte Carlo sampling from 100 samples every 100 iterations.
Revised: The ELBOs are estimated via Monte Carlo sampling from 100 samples every 100 iterations.

Original: In figure 2(a), both GM-ADVI and the normalizing flow converge to much better (lower) loss compared to FR-ADVI and MF-ADVI.
Revised: In figure 2(a), both GM-ADVI and the normalizing flow converge to much better (lower) loss compared to FR-ADVI and MF-ADVI.

Original: MF-ADVI converges to suboptimal loss due to false temporal independence assumption.
Revised: MF-ADVI converges to suboptimal loss due to false temporal independence assumption.

Original: Although potentially highly expressive, FR-ADVI fails to converge after 50,000 iterations, likely due to having too many irrelevant parameters that make it susceptible to noise and local optima.
Revised: Although potentially highly expressive, FR-ADVI fails to converge after 50,000 iterations, likely due to having too many irrelevant parameters that make it susceptible to noise and local optima.

Original: Compared to the normalizing flow,
Revised: Compared to the normalizing flow,

Original: Gauss-Markov ADVI is much less noisy and still converges to better variational bounds, as highlighted in figure 2(b).
Revised: Gauss-Markov ADVI is much less noisy and still converges to better variational bounds, as highlighted in figure 2(b).

Original: GM-ADVI more accurately captures the variance of the latent state posterior.
Revised: GM-ADVI more accurately captures the variance of the latent state posterior.

Original: Figure 3 compares the variational posterior of the latent states q ( x ) recovered by the variational methods against the true posterior p ( x | y ; θ ) computed by the Kalman smoother on the SCON-C model.
Revised: Figures 3 and 4 compare the variational posterior of the latent states q ( x ) recovered by the variational methods against the true posterior p ( x | y ; θ ) computed by the Kalman smoother on the SCON-C and

Original: For the
Revised: SCON-SS models.

Original:  Kalman smoother, the figure shows analytical means and standard deviations.
Revised: For Kalman smoother, the figure shows analytical means and standard deviations.

Original: For the variational methods, we use empirical estimates from 1,000 samples drawn from the variational posterior q ( x ) .
Revised: For the variational methods, we use empirical estimates from 1,000 samples drawn from the variational posterior q ( x ) .

Original: Except for FR-ADVI which is highly noisy and  overfits the observations, all other methods  perform equally well at estimating the posterior mean in figure 3(a).
Revised: Except for FRADVI which is highly noisy and clearly overfits the observations, all other methods may seem to perform reasonably well at estimating the posterior mean in figure 3(a).

Original: However, Gauss-Markov ADVI is the only approach that captures the appropriate amount of uncertainty in the posterior distribution, matching the Kalman 
Revised: However, figure 3(b) highlights that GM-ADVI mean estimates are still the closest to the true Kalman smoother estimates; other methods are biased by noisy observations and insufficiently smooth. More importantly, GM-ADVI is the only approach that captures the appropriate amount of uncertainty in the posterior distribution, matching the Kalman smoother variance on all three states, as shown in figure 4(a). The full-rank approximation systematically overestimates [Figure]

Original: Gauss-Markov ADVI
Revised: Mean-field ADVI

Original: Normalizing Flow  0.
Revised: Normalizing Flow Kalman Smoother 0.

Original: S 0.
Revised: S 0.

Revised: S 0.

Revised: D 0.

Original: D      hour 0.
Revised: D      hour 0.

Original: M smoother variance on all three states, as shown in figure 3(b).
Revised: [Figure] hour [Figure] [Figure]

Revised: [Figure]

Original: The full-rank approximation systematically overestimates variance, while the mean-field and normalizing flow approximations systematically underestimate variance.
Revised: [Figure] [Figure] [Figure] variance, while the mean-field and normalizing flow approximations systematically underestimate variance.

Original: Figurehighlights the difference in the posterior standard deviation between GM-ADVI and the normalizing flow on the SCONSS model.
Revised: Figure 4(b) shows the posterior standard deviation of GM-ADVI and the normalizing flow on the SCON-SS model.

Original: We do not show the posterior mean comparison since both methods perform similarly.

Original: GM-ADVI estimates the posterior of diffusion noise parameters more accurately.
Revised: GM-ADVI estimates the posterior of diffusion noise parameters more accurately.

Original: Figure 5 shows the marginal distributions of the variational posterior on select SCON-C model parameters θ .
Revised: Figure 5 shows the marginal distributions of the variational posterior on select SCON-C model parameters θ .

Original: For all methods, we use the full-rank variational approximation on θ with rescaled sigmoid transformation to enforce the lower and upper bounds.
Revised: For all methods, we use the full-rank variational approximation on θ with rescaled sigmoid transformation to enforce the lower and upper bounds.

Original: We expect posteriors to shift away from the priors (blue) toward the true θ value (gray vertical line), but still display uncertainty.
Revised: We expect posteriors to shift away from the priors (blue) toward the true θ value (gray vertical line), but still display uncertainty.

Original: For all drift parameters (the first three rows), all methods seem to converge to similar posterior distributions.
Revised: For all drift parameters (the first three rows), all methods seem to converge to similar posterior distributions.

Original: All methods other than GM-ADVI seem to overestimate diffusion parameters ( c S , c D , c M ) compared to their true values.Fig-
Revised: All methods other than GM-ADVI seem to overestimate diffusion parameters ( c S , c D , c M ) compared to their true values.

Original: [Figure] ure 6 shows the marginal distributions of the variational posterior on SCON-SS diffusion parameters ( s S , s D , s M ).
Revised: Figure 6 shows the marginal distributions of the variational posterior on SCON-SS diffusion parameters ( s S , s D , s M ).

Original: We do not show the drift parameter comparison since both methods infer similar posterior distributions.
Revised: We do not show the drift parameter comparison since both methods infer similar posterior distributions.

Original: Here, we see more clearly that the normalizing flow still overestimates diffusion parameters compared to GM-ADVI.
Revised: Here, we see more clearly that the normalizing flow still overestimates diffusion parameters compared to GM-ADVI.

Original: GM-ADVI is more practical for scientific applications.
Revised: GM-ADVI is more practical for scientific applications.

Original: Compared to the normalizing flow, all ADVI methods share equal advantage of being simpler and more easily integrated to PPLs to enable automatic inference.
Revised: Compared to the normalizing flow, all ADVI methods share equal advantage of being simpler and more easily integrated to PPLs to enable automatic inference.

Original: They require fewer tuning hyperparameters, such as number of layers and choice of neural network architectures, among others.
Revised: They require fewer tuning hyperparameters, such as number of layers and choice of neural network architectures, among others.

Original: Although we show no direct runtime comparison in this paper, all ADVI experiments were run on the CPU, while the normalizing flow experiments required a GPU, which may not be easily accessible for many scientific applications.
Revised: Although we show no direct runtime comparison in this paper, all ADVI experiments were run on the CPU, while the normalizing flow experiments required a GPU, which may not be easily accessible for many scientific applications.

Original: Compared to other ADVI approximations, Gauss-Markov offers balance between expressivity and efficiency: it is expressive without the extraneous parameters, while maintaining linear time and space.[Figure] [Figure]
Revised: Compared to other ADVI approximations, Gauss-Markov offers balance between expressivity and efficiency: it is expressive without the extraneous parameters, while maintaining linear time and space.

Original: [Figure]
Revised: [Figure]

Original: 5.2. Meta-Analysis Data
Revised: 5.2. Meta-Analysis Data

Original: Finally, we apply the proposed Gauss-Markov ADVI method on meta-analysis data used in (Xie et al., 2020).
Revised: Finally [Figure] we apply the proposed Gauss-Markov ADVI method on meta-analysis data used in (Xie et al., 2020) [Figure]

Original: The dataset was compiled from 27 soil warming studies that measured CO 2 .
Revised: The dataset was compiled from 27 soil warming studies that measured CO 2 .

Original: The pooled data consist of annual COresponse ratios over a period of 13 years.
Revised: The pooled data consist of annual COresponse ratios over a period of 13 years.

Original: Each response ratio is calculated by dividing the annual CO 2 mean following warming perturbation by CO 2 measured at prewarming steady state.
Revised: Each response ratio is calculated by dividing the annual CO 2 mean following warming perturbation by CO 2 measured at prewarming steady state.

Original: We consider the stochastic variants of the two models explored by Xie et al.
Revised: We consider the stochastic variants of the two models explored by Xie et al.

Original: (2020): SCON-SS and
Revised: (2020): SCON-SS and

Original: SAWB-SS.
Revised: SAWB-SS.

Original: We focus on Gauss-Markov ADVI in this case study and drop comparison with the remaining variational methods since they perform poorly on simulated data.
Revised: We focus on Gauss-Markov ADVI in this case study and drop comparison with the remaining variational methods since they perform poorly on simulated data.

Original: We discretize the 13-year timespan into T = 220 discrete time steps.
Revised: We discretize the 13-year timespan into T = 220 discrete time steps.

Original: Following Xie et al.
Revised: Following Xie et al.

Original: (2020), annual response ratios are assumed to be “collected” at the halfway point of each year.
Revised: (2020), annual response ratios are assumed to be “collected” at the halfway point of each year.

Original: We modify our observation likelihood to use the CO 2 response ratios as observations (details in the supplement).
Revised: We modify our observation likelihood to use the CO 2 response ratios as observations (details in the supplement).

Original: As before, we infer latent states x and model parameters θ jointly.
Revised: As before, we infer latent states x and model parameters θ jointly.

Original: We run  20,000 inference iterations with 50 samples using AdaGrad with learning rate 0.1.
Revised: We run [Figure] [Figure] 20,000 inference iterations with 50 samples using AdaGrad with learning rate 0.1.

Original: We estimate the posterior of COresponse ratios by drawing 5,000 samples of x ( b ) ∼ q ( x ) and θ ( b ) ∼ q ( θ ) , and evaluate the CO 2 response ratio corresponding to each sample b as a function of x ( b ) and θ ( b ) .
Revised: We estimate the posterior of COresponse ratios by drawing 5,000 samples of x ( b ) ∼ q ( x ) and θ ( b ) ∼ q ( θ ) , and evaluate the CO 2 response ratio corresponding to each sample b as a function of x ( b ) and θ ( b ) .

Original: Figure 7 shows the variational posterior of CO 2 response ratios.
Revised: Figure 7 shows the variational posterior of CO 2 response ratios.

Original: With only a total of 13 data points, this meta-analysis dataset is substantially more sparse than the simulated data.
Revised: With only a total of 13 data points, this meta-analysis dataset is substantially more sparse than the simulated data.

Original: However, the original simulation (Xie et al., 2020) using Hamiltonian Monte Carlo (Hoffman et al., 2014) and deterministic
Revised: However, the original simulation (Xie et al., 2020) using Hamiltonian Monte Carlo (Hoffman et al., 2014) and deterministic

Original: ODE solver still took multiple weeks to run.
Revised: ODE solver still took multiple weeks to run.

Original: In comparison, our approach took less than an hour on the CPU.
Revised: In comparison, our approach took less than an hour on the CPU.

Original: 6. Conclusion
Revised: 6. Conclusion

Original: As more effort is being made in biogeochemistry and other scientific fields to collect denser observations over longer periods of study, the need for faster and more scalable inference algorithms also becomes more crucial.
Revised: As more effort is being made in biogeochemistry and other scientific fields to collect denser observations over longer periods of study, the need for faster and more scalable inference algorithms also becomes more crucial.

Original: For example, the ongoing Harvard Forest study (Melillo et al., 2017) that began in 1991 currently contains over 500 CO 2 observations over the span of 30 years and continues to grow.
Revised: For example, the ongoing Harvard Forest study (Melillo et al., 2017) that began in 1991 currently contains over 500 CO 2 observations over the span of 30 years and continues to grow.

Original: We have developed Gauss-Markov ADVI for sequential latent variable models that leverages the Markov structure of stochastic differential equations to arrive at an algorithm that is both expressive and scales linearly in the sequence length.
Revised: We have developed Gauss-Markov ADVI for sequential latent variable models that leverages the Markov structure of stochastic differential equations to arrive at an algorithm that is both expressive and scales linearly in the sequence length.

Original: We demonstrate its effective use on the time discretized soil biogeochemical models.
Revised: We demonstrate its effective use on the time discretized soil biogeochemical models.

Original: Compared to other variational methods, Gauss-Markov ADVI converges faster and more accurately captures uncertainty in the posterior.
Revised: Compared to other variational methods, Gauss-Markov ADVI converges faster and more accurately captures uncertainty in the posterior.

Original: It provides a compelling inference method for complex sequence models with interpretable latent variables and meaningful priors.
Revised: It provides a compelling inference method for complex sequence models with interpretable latent variables and meaningful priors.

Original: It is designed to enable easy integration with PPLs, and in future work we would like to implement our algorithm in
Revised: It is designed to enable easy integration with PPLs, and in future work we would like to implement our algorithm in

Original: PPLs such as Stan or Pyro, to allow for convenient use by the scientific community.
Revised: PPLs such as Stan or Pyro, to allow for convenient use by the scientific community.

Original: In the meantime, our Python code will be made available online.
Revised: In the meantime, our Python code will be made available online.

Revised: Acknowledgements

Revised: This research was supported in part by a UC Irvine ICS Exploration Research Award, and by NSF Robust Intelligence

Revised: Award No. IIS-1816365.

Original: A. Transition Likelihood
Revised: A. Transition Likelihood

Original: The transition likelihod is given by:
Revised: The transition likelihod is given by:

Original: [Equation]
Revised: [Equation]

Original: where TN ( µ, Σ , a, b ) denotes the truncated Gaussian density with mean µ , covariance Σ , lower bound a , and upper bound b .
Revised: where TN ( µ, Σ , a, b ) denotes the truncated Gaussian density with mean µ , covariance Σ , lower bound a , and upper bound b .

Original: The diffusion matrix in the constant diffusion model (SCON-C and SAWB-C) is β ( x t − 1 , θ )
Revised: The diffusion matrix in the constant diffusion model (SCON-C and SAWB-C) is β ( x t − 1 , θ )

Original: =
Revised: =

Original: Iσ .
Revised: Iσ .

Original: In the state-scaling model (SCON-SS and SAWB-SS), it is β ( x t − 1 , θ ) =
Revised: In the state-scaling model (SCON-SS and SAWB-SS), it is β ( x t − 1 , θ ) =

Original: I ( σ ⊙ √ x ) , where ⊙ denotes elementwise multiplication and σ ∈ θ is a model parameter.
Revised: I ( σ ⊙ √ x ) , where ⊙ denotes elementwise multiplication and σ ∈ θ is a model parameter.

Original: We define the drift function α below.
Revised: We define the drift function α below.

Original: A.1. SCON Model
Revised: A.1. SCON Model

Original: The drift vector in the SCON model obeys the following dynamics:
Revised: The drift vector in the SCON model obeys the following dynamics:

Original: [Equation]
Revised: [Equation]

Original: with latent states x t =
Revised: with latent states x t =

Original: ( S t , D t , M t ) and external input rate ( I S,t , I D,t ) .
Revised: ( S t , D t , M t ) and external input rate ( I S,t , I D,t ) .

Original: The model parameters consist of both drift and diffusion parameters θ = ( u M , a DS , a SD , a M , a MS , k S, ref , k D, ref , k M, ref , Ea S , Ea D , Ea M , σ S , σ D , σ M ) .
Revised: The model parameters consist of both drift and diffusion parameters θ = ( u M , a DS , a SD , a M , a MS , k S, ref , k D, ref , k M, ref , Ea S , Ea D , Ea M , σ S , σ D , σ M ) .

Original: The k linear first-order decay parameters obey Arrhenius temperature dependence such that:
Revised: The k linear first-order decay parameters obey Arrhenius temperature dependence such that:

Original: [Equation]
Revised: [Equation]

Original:  where R is the ideal gas constant 8.314 J · K − 1 · mol − 1 , T t is the temperature at time t , and T ref specifies a “reference” equilibrium temperature which we set at 283 K (for simulated data) or 283.15 K (for meta-analysis data).
Revised: [Equation] where R is the ideal gas constant 8.314 J · K − 1 · mol − 1 , T t is the temperature at time t , and T ref specifies a “reference” equilibrium temperature which we set at 283 K (for simulated data) or 283.15 K (for meta-analysis data).

Original: A.2. SAWB Model
Revised: A.2. SAWB Model

Original: The drift vector in the SAWB model obeys the following dynamics:
Revised: The drift vector in the SAWB model obeys the following dynamics:

Original: [Equation]
Revised: [Equation]

Revised: 

Original: [Equation]
Revised: [Equation]

Original: with latent states x t =
Revised: with latent states x t =

Original: ( S t , D t , M t , E t ) and external input rate ( I S,t , I D,t ) .
Revised: ( S t , D t , M t , E t ) and external input rate ( I S,t , I D,t ) .

Original: The model parameters consist of both drift and diffusion parameters θ = ( u Q,ref , Q, a MS , K D , K U , V D, ref , V U, ref , Ea V D , Ea V U , r M , r E , r L , σ S , σ D , σ M , σ E ) .
Revised: The model parameters consist of both drift and diffusion parameters θ = ( u Q,ref , Q, a MS , K D , K U , V D, ref , V U, ref , Ea V D , Ea V U , r M , r E , r L , σ S , σ D , σ M , σ E ) .

Original: The V reaction rate parameters are forced by temperature via Arrhenius temperature dependence such that:
Revised: The V reaction rate parameters are forced by temperature via Arrhenius temperature dependence such that:

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The u Q carbon use efficiency fraction follows linear temperature dependence such that:
Revised: The u Q carbon use efficiency fraction follows linear temperature dependence such that:

Original: [Equation]
Revised: [Equation]

Original: B. Observation Likelihood
Revised: B. Observation Likelihood

Original: The observation likelihood is given by:
Revised: The observation likelihood is given by:

Original: [Equation]
Revised: [Equation]

Original: In all experiments, we fix lower and upper bounds ( a y , b y ) .
Revised: In all experiments, we fix lower and upper bounds ( a y , b y ) .

Original: In experiments with simulated data, our observations are noisy measurements of the state variables and CO 2 .
Revised: In experiments with simulated data, our observations are noisy measurements of the state variables and CO 2 .

Original: We fix the observation covariance Σ y , and the observation mean is computed as a function of the state variables:
Revised: We fix the observation covariance Σ y , and the observation mean is computed as a function of the state variables:

Original: [Equation]
Revised: [Equation]

Original: In the meta-analysis data, we observe only the CO 2 response ratios.
Revised: In the meta-analysis data, we observe only the CO 2 response ratios.

Original: We assume prior Σ y ∼ HalfCauchy (1) for the observation covariance and infer its posterior.
Revised: We assume prior Σ y ∼ HalfCauchy (1) for the observation covariance and infer its posterior.

Original: The observation mean is:
Revised: The observation mean is:

Original: [Equation]
Revised: [Equation]

Original: 16) where ˆ r CO 2 = g (ˆ x, θ ) is the CO 2 computed at pre-warming steady state.
Revised: [Equation] where ˆ r CO 2 = g (ˆ x, θ ) is the CO 2 computed at pre-warming steady state.

Original: We define CO 2 and steady state solutions ˆ x in
Revised: We define CO 2 and steady state solutions ˆ x in

Original: SCON and SAWB models below.
Revised: SCON and SAWB models below.

Original: B.1. SCON Model
Revised: B.1. SCON Model

Original: The CO 2 at time t in SCON is given by the equation:
Revised: The CO 2 at time t in SCON is given by the equation:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The SCON respiration can be thought of as the sum of the fractions of carbon fluxes not transferred into other state variables at t .
Revised: The SCON respiration can be thought of as the sum of the fractions of carbon fluxes not transferred into other state variables at t .

Original: The steady state solutions for the state variables in SCON are:
Revised: The steady state solutions for the state variables in SCON are:

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Original: B.2. SAWB Model
Revised: B.2. SAWB Model

Original: The CO 2 at time t in SAWB is given by the equation:
Revised: The CO 2 at time t in SAWB is given by the equation:

Original: [Equation]
Revised: [Equation]

Original:  The SAWB respiration can be thought of as the fraction of carbon flux out of the DOC pool that is not absorbed by the microbial biomass MBC.
Revised: [Equation] The SAWB respiration can be thought of as the fraction of carbon flux out of the DOC pool that is not absorbed by the microbial biomass MBC [Equation]

Original: The steady state solutions for the state variables in SAWB are:
Revised: The steady state solutions for the state variables in SAWB are:

Original: [Equation]
Revised: [Equation]

Original: [Equation]