Probabilistic Transformer For Time Series Analysis

1 Introduction . Generative modeling of multivariate time series is a challenging problem with wide-ranging applications in demand forecasting [ 15 , 76 ] , autonomous driving [ 2 , 16 ] , robotics [ 29 , 67 ] , and health care [ 20 , 21 , 59 ] . Despite remarkable progress in recent years , models that predict high-dimensional future observations from a few past examples have remained intractable , partly due to the complex , non-deterministic temporal dynamics across long-distance time steps . Given a sequence of human poses , for example , such models must internally figure out the involved dynamics of various body components across space and time while maintaining the inherent uncertainty of multiple plausible futures , even though only one such future is observed . Among proposed probabilistic approaches , state space models ( SSMs ) provide a principled framework for learning and drawing inference from sequential inputs [ 27 , 66 ] . While autoregressive models feed its predictions back into the dynamics model without any compressed representation of data , SSMs model stochastic transitions between abstract states using latent variables , allowing for efficient state-to-state sampling without the need to render high-dimensional observations . Gaussian linear dynamical systems ( LDSs ) , one of the best known SSMs [ 92 ] , for example , postulate linear state transitions and enjoy exact inference via the celebrated Kalman filter algorithm . While early extensions of LDSs focus on linearization [ 46 ] and unscented transform [ 88 ] , recent work that marry state space models with deep neural networks offers much more flexibility to model complex dependencies across different time steps . Some approaches retain the Markovian dynamics of LDSs and only replace their linear observation models with feed-forward networks [ 23 , 31 , 47 , 71 ] , whereas others favor nonlinear state transitions and parametrize such dependencies via recurrent neural networks ( RNNs ) [ 22 , 23 , 30 , 39 , 51 , 75 ] . Despite differences , both Markovian transitions and RNNs are often not capable of capturing long-range dependencies in highly structured sequential inputs [ 36 , 100 ] , limiting the capacity of the corresponding SSMs . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this work , we propose to combine the complementary strengths of SSMs and transformer architectures [ 85 ] , a powerful mechanism for modeling long-term interactions that enjoys success across a variety of sequence modeling tasks [ 26 , 48 , 99 ] . In contrast to most SSMs , our models make extensive use of attention mechanism [ 5 , 85 ] between latent variables to model non-Markovian dynamics ( see Figure 1 ) . Compared to transformer-based methods , our models are probabilistic , nonautoregressive in a similar fashion to LDSs , and capable of generating diverse long-term forecasts with uncertainty estimates . Our main contributions are threefold . First , we propose novel SSMs based on transformer architectures for multivariate time series , which include generative models and inference procedures based on variational inference [ 49 , 74 ] . Second , we extend our models to include several layers of stochastic latent variables organized in a hierarchy for further expressiveness . Third , we conduct extensive experiments on time series forecasting and human motion prediction and demonstrate that our Probabilistic Transformer ( ProTran ) performs remarkably well compared to various state-of-theart baselines . 2 Preliminaries . 2.1 Variational State Space Models . Let { x ( i ) 1 : Ti } N i=1 consist of N univariate time series where x ( i ) 1 : Ti = ( x ( i ) 1 , x ( i ) 2 , · · ·x ( i ) Ti ) and x ( i ) t denotes the vaue of the i-th time series at time t. We consider the multivariate form x1 : T = ( x1 , x2 , . . . , xT ) where xt = ( x ( 1 ) t , . . .x ( N ) t ) ∈ RN . Conditioning on observed values up to time C , we aim to produce distributional forecasts into the future p ( xC+1 : T |x1 : C ) . For clarity , we refer to x1 : C and xC+1 : T as contexts and targets , respectively . We are interested in probabilistic models parametrized by θ of the form pθ ( x1 : T |x1 : C ) = ∫ pθ ( x1 : T | z1 : T ) pθ ( z1 : T |x1 : C ) dz1 : T ( 1 ) where z1 : T = ( z1 , z2 , . . . , zT ) denotes the corresponding sequence of latent variables , sometimes referred to as states . In other words , we assume a generative model that can be decomposed into a transition model pθ ( z1 : T |x1 : C ) between the latent variables conditioned on the contexts , and an emission model pθ ( x1 : T | z1 : T ) from the latent variables to observable outputs . In particular , we further impose several assumptions on both models : 1 pθ ( z1 : T |x1 : C ) = T∏ t=1 pθ ( zt | z1 : t−1 , x1 : C ) , pθ ( x1 : T | z1 : T ) = T∏ t=1 pθ ( xt | zt ) . ( 2 ) As demonstrated in Figure 1 ( b ) , the latent variable zt+1 depends not only on zt but also on all of its preceding latent variables , including zt−1 , in contrast to linear dynamical systems ( LDSs ) . In addition , the transition and emission models allow for non-linearity via neural network parametrizations . These assumptions aim to maximize model capacity for real-world applications with complex emissions or temporal dependencies . However , neither x1 : t−1 nor z1 : t−1 are included in the emission model p ( xt | z1 : T , x1 : C ) . Such assumptions are important , as it has been argued previously that a leakage of information from the latent space in autoregressive models can hinder long-term predictions [ 23 , 47 ] . While all ground truth observations are available during training , the entire sequence has to be generated sequentially at test time , making the dependencies on x1 : t−1 prone to accumulated errors over multiple time steps . By letting the latent variable zt capture all information needed to render xt , we also avoid the computational costs associated with repeatedly decoding and encoding xt in multi-step predictions . The inclusion of nonlinear state transitions and observation models necessarily requires approximate inference . We follow the stochastic variational inference framework [ 49 , 74 ] and assume that the variational posterior parametrized by φ can be decomposed auto-regressively as qφ ( z1 : T |x1 : T ) =∏ t qφ ( zt | z1 : t−1 , x1 : T ) , which leads to a lower bound on the log likelihood : log pθ ( x1 : T |x1 : C ) ≥ T∑ t=1 ( Eq [ log pθ ( xt|zt ) ] − KL ( qφ ( zt|z1 : t−1 , x1 : T ) ‖ pθ ( zt|z1 : t−1 , x1 : C ) ) ) , ( 3 ) where KL is the Kullback-Leibler divergence . For computational stability , we assume homoscedasticity and choose Laplace distribution with scale parameter β as a parametric form for pθ ( xt | zt ) , i.e . we optimize for L1 reconstruction loss with a cross-validated factor β for the KL term , following similar variational autoencoder ( VAE ) work [ 24 , 41 , 86 ] . Such an assumption does not necessarily limit the capacity of our models , as powerful stochastic transitions and flexible emission models can theoretically characterize arbitrary noise covariance [ 66 ] . Incorporating structured probabilistic outputs such as Gaussian copulas [ 75 ] or normalizing flows [ 23 ] can potentially further improve our model performance . 2.2 Transformer Architectures . Central to our models and other transformer-based approaches [ 48 , 85 ] is the notion of attention [ 5 ] , which allows the models to focus on important parts within a context . Multi-head attention , for example , maps a sequence of queries Q ∈ R ` q×d of length ` q to a sequence of outputs O = [ O1 , . . . , OH ] ∈ R ` q×d of the same size by attending over given ` k key-value pairs K ∈ R ` k×d , V ∈ R ` k×d : Oh = Attention ( Qh , Kh , Vh ) = Softmax ( QhK T h√ d ) Vh , ( 4 ) where Qh = QW Q h , Kh = KW K h , Vh = VW V h are projected queries , keys , and values corresponding to head h ∈ [ 1 , H ] with learning parameters WQh , WKh , WVh , respectively . In case Q = K = V , we refer to such an attention mechanism as self-attention . Given fully observed sequences of inputs , the mapping can be computed efficiently without any imposed sequential order often seen in recurrent neural networks [ 19 , 42 ] . More importantly , the direct connections between long-distance time steps are baked into the mechanism as information from previous time steps is easily accessible without being compressed into a fixed representation , easing optimization and learning of long-term dependencies [ 5 , 85 ] . 1For notational simplicity , we assume x0 = z1:0 = ∅ and p ( z |x0 ) = p ( z ) . Without recurrence , Transformer [ 85 ] encodes information about each time step t with pred efined sinusoidal positional embeddings Position ( t ) = [ pt ( 1 ) , . . . , pt ( d ) ] ∈ Rd where the i-th embedding is given by pt ( i ) = sin ( t · ci/d ) for even i and pt ( i ) = cos ( t · ci/d ) for odd i and c is some large constant . Empirical results show that such positional embeddings are also important to our models . 3 Probabilistic Transformer . In this section , we first present our single-layered model and subsequently its multi-layered extension for a hierarchy of stochastic latent variables . As alluded earlier , our model consists of a generative model and an inference model that share information and parameters extensively . 3.1 Single-Layered Probabilistic Transformer . Generative Model . Given some contexts x1 : C , we first apply a linear projection and combine it with a positional embedding to obtain h1 : C ∈ Rd , i.e . ht = LayerNorm ( MLP ( xt ) + Position ( t ) ) , ( 5 ) where LayerNorm and MLP denote layer normalizations [ 4 ] and multi-layer perceptrons , respectively . While a traditional transformer model often dedicates an entire encoder for the same purpose [ 55 , 72 ] , we find such a simple mapping works sufficiently well in conjunction with the contextattention module of the corresponding decoder . As implied in Equation ( 2 ) , our latent dynamics decomposes auto-regressively . At each time step , we parametri ze the distribution pθ ( zt | z1 : t−1 , x1 : C ) by a Gaussian with parameters resulting from two sequential steps of attention : a self-attention over the previously inferred states z1 : t−1 and another attention over the projected contexts h1 : C . These two operations mirror those found in the decoder of Transformer [ 85 ] , with the stochastic latent variables replacing its decoder inputs . Unfortunately , using stochastic samples of zt as attention queries is problematic , as purely stochastic transitions make it difficult for the model to reliably retain information across multiple time steps [ 17 , 30 , 39 ] . We therefore encapsulate the latent variables in hidden representations wt that also has a deterministic component . Combined with the attention steps , such representations help model long-range temporal dependencies while accounting for the stochasticity of future observations . Starting with a learnable , context-agnostic representation w0 , we recursively update wt using a stochastic sample from pθ ( zt | z1 : t−1 , x1 : C ) and th e positional embedding for the current time step t. The generating process for the time step t can be summarized by the following pseudocode : w̄t = LayerNorm ( wt−1 + Attention ( wt−1 , w1 : t−1 , w1 : t−1 ) ) ( 6 ) ŵt = LayerNorm ( w̄t + Attention ( w̄t , h1 : C , h1 : C ) ) ( 7 ) zt = Sample ( N ( zt ; MLP ( ŵt ) , Softplus ( MLP ( ŵt ) ) ) ) ( 8 ) wt = LayerNorm ( ŵt + MLP ( zt ) + Position ( t ) ) , ( 9 ) where Sample and Softplus are the Gaussian sampling and approximating rectifier operators . Each stochastic sample of w1 : T is then mapped to a sequence of x1 : T via a multi-layer perceptron . We emphasize that our generation procedure in the latent space is more efficient than others in the observation space , which requires encoding and decoding high-dimensional inputs repeatedly . Inference Model . We parametrize the approximate posterior qφ ( zt | z1 : t−1 , x1 : T ) at time step t in a simi lar fashion to the prior pθ ( zt | z1 : t−1 , x1 : C ) . Indeed , these parametrizations share most parameters and are done simultaneously in the same recursive loop , following the exact same steps in Equation ( 6 ) and Equation ( 7 ) ( see Figure 1 ) . We note that similar sharing techniques between the generative and inference processes have emerged as a common theme among recent successful VAE models [ 17 , 62 , 83 ] . While the prior only has access to the conditioning observations x1 : C , the approximate posterior should take into account all observations during training , including the targets xC+1 : T . Due to the inherent unidirectional aspect of RNNs , previous work that uses RNNs to parametrize the approximate posterior often disregards such a property [ 22 , 30 , 51 ] and often resorts to a filtering routine p ( zt | z1 : t−1 , x1 : t ) . In contrast , our inference procedure resembles more of the smoothing process of LDSs , factoring in both past and future observations via another application of self-attention : kt = Attention ( h1 : T , h1 : T , h1 : T ) ) ( 10 ) zt = Sample ( N ( zt ; MLP ( [ ŵt , kt ] ) , Softplus ( MLP ( [ ŵt , kt ] ) ) ) . ( 11 ) Here , we replace Equation ( 8 ) in the generative model with Equation ( 11 ) , where the hidden representation kt summarizing all information relevant to the current tim estep t has been concatenate to the latent-and-context-aware representation ŵt preceding the Gaussian parametrization . The generative model and the inference model are trained end-to-end with a single stochastic variational inference objective stated in Equation ( 3 ) . Such a variational bound includes the reconstruction loss for x1 : C and the KL term for z1 : C . Alternatively , we can exclude these terms from the objective , which is equivalent to starting the inference process from t = C + 1 instead of t = 1 . Our models incur a time complexity of O ( T 2d ) and a memory cost of O ( T 2d ) , where T is the total sequence length and d is the dimensionality of the latent space . The recursive latent dynamics also does not allow use the take full advantange of parallelizable attentions . However , we find that our models are still efficient in practice , especially for reasonably small values of T .

Summary: The submission presents a probabilistic extension to the Transformer by augmenting the intermediate Transformer representations via stochastic latent variables. The attention mechanism allows for accessing the past and the future steps directly in contrast to the indirect hidden state of RNN-based state-space models (SSMs). The building blocks of the Transformer remain untouched. The attention operation is applied to the deterministic representations as in the Transformer. The random variables are then calculated by using these updated deterministic representations. The authors also propose a hierarchical variant of their model to enable stacking several layers as commonly observed in the Transformer architectures.