Flow-based & Recurrent MDNs (FRMDN)

• Flow-based and Recurrent MDNs (FRMDNs) are sequence models that combine recurrent mixture density networks with normalizing flows to capture complex, multimodal data distributions.
• They parameterize predictive densities by transforming outputs via invertible affine-coupling layers, enabling efficient computation of Jacobian determinants.
• Empirical results demonstrate that FRMDNs outperform traditional RMDNs in image and audio tasks by providing state-of-the-art sequence likelihoods.

Flow-based and Recurrent Mixture Density Networks (FRMDN) represent a class of sequence models that augment classical recurrent mixture density networks (RMDNs) by integrating invertible transformations via normalizing flows into their probabilistic output heads. Through this hybridization, FRMDNs achieve highly expressive, tractable sequence likelihoods and offer state-of-the-art performance on diverse generative modeling benchmarks, notably in vision and audio settings. Closely related architectures, such as Recurrent Flow Networks (RFN), further expand this paradigm by disentangling spatial and temporal uncertainty through the joint parameterization of RNN hidden states, stochastic latent variables, and flow-based emission distributions.

1. Foundations: Recurrent Mixture Density Networks and Flow-based Extensions

The core of an RMDN consists of an RNN (typically an LSTM) that processes a sequence of inputs $x_{\leq t}$ and previous targets $y_{\leq t}$ to produce hidden state vectors $h_t$ at each time step. From this, parameters of a Gaussian mixture model (GMM)—mixture weights $\pi_{t,k}$, means $\mu_{t,k}$, and covariance structures $\Sigma_{t,k}$—are generated, directly modeling the predictive density $p(y_{t+1} | h_t)$ at each step. While this architecture is expressive, its ability to capture complex, nonlinear, and multimodal targets remains fundamentally limited by the parametric form of the GMM.

FRMDNs address this expressivity bottleneck by modeling the density not directly on $y_{t+1}$ but rather on a nonlinearly transformed version $z_{t+1} = f(y_{t+1})$, where $f$ is a composition of invertible affine-coupling transformations, i.e., a normalizing flow of RealNVP type. The output density under the change of variables is

$$p(y_{t+1} | h_t) = p(z_{t+1} | h_t) \cdot \prod_{n=1}^N \left| \det \frac{\partial f_n^{-1}}{\partial z_{n-1}} \right|,$$

where $p(z_{t+1}|h_t)$ is the standard RNN-GMM density and the Jacobian factors are efficiently computable due to the flow’s triangular structure (Razavi et al., 2020).

2. Model Architecture: Parameterization and Flow Construction

At each time step, the RNN yields a fixed-dimensional vector $h_t$. A final linear projection is split into four blocks to parameterize the GMM:

• $o_t^{\pi} \in \mathbb{R}^K$: Mixture logits, mapped to weights $\pi_{t,k}$ via softmax.
• $o_t^{\mu} \in \mathbb{R}^{K \times d}$: Component means $\mu_{t,k}$.
• $o_t^D \in \mathbb{R}^{K \times d}$: Diagonal precision factors $D_{t,k}$, exponentiated for positive definiteness.
• $o_t^U \in \mathbb{R}^{K \times d \times d'}$: Low-rank factors $U_{t,k}$.

Covariance matrices are parameterized as $$\Sigma_{t,k}^{-1} = D_{t,k} + U_{t,k} U_{t,k}^{\top}$$, with $d' \ll d$ in practice, balancing expressivity and computational tractability.

The normalizing flow $f$ is composed of $N$ affine-coupling layers, each splitting the $d$-dimensional variable and applying invertible transformations:

• For each block $n$:

$$z_{n,1:d'} = z_{n-1,1:d'}, \quad z_{n,(d'+1):d} = z_{n-1,(d'+1):d} \odot \exp[s_n(z_{n-1,1:d'})] + t_n(z_{n-1,1:d'})$$

where $s_n(\cdot)$ and $t_n(\cdot)$ are neural networks and the Jacobian is lower triangular, yielding efficient determinant computation.

3. Training, Inference, and Implementation Details

FRMDNs are trained via direct maximization of the observed-sequence log-likelihood:

$$L = -\sum_{q=1}^Q \sum_{t=0}^{T_q-1} \left[ \log\sum_{k=1}^K \pi_{t,k} \mathcal{N}(f(y_{t+1}^q); \mu_{t,k}, \Sigma_{t,k}) + \sum_{n=1}^N \log \left| \det \frac{\partial f_n^{-1}}{\partial z_{n-1}} \right| \right],$$

with gradients efficiently propagated through both the RNN-GMM and flow layers due to the closed-form Jacobian. Optimization utilizes Adam or RMSProp with learning-rate scheduling, and covariance entries are clipped for stability during image-sequence modeling (Razavi et al., 2020).

Inference proceeds by:

1. Computing mixture parameters from $h_t$.
2. Sampling $z_{t+1}$ from the Gaussian mixture.
3. Applying the inverse flow $f^{-1}$ to obtain $y_{t+1}$.

Implemented RNNs utilize a single-layer LSTM with 256 hidden units (for image tasks), or several thousand units (in speech). Flows employ $N=2$ affine coupling layers with $s,t$ networks set as multilayer perceptrons utilizing LeakyReLU or Tanh activations.

4. Generalizations: Stochastic Latent States and Alternative Flow Parameterizations

Recurrent Flow Networks (RFN) (Gammelli et al., 2020) represent a complementary direction by further introducing stochastic latent variables $z_t$ alongside deterministic RNN states in the recurrent backbone. The generative process factorizes over time as:

$$p(x_{1:T}, z_{1:T}, h_{1:T}) = \prod_{t=1}^T p(h_t | h_{t-1}, u_t) \, p(z_t | z_{t-1},h_t) \, p(x_t | z_t, h_t).$$

The emission $p(x_t | z_t,h_t)$ leverages a conditional normalizing flow whose base distribution's parameters and coupling layers are functions of both $z_t$ and $h_t$. Variational inference is enabled with an amortized posterior $q_\phi(z_t | z_{t-1}, h_t, x_t)$ within a filtering-style ELBO objective.

This architecture allows explicit modeling of both temporal uncertainty (via latent $z_t$) and spatial complexity in emissions (via flow-based outputs), permitting separation and joint control of sequence dynamics and multimodal spatial distributions, as empirically validated in spatiotemporal mobility data (Gammelli et al., 2020).

5. Empirical Results and Comparative Performance

FRMDNs have demonstrated consistent and significant improvements over baseline RMDNs and other flow-based generative models across multiple domains:

• Image-sequence modeling (CarRacing, SuperMario latent spaces, $d=32$): For $K=5$ and low-rank $d'=3$, FRMDN achieves best reported test NLLs (CarRacing: 2.35 nats/time-step, SuperMario: 1.28 nats/time-step), outperforming RMDN variants (improvements up to 0.05 nats/time-step at low $K$).
• Speech generation (Blizzard, TIMIT, Accent datasets): With $K=20$ mixture components and diagonal precisions, FRMDN surpasses RMDN by thousands of nats (e.g., Blizzard: FRMDN $-7,961$ vs. RMDN $-4,272$).
• Single-image density (MNIST, CIFAR-10): On MNIST, FRMDN achieves test NLL below zero ($-89.9$ nats), outperforming masked autoregressive flows (MAF $1,092.3$ nats) and RMDN ($1,189.9$ nats); on CIFAR-10, FRMDN ($-3,159.8$ nats) exceeds MAF and RealNVP by 20–120 nats.

RFN demonstrates state-of-the-art log-likelihood on spatiotemporal urban mobility datasets such as NYC taxi and Copenhagen bike-share, outperforming conventional MDN-RNNs, VRNNs, and ConvLSTMs especially for long rollout horizons and in capturing intricate, non-Gaussian spatial features (Gammelli et al., 2020).

6. Theoretical and Practical Significance

The integration of normalizing flows within recurrent mixture density networks yields two principal benefits:

• Expressivity: By modeling the density on a flexible, learned nonlinear transformation of the targets, FRMDNs capture highly non-Gaussian, multimodal, and nonlinear structure beyond the direct GMM parameterization.
• Tractability: The use of affine-coupling flows (e.g., RealNVP) provides computationally efficient, analytically tractable likelihoods, enabling exact sequence-level log-likelihood evaluation and stable maximum likelihood training.

Low-rank plus diagonal covariance parameterization balances computational efficiency with the ability to model complex covariances, aiding scalability to high-dimensional targets, as in image sequences.

An implication is that FRMDN and related architectures facilitate both improved generative quality and reliable sequence likelihoods for downstream use in model-based reinforcement learning, audio waveform modeling, and empirical sequence density estimation.

7. Extensions, Limitations, and Future Directions

Extensions such as those in RFN (Gammelli et al., 2020) suggest opportunities to further separate sources of sequence variability by embedding stochastic latent states and leveraging conditional flows in the emission process. The consistent likelihood benefits reported for both vision and audio, as well as continuous spatial domains, indicate broad applicability.

Limitations include the computational cost associated with deep flows for large target dimensions, and the potential need for careful hyperparameter tuning (mixture count $K$, flow depth $N$, low-rank size $d'$). A plausible implication is the growing importance of hybrid architectures that combine the sample efficiency and uncertainty modeling of RNN-MDNs with the density estimation capacity of normalizing flows.

Overall, FRMDNs and variants constitute a robust and general framework for high-capacity, tractable sequential density modeling, effectively bridging autoregressive and flow-based generative paradigms (Razavi et al., 2020, Gammelli et al., 2020).