# Transformer Neural Autoregressive Flows

Massimiliano Patacchiola
University of Cambridge
<EMAIL_ADDRESS>
&Aliaksandra Shysheya
University of Cambridge
<EMAIL_ADDRESS>
Katja Hofmann
Microsoft Research
<EMAIL_ADDRESS>
&Richard E. Turner
University of Cambridge
<EMAIL_ADDRESS>

###### Abstract

Density estimation, a central problem in machine learning, can be performed
using Normalizing Flows (NFs). NFs comprise a sequence of invertible
transformations, that turn a complex target distribution into a simple one, by
exploiting the change of variables theorem. Neural Autoregressive Flows (NAFs)
and Block Neural Autoregressive Flows (B-NAFs) are arguably the most perfomant
members of the NF family. However, they suffer scalability issues and training
instability due to the constraints imposed on the network structure. In this
paper, we propose a novel solution to these challenges by exploiting
transformers to define a new class of neural flows called Transformer Neural
Autoregressive Flows (T-NAFs). T-NAFs treat each dimension of a random
variable as a separate input token, using attention masking to enforce an
autoregressive constraint. We take an amortization-inspired approach where the
transformer outputs the parameters of an invertible transformation. The
experimental results demonstrate that T-NAFs consistently match or outperform
NAFs and B-NAFs across multiple datasets from the UCI benchmark. Remarkably,
T-NAFs achieve these results using an order of magnitude fewer parameters than
previous approaches, without composing multiple flows.

##  1 Introduction

Normalizing Flows (NFs) have emerged as a powerful paradigm in probabilistic
modeling and machine learning, offering a versatile framework for mapping
between probability density functions through invertible transformations. NFs
are suitable for applications which require models of complex probability
distributions that support computationally efficient sampling and density
evaluation. Consequently, NPs are useful in density modelling problems and
also for variational inference in which they are used to specify approximate
posterior distributions.

Most of the research related to NFs focuses on designing expressive
transformations that conform to practical constraints. NFs specify an
bijective mapping from a simple base distribution to a target distribution so
that the transformation of densities formula can be used to specify the
density over the target variables. The transformation of densities formula
requires the computation of the determinant of the Jacobian of this mapping
and so NF models are structured so these objects are simple to compute. For
example, one notable class of NFs, known as Autoregressive Flows (AFs),
decompose a joint distribution into a product of univariate conditionals.
These transformations possess a lower triangular Jacobian, making it
computationally tractable to compute their determinants, which is an essential
requirement for the application of the change of variables theorem. Kingma et
al. (2016) introduced the concept of inverse autoregressive flows (IAFs),
refining AFs by employing a composition of trivially invertible affine
transformations to model each conditional, paving the way for more efficient
computation while retaining expressiveness.

Recent advances in the NF landscape include the introduction of Neural
Autoregressive Flows (NAFs, Huang et al. 2018) and Block Neural Autoregressive
Flows (B-NAFs, De Cao et al. 2020) NAFs and B-NAFs replace the IAF’s
transformation with a learned bijection implemented via a strictly monotonic
neural network. This fundamental paradigm shift makes NAFs and BNAFs more
flexible than IAFs, as they are universal approximators of real and continuous
distributions. However, despite these advancements, both NAFs and B-NAFs
encounter challenges in scalability and efficiency, primarily due to the
intricate structures and constraints imposed on their networks. The
scalability issue becomes particularly pronounced when dealing with high-
dimensional data, where the number of parameters in the network increases
significantly. Additionally, the training instability often observed in these
models can be attributed to the strict monotonicity and masking constraints,
which can limit the expressive power and learning dynamics of the networks.
These challenges necessitate a more flexible and scalable approach to enhance
the practical applicability of NFs in complex, real-world scenarios.

In response to these limitations, we introduce Transformer Neural
Autoregressive Flows (T-NAFs), a novel approach that leverages the power of
transformer architectures within the realm of NFs. T-NAFs treat each dimension
of a random variable as a separate input token and utilize attention masking
to enforce an autoregressive constraint. This design choice allows T-NAFs to
handle high-dimensional data more efficiently, as the transformer architecture
inherently scales better with dimensionality and does not require the same
level of parameterization as traditional feedforward networks. Moreover,
T-NAFs offer enhanced flexibility and stability in training by freeing the
majority of the parameters from the constraints typically imposed in NAFs and
B-NAFs. By incorporating transformers, T-NAFs not only address the scalability
and efficiency issues but also open up new avenues for innovation in the field
of NFs, particularly in complex density estimation tasks.

The main contributions of this paper can be summarized as follows:

  1. 1.

We introduce Transformer Neural Autoregressive Flows (T-NAFs), a novel model
that effectively combines transformers with neural autoregressive flows for
density estimation, utilizing attention masking to enforce autoregression
efficiently.

  2. 2.

We demonstrate remarkable flexibility in handling various invertible
transformations and showcase significant improvements in parameter efficiency
and scalability compared to previous models like NAFs and B-NAFs.

  3. 3.

Through extensive ablation studies and empirical evaluations across multiple
benchmarks, we provide key insights into the adaptability, complexity, and
effectiveness of the model.

##  2 Previous work

In this section we delve into the extensive body of research surrounding NFs,
a domain that has witnessed substantial developments in the last years. For a
comprehensive and detailed understanding of the progress in this field, we
refer the reader to recent review papers of Kobyzev et al. (2019) and
Papamakarios et al. (2021).

One broad line of research views NFs as a finite number of tractable simple
transformations and hence focuses on developing more expressive
transformations that have a tractable inverse and Jacobian determinant.
Rezende & Mohamed (2015) were the first to introduce two parametric families
of such transformations, i.e. the planar and the radial flows. Another well-
known representatives of this research direction are Autoregressive Flows
(Section 3.2), which are particularly relevant to our work. Dinh et al.
(2016); Kingma et al. (2016) proposed to build Autoregressive Flows with
simple affine transformations. To improve the limited expressivity of the
affine autoregressive transformations, several non-affine transformations were
proposed: Huang et al. (2018); De Cao et al. (2020) use a monotonic multi-
layered perceptron, while Jaini et al. (2019); Wehenkel & Louppe (2019) use an
integral of some positive function represented as neural network. Although
these transformations can be made arbitrary flexible, their major drawback is
that they do not have an analytic inverse. To overcome this issue, different
spline-based transformations were proposed in Müller et al. (2018); Durkan et
al. (2019); Dolatabadi et al. (2020).

Another important design choice for Autoregressive Flows is how to implement
the conditioning on observed variables. Even though there only constraint
imposed on conditioning function is the autoregressive one, in practice it is
still important to choose a parameter-efficient function that could
accommodate inputs of variable size. For instance, Oliva et al. (2018); Kingma
et al. (2016) use recurrent neural network (RNN) as a conditioning function.
Even though parameter-efficient, RNNs require sequential computation, which
can be prohibitively expensive for longer sequences. As an alternative to
RNNs, the idea of masking can be used: usually a standard feedforward network
is trained, but some of the connections are zeroed to preserve the
autoregressive structure of the conditioning function. Germain et al. (2015)
proposed a masking procedure for fully-connected neural networks. Masking was
picked up in many autoregressive flows models, including Kingma et al. (2016);
De Cao et al. (2020); Huang et al. (2018); Papamakarios et al. (2017).

Transformers have been used for distribution modeling in Fakoor et al. (2020).
The method exploits a transformer encoder that takes the input sequence
$x_{1},x_{2},\dots,x_{i}$ and predicts the $i$-th conditional distribution
$p(x_{i}|x_{1},x_{2},\dots,x_{i-1})$. The output is a categorical distribution
(discrete variables, or a mixture of Gaussians (continous variables). The main
issue with this approach is that it is not possible to get an exact estimate
of the likelihood of a sample, as the transformations are not bijective.
Moreover, if the input sequence is long the output of the model will be
disproportionally large, since each input is mapped to multiple outputs (e.g.
categories or parameters of the mixture).

##  3 Background

###  3.1 Normalizing Flows

A Normalizing Flow (NF) is an invertible (bijective) funciton
$f:\mathcal{X}\rightarrow\mathcal{Y}$ between two continuous random variables.
Since $f$ is invertible, it is possible to use the change of variable theorem
to translate between the two densities $p_{X}(\mathbf{x})$ and
$p_{Y}(\mathbf{y})$

$p_{Y}(\mathbf{y})=p_{X}(\mathbf{x})\ \big{\lvert}\text{det}\
\mathbf{J}_{f(\mathbf{x})}\big{\rvert}^{-1}\quad\text{where}\quad\mathbf{J}_{f(\mathbf{x})}=\frac{\partial\
f(\mathbf{x})}{\partial\ \mathbf{x}}.$ (1)

The determinant of the Jacobian $\mathbf{J}_{f(\mathbf{x})}$ accounts for
local expansions/contractions of $\mathcal{X}$ at $\mathbf{x}$. Estimating
this determinant is one of the major bottlenecks to overcome, since typically
this has complexity $\mathcal{O}(N^{3})$ for a Jacobian matrix of size
$N\times N$. One way to deal with this issue is to enforce a lower-triangular
structure on the Jacobian. The determinant of a lower-triangular matrix can be
estimated by multiplying the elements on the main diagonal, which has cost
$\mathcal{O}(N)$. The lower-triangular structure of the Jacobian can be
enforced through an autoregressive factorization, as detailed in the next
subsection.

###  3.2 Autoregressive Flows

For a multivariate random variable $\mathbf{x}\in\mathbb{R}^{K}$ sampled from
$p_{X}(\mathbf{x})$ we can estimate the join distribution with an
autoregressive factorization by exploiting the chain rule

$p_{X}(\mathbf{x})=p_{X_{1}}(x_{1})p_{X_{2}}(x_{2}|x_{1})\dots
p_{X_{K}}(x_{K}|x_{K-1},\dots,x_{1}).$ (2)

Exploiting this structure it is possible to estimate an arbitrary $y_{i}$ by
defining an autoregressive conditioner $c$ and an invertible transformation
$t$, such that

$y_{i}=f(x_{1},\dots,x_{i})=t(c(x_{1},\dots,x_{i-1}),x_{i}).$ (3)

Previous work, has investigated different ways of parametrizing $c$ and $t$. A
line of work focused on defining simple affine transformations such as
$t(\mu,\sigma,x_{i})=\mu+\sigma x_{i}$ in RealNVP (Dinh et al., 2016) or
$t(\mu,\sigma,x_{i})=\sigma x_{i}+(1-\sigma)\mu$ in IAF (Kingma et al., 2016).
While those transformations are trivial to invert, they suffer of scarce
flexibility which makes them inadequate to model complex distributions. In
order to overcome this issue another line of work has proposed to use neural
networks to represent the invertible transformation and/or the conditioner. In
the next subsection we will describe these particular models more in detail.

###  3.3 Neural Autoregressive Flows

Neural Autoregressive Flows In Neural Autoregressive Flows (NAFs, Huang et al.
2018) Equation 3 is reframed as follows

$t(c(x_{1},\dots,x_{i-1}),x_{i})=\text{NN}\big{(}x_{i};\
\bm{\psi}_{i}=c(x_{1},\dots,x_{i-1};\ \bm{\theta}_{i})\big{)}$ (4)

where NN is a feed-forward neural network and $\bm{\psi}_{i}$ are pseudo-
parameters generated by a conditioner network (parametrized by $\bm{\theta}$),
those pseudo-parameters represent the weights and biases of the neural
network. The conditioner network is implemented using a masked feed-forward
model as in MADE (Germain et al., 2015), while the invertible transformation
is implemented as a deep sigmoidal flow (DSF) or a deep dense sigmoidal flow
(DDSF). Both DSF and DDSF exploit sigmoidal inflection points in a fully-
connected model to induce multi-modality. A major limitation of NAFs is the
use of MADE in the conditioner. With MADE the number of parameters in the
network grows with the dimensionality of the random variable, which makes the
computational cost prohibitive expensive. Moreover, MADE exploits a zero-
masking strategy to impose the autoregressive constraint and to parallelize
the computation in the forward pass, which results in a large portion of
unused weights that affect the parameter count without having any concrete
utility.

Block Neural Autoregressive Flows In order to overcome some of the limitations
of NAFs, De Cao et al. (2020) proposed Block Neural Autoregressive Flows
(B-NAFs). The main idea of B-NAF is to directly parametrize the transformation
$t$ without the conditioner $c$ by using a single neural network

$t(c(x_{1},\dots,x_{i-1}),x_{i})=\text{NN}\big{(}x_{1},\dots,x_{i-1};\
\bm{\theta}_{i}\big{)}.$ (5)

The neural network is a feed-forward model, carefully designed to be both
autoregressive and strictly monotone (bijective). The autoregressive constrain
is imposed by using masking as in MADE (Germain et al., 2015) and the
monotonic constraint is imposed by using strictly positive weights and
invertible activation functions. While B-NAFs are more effective than NAFs on
a variety of datasets (De Cao et al., 2020) they still have two major
limitations. The first limitation is that, similarly to NAFs, B-NAFs rely on
masked feed-forward networks (MADE) that scale poorly with the dimensionality
of the random variable and results in unused parameters due to zero-masking.
The second limitation is that enforcing monotonicity in a large model leads to
training instability.

In the next section we describe how it is possible to overcome the limitations
of NAFs and B-NAFs by exploiting transformers as conditioning mechanisms,
presenting a new type of neural autoregressive flow.

##  4 Transformer Neural Autoregressive Flows

In order to overcome the limitations of NAF and B-NAF we propose a new type of
flow called Transformer Neural Autoregressive Flow (T-NAF). The main idea of
T-NAF is to allocate most of the computational capacity to the conditioner,
using a transformer neural network (Vaswani et al., 2017). This design choice
has a series of advantages: (i) we can define an autoregressive constrain by
simply using an autoregressive mask in the attention layer, contrary to MADE
this does not result in unused weights; (ii) the parameters of the transformer
are shared across all the dimensions of the random variable, which
significantly reduces the number of total parameters in the model; (iii) the
conditioner does not need to be monotonic, which makes training significantly
more stable; and (iv) we can exploit all the techniques that have been
developed in the last few years to boost the performance of vision and
language transformers (e.g. FlashAttention Dao et al. 2022, PagedAttention
Kwon et al. 2023).

In T-NAFs we define the invertible transformation and conditioner as follows

$\displaystyle t(c(x_{1},\dots,x_{i-1}),x_{i})=t\big{(}x_{i};\
\bm{\psi}_{i}=\text{TN}(x_{1},\dots,x_{i-1},i;\ \bm{\theta})\big{)},$ (6)

where TN is a Transformer Network parametrized by $\bm{\theta}$, that takes as
input the random variable and the index $i$ (implicitly included via
positional encoding) and generates the parameters $\bm{\psi}_{i}$ of the
transformation $t$. We do not impose any particular constraint on the
transformation $t$, which can be any type of invertible transformation. In
Section 4.2 we describe in more details some of the transformations we have
used in our experiments. Empirically we have observed that a compact monotonic
neural network (taking the variable $x_{i}$ as input and producing
$y_{i}\in[0,1]$ as output) offers the best trade-off between performance and
model overhead. This is equivalent to a parametric Cumulative Distribution
Function (CDF). Crucially, because of the single input-output structure, this
network does not require MADE masking and it is stable to train due to its
limited size.

Importantly, in Equation 6 there is a single set of parameters $\bm{\theta}$
in the conditioner that are shared across all the dimensions of the random
variable, while in both NAFs and B-NAFs there is a separate set of parameters
$\bm{\theta}_{i}$ for each dimension $x_{i}$. This means that T-NAFs requires
significantly less parameters (see the experimental section for more details).

Figure 1: Graphical representation of a T-NAF model. Left: the architecture
includes a transformer neural network conditioner TN and an invertible
transformation $t$. In $\text{TN}(x_{1},\dots,x_{D},i;\bm{\theta})$ each
dimension of the random variable is linearly projected to get an embedding
which is added to a learnable position vector. A sequence of $L$ transformer
layers is used to produces hidden embeddings
$\mathbf{h}_{1},\dots,\mathbf{h}_{D}$ that are passed through a projection
head to generate pseudo-parameters $\bm{\psi}_{1},\dots,\bm{\psi}_{D}$. The
pseudo-parameters are used as part of an invertible transformation
$t(x_{i};\bm{\psi}_{i})$. Right: detailed schematic of a transformer encoder
layer. Inputs are passed through a normalization layer, a Multy Head Attention
(MHA) layer with autoregressive mask, another normalization layer, and an MLP.

###  4.1 Overview of the architecture

The overall design of the proposed transformer is similar to the Vision
Transformer (ViT, Dosovitskiy et al. 2020) with two key differences: (i) we
replace the convolutional projections with simple linear projections, as we
are not dealing with image patches; (ii) we replace the classification head
(taking as input only the first embedding in the sequence) with a projection
head that generates the pseudo-parameters (for all the embeddings in the
sequence). Figure 1 showcase a high-level overview of the T-NAF architecture.

In the initial stage of the pipeline, each dimension of the random variable is
linearly projected to get inputs embeddings, those are added to learnable
positional embeddings and passed to a series of $L$ transformer encoder
layers. The first element of the autoregressive factorization is the marginal
distribution $p(x_{1})$, which does not require any conditioning. To impose
this constraint we use a learnable input embedding as first element of the
sequence. This guarantees that the first set of pseudo-parameters
$\bm{\psi}_{1}$ is generated without having access to any element of the
sequence. This is similar in spirit to the Beginning of Sequence (BoS) token
or class token used in a variety of transformer models (Devlin et al., 2018;
Dosovitskiy et al., 2020).

A detailed representation of the encoder layer is provided on the right-hand
side of Figure 1. The encoder layer performs the following operations:
normalization via layer-norm, attention estimation via Multi-Head Attention
(MHA), normalization via layer-norm, forward pass in a Multi-Layer Perceptron
(MLP). Skip connections are used from the input to the output of the MHA, and
from the output of the MHA to the output of the encoder. An autoregressive
mask is used in the MHA to enforce the autoregressive constraint. The output
of the encoder is a series of hidden embedding vectors $\mathbf{h}_{i}$, which
are passed to a projection head to generate the pseudo-parameters
$\bm{\psi}_{i}$. The pseudo-parameters are used to define the invertible
transformation $y_{i}=t(x_{i};\ \bm{\psi}_{i})$. The projection head can be a
set of linear layers that takes as input the embeddings $\mathbf{h}_{i}$ and
generates the pseudo-parameters $\bm{\psi}_{i}$, or it can simply be an
identity function (see Section 4.2 for additional details).

###  4.2 Transformation

The modularity of the architecture allows using a large and diversified set of
transformations, below we detail some of them.

Affine transformation One of the simplest transformation is the affine one. An
affine transformation is defined by a scale $\mu_{i}$ and shift $\sigma_{i}$
parameter, grouped into a set of pseudo-parameters
$\bm{\psi}_{i}=\\{\mu_{i},\sigma_{i}\\}$, that are obtained as follows

$t_{\text{affine}}(x_{i};\ \bm{\psi}_{i})=\mu_{i}+\sigma_{i}\
x_{i}\quad\text{with}\quad\mu_{i}=\text{LIN}_{\mu}(\mathbf{h}_{i};\bm{\theta}_{\mu}),\quad\sigma_{i}=\text{LIN}_{\sigma}(\mathbf{h}_{i};\bm{\theta}_{\sigma}),$
(7)

where $\text{LIN}_{\mu}$ and $\text{LIN}_{\sigma}$ are linear layers
parametrized by $\bm{\theta}_{\mu}$ and $\bm{\theta}_{\sigma}$ used to
generate the scale $\mu_{i}$ and shift $\sigma_{i}$. The linear layers of the
projection head are shared across all the output tokens. Note that, it is
possible to concatenate multiple affine transformations via function
composition, in this case a separate set of linear layers can be assigned to
each transformation.

CDF transformation A more sophisticated transformation is given by a
parametric Conditional Distribution Function (CDF) that maps $x_{i}$ to a
uniform distribution. Following previous work, we can define the parametric
CDF as a neural network with positive weights (Archer & Wang, 1993; Sill,
1997; Daniels & Velikova, 2010; De Cao et al., 2020). In this case the pseudo-
parameters are the weights and biases of the neural network
$\bm{\psi}_{i}=\\{\mathbf{w}^{(1)}_{i},\mathbf{b}^{(1)}_{i},\mathbf{w}^{(2)}_{i},b^{(2)}_{i}\\}$
which are used as part of the following transformation

$t_{\text{CDF}}(x_{i};\
\bm{\psi}_{i})=\text{sig}\left(\text{tanh}\left(\exp\mathbf{w}^{(1)}_{i}\
x_{i}+\mathbf{b}^{(1)}_{i}\right)^{T}\exp\mathbf{w}^{(2)}_{i}+b^{(2)}_{i}\right),$
(8)

where $\text{sig}(\cdot)$ and $\text{tanh}(\cdot)$ are the sigmoid and
hyperbolic tangent functions. Similarly to Equation 7, the weights and biases
are generated via dedicated linear layers
$\mathbf{w}^{j}_{i}=\text{LIN}_{\mathbf{w}^{j}}(\mathbf{h}_{i};\bm{\theta}_{\mathbf{w}^{j}})$
and
$\mathbf{b}^{j}_{i}=\text{LIN}_{\mathbf{b}^{j}}(\mathbf{h}_{i};\bm{\theta}_{\mathbf{b}^{j}})$
as part of the projection head. The exponential function used in Equation 8
ensures that the weights are strictly-positive enforcing the monotonic
constraint (softplus can also be used), this and the use of invertible
activation functions makes the transformation bijective. It is possible to
invert $t_{\text{CDF}}$ since it has a single input, a single output, and is
bijective. The inverse of $t_{\text{CDF}}$ is an inverse-CDF, which can be
used for inversion sampling following two steps: (i) first sample a value
$y_{i}$ from the uniform distribution, (ii) find the corresponding $x_{i}$ via
a root-finding method (e.g. the bisection method).

Shared-CDF transformation In the CDF transformation defined in the previous
paragraph, the pseudo-paramters are generated by using the projection head and
the transformer embeddings. An alternative approach is to define the CDF
neural network in advance, using a fixed set of weights that are shared across
all the outputs. In this case, $y_{i}$ is obtained via a conditional CDF, that
takes as input $x_{i}$ and the transformer embeddings $\mathbf{h}_{i}$.
Assuming that the conditional CDF is modeled using a neural network with a
single hidden layer, we can formalize this particular case defining a set of
pseudo-parameters $\bm{\psi}_{i}$ and a set of shared-parameters $\bm{\phi}$:

$\bm{\psi}_{i}=\left\\{\mathbf{h}_{i}\right\\}\quad\text{and}\quad\bm{\phi}=\left\\{\mathbf{w}^{(1)},\mathbf{w}^{(2)},\mathbf{b}^{(1)},b^{(2)},\mathbf{\hat{W}}^{(1)},\mathbf{\hat{w}}^{(2)}\right\\}.$
(9)

The pseudo-parameters coincide with the embeddings produced by the transformer
conditioner
$\mathbf{h}_{i}=\text{TN}\left(x_{1},\dots,x_{i-1},i;\bm{\theta}\right)$. The
shared-parameters include
$\mathbf{w}^{(1)},\mathbf{w}^{(2)},\mathbf{b}^{(1)},b^{(2)}$ which are related
to the input $x_{i}$, and $\mathbf{\hat{W}}^{(1)},\mathbf{\hat{w}}^{(2)}$
which are related to the conditional input $\mathbf{h}_{i}$. The shared-CDF
transformation is defined as follows

$t_{\text{S-CDF}}(x_{i};\
\bm{\psi}_{i},\bm{\phi})=\text{sig}\left(\text{tanh}\left(\exp\mathbf{w}^{(1)}\
x_{i}+\mathbf{\hat{W}}^{(1)}\
\mathbf{h}_{i}+\mathbf{b}^{(1)}\right)^{T}\exp\mathbf{w}^{(2)}+\mathbf{\hat{w}}^{(2)}\mathbf{h}_{i}+b^{(2)}\right),$
(10)

where we have discarded any reference to the projection head, as we can assume
it corresponds to the identity function. Note the crucial distinction between
positive-constraint weights $\exp\mathbf{w}^{(1)},\exp\mathbf{w}^{(2)}$ used
for the input and the unconstrained weights
$\mathbf{\hat{W}}^{(1)},\mathbf{\hat{w}}^{(2)}$ used in the conditioning
mechanism. The main advantage of the shared-CDF approach is that the
transformer needs to generate just the vector of embeddings $\mathbf{h}_{i}$
for each token, while in the CDF transformation it needs to generate the
weights and biases of the entire CDF network. Therefore, it is more convenient
to use a shared-CDF transformation when the size of the CDF neural network is
large.

Spline transformation Following Durkan et al. 2019, we also incorporated a
spline transformation based on monotonic rational-quadratic splines. Besides
being comparable to the CDF transformation in terms of performance and
capacity (Durkan et al., 2019), the spline transformation has the advantage of
being easily invertible. Each block of this transformation consists of a
elementwise rational-quadratic spline

$t_{\text{spline}}(x_{i};\
\bm{\psi}_{i})=g(x_{i};\bm{\psi}_{i}=\text{TN}\left(x_{1},\dots,x_{i-1},i;\bm{\theta}\right))$
(11)

where $g(\cdot;\ \bm{\psi})$ is a monotonic rational-quadratic spline function
parameterized by a set of pseudo-parameters $\bm{\psi}$ generated by the
transformer conditioner (see Durkan et al., 2019 for the details on the
parameterization of the spline). To ensure cross-variable interactions, the
elementwise function is commonly followed by an invertible linear
transformation across variables

$t_{\text{linear}}(\mathbf{x};\
\bm{\phi})=\mathbf{P}\mathbf{L}\mathbf{U}\mathbf{x}$ (12)

with the set of shared-parameters
$\bm{\phi}=\\{\mathbf{P},\mathbf{L},\mathbf{U}\\}$, where $\mathbf{P}$ is a
fixed permutation matrix, $\mathbf{L}$ is lower triangular with ones on the
diagonal, and $\mathbf{U}$ is upper triangular. We generate a set of
$\bm{\psi}_{i}^{j}$ from the initial variables $x_{1},\dots,x_{i-1}$ at each
spline block transformation $j$. However, this means that the suggested
$t_{\text{linear}}(\mathbf{x};\ \bm{\phi})$ breaks the autoregressive
constraint and the determinant of the Jacobian cannot be computed in an
efficient way. To fix this issue, we simplify the linear transformation to

$t_{\text{linear}}(\mathbf{x};\ \bm{\phi})=\mathbf{L}\mathbf{x}$ (13)

with $\bm{\phi}=\\{\mathbf{L}\\}$ for all spline transformation blocks except
for the first one. As a result, each block of spline transformation is a
composition of $t_{\text{spline}}$ and $t_{\text{linear}}$, where
$t_{\text{spline}}$ is applied elementwise and $t_{\text{linear}}$ acts on the
whole set of variables.

##  5 Experiments

###  5.1 Density Estimation

In this section we evaluate the performance of our method on density
estimation with real-world data, using classic benchmarks from the UCI
repository. Following standard practice, we evaluate our model on four
datasets from UCI (POWER, GAS, HEPMASS, MINIBOONE) and on BSDS300 (Martin et
al., 2001) a dataset obtained by extracting random patches from the homonym
datasets of natural image. We compare against a variety of methods: RealNVP
(Dinh et al., 2016), Glow (Kingma & Dhariwal, 2018), MADE (Germain et al.,
2015), MAF (Papamakarios et al., 2017), TAN (Oliva et al., 2018), FFORJD
(Grathwohl et al., 2018), SOS (Jaini et al., 2019), RQ-NSF (Durkan et al.,
2019). We also include neural autoregressive models such as NAF (Huang et al.,
2018) and B-NAF (De Cao et al., 2020). The results for other methods are taken
from the corresponding papers. Note that some methods (FFJORD, B-NAF and TAN)
perform a hyper-parameter search over each dataset while we used the same
hyper-parameters in all experiments. We used the following configuration for
T-NAF: one flow, three or five layers in the transformer encoder, embeddings
of size 32, eight attention heads, one hidden layer (64 units) in the encoder
MLPs, the projection head generates the parameters of a CDF network with one
hidden layer (128 units).

Results The results are reported in Table 1. Overall, T-NAF is able to
outperform all the other methods on HEPMASS using the larger model with five
layers. On BSDS300 and GAS T-NAF obtained the second highest score after TAN
and RQ-NSF, but without performing any ad-hoc hyper-parameter search and
without using multiple flows. The only dataset where other methods outperform
T-NAF is MINIBOONE. T-NAF is more efficient than other models for two reasons:
a reduced number of flows and reduced number of model parameters. In terms of
number of flows, most of the other methods concatenate multiple
transformations and permute the order of the variable in between. While this
attenuates possible issues due to the variable ordering, it comes at the price
of an increased overhead and parameter count. T-NAF achieves top performance
using only one flow, therefore substantially reducing the complexity of the
model.

Table 1: Log-Likelihood (higher is better) for 4 datasets from UCI and for
BSDS300. Average and standard deviation over 3 seeds. We report number of
dimensions ($D$) and dataset size ($N$, Millions). T-NAF obtains some of the
best results by using just one flow and no ad hoc architectures. ∗Methods
using ad hoc architectures on each dataset. Best result in bold, second-best
underlined.

POWER GAS HEPMASS MINIBOONE BSDS300 $D=6$ $D=8$ $D=21$ $D=43$ $D=63$ Model
Flows $N=2.0$ $N=1.0$ $N=0.5$ $N=0.04$ $N=1.3$ RealNVP 5 -0.02$\pm$0.01
4.78$\pm$1.80 -19.62$\pm$0.02 -13.55$\pm$0.49 152.97$\pm$0.28 RealNVP 10
0.17$\pm$0.01 8.33$\pm$0.14 -18.71$\pm$0.02 -13.84$\pm$0.52 153.28$\pm$1.78
Glow n/a 0.17$\pm$0.01 8.15$\pm$0.40 -18.92$\pm$0.08 -11.35$\pm$0.07
155.07$\pm$0.03 MADE MoG n/a 0.40$\pm$0.01 8.47$\pm$0.02 -15.15$\pm$0.02
-12.27$\pm$0.47 153.71$\pm$0.28 MAF Affine 5 0.14$\pm$0.01 9.07$\pm$0.02
-17.70$\pm$0.02 -11.75$\pm$0.44 155.69$\pm$0.28 MAF Affine 10 0.24$\pm$0.01
10.08$\pm$0.02 -17.73$\pm$0.02 -12.24$\pm$0.45 154.93$\pm$0.28 MAF MoG 5
0.30$\pm$0.01 9.59$\pm$0.02 -17.39$\pm$0.02 -11.68$\pm$0.44 156.36$\pm$0.28
FFJORD∗ n/a 0.46$\pm$0.01 8.59$\pm$0.12 -14.92$\pm$0.08 -10.43$\pm$0.04
157.40$\pm$0.19 SOS 7 0.60$\pm$0.01 11.99$\pm$0.41 -15.15$\pm$0.10
-8.90$\pm$0.11 157.48$\pm$0.41 NAF-DDSF 5 0.62$\pm$0.01 11.91$\pm$0.13
-15.09$\pm$0.40 -8.86$\pm$0.15 157.73$\pm$0.04 NAF-DDSF 10 0.60$\pm$0.02
11.96$\pm$0.33 -15.32$\pm$0.23 -9.01$\pm$0.01 157.43$\pm$0.30 B-NAF∗ 5
0.61$\pm$0.01 12.06$\pm$0.02 -14.71$\pm$0.38 -8.95$\pm$0.07 157.36$\pm$0.03
TAN∗ 5 0.60$\pm$0.01 12.06$\pm$0.02 -13.78$\pm$0.02 -11.01$\pm$0.48
159.80$\pm$0.07 RQ-NSF (AR) 10 0.66$\pm$0.01 13.09$\pm$0.02 -14.01$\pm$0.03
-9.22$\pm$0.48 157.31$\pm$0.28 T-NAF/3 (ours) 1 0.63$\pm$0.01 12.01$\pm$0.02
-14.88$\pm$0.37 -11.59$\pm$0.67 157.83$\pm$0.10 T-NAF/5 (ours) 1 0.54$\pm$0.01
12.27$\pm$0.01 -13.20$\pm$0.26 -10.67$\pm$0.06 159.41$\pm$0.04

Parameter efficiency We compared T-NAF and B-NAF in terms of number of
parameters vs. test log-likelihood on the UCI and BSDS300 datasets. By
assessing the total parameter count required by each model, we aim to
demonstrate how T-NAFs optimize parameter usage. This evaluation offers
insights into the practical implications of employing transformer-based
architectures in probabilistic modeling, especially in scenarios where
computational resources and model simplicity are important. For B-NAF we
report the parameters and performance of the best model as specified in De Cao
et al. (2020), that is: MLP with two hidden layers, number of hidden units
defined as $D\times 40$, and five flows. For our T-NAF model we report the
number of parameters and performance for the variant based on a CDF
transformation and a transformer with 3 and 5 layers. For a fair comparison we
also include the T-NAF pseudo-parameters $\bm{\psi}_{1},\dots,\bm{\psi}_{D}$
generated during the inference step. The results are summarized in Figure 2.
The figure shows that T-NAF offers a better trade-off in terms of parameter
count vs. performance. In particular, the gap between the two methods gets
larger as the dimensionality of the input increases.

###  5.2 Ablations

In this section, we provide a detailed ablation study to investigate the
impact of different components of our T-NAF on performance. A critical aspect
of T-NAF is the use of different invertible transformations and their effects
on the model’s density estimation capabilities. Our experiments primarily
focus on contrasting the performance of three distinct types of
transformations: Conditional Distribution Function (CDF), shared CDF, and
Spline transformations. Each of these transformations offers unique
characteristics: CDF transformations are expressive and allow mapping inputs
to a uniform distribution, potentially capturing complex patterns in the data
more effectively; Spline transformations can offer high levels of flexibility
and are expected to model intricate data distributions with greater precision.
Additionally, the experiments compare the Shared-CDF approach, where a fixed
set of weights is shared across outputs, with the standard CDF approach, where
the transformer generates the parameters of the neural-CDF separately for each
output. This comparison is crucial to evaluate the benefits of parameter
sharing versus individual parameterization in the context of density
estimation tasks. Furthermore, we explore the impact of varying the number of
layers in the transformer encoder on the overall performance of the model.
This exploration helps in understanding the trade-off between model complexity
(in terms of depth) and performance gains. A deeper model might capture more
complex relationships in the data but at the cost of increased computational
requirements and potential overfitting.

Figure 2: Trade-off between number of parameters and performance. Vertical
axis represents the number of parameters in log-scale (lower is better), and
horizontal axis represent the log-likelihood on the test set (higher is
better). The optimal trade-off is represented by points in the bottom-right
corner. Overall T-NAF offers a better trade-off w.r.t. B-NAF; the gap in
number of parameters between T-NAF and B-NAF gets larger as the number of
input dimensions $D$ increases.

Results A summary of the results is presented in Table 2. The standard CDF
model with five layers consistently outperforms the remaining models across
most datasets. The Spline Flow with five layers is competitive on three
datasets (POWER, HEPMASS, and BSDS300) while suffering the most on MINIBOONE.
Overall, the increase in transformer layers from three to five show a
noticeable improvement across the board. This outcome indicates that
additional layers might help capture more complex dependencies. Only in one
case additional layers seems to reduce the performance (POWER dataset). Using
a shared CDF over the standard one does not seem to give particular benefits,
there is only a marginal advantage on MINIBOONE and BSDS300 when compared with
the CDF with three layers.

Table 2: Ablation experiments on head type. CDF: the parameters of a CDF are
generated by the transformer for each output separately; Shared-CDF:
conditional CDF with parameters shared across outputs and conditional input
provided by the transformer embeddings; Spline Flow: the parameters of a
spline flow are generated by the transformer for each output separately. Log-
Likelihood (higher is better) for 4 datasets from UCI and for BSDS300. Average
and standard deviation over 3 seeds. We report number of dimensions ($D$) and
dataset size ($N$, Millions).

POWER GAS HEPMASS MINIBOONE BSDS300 $D=6$ $D=8$ $D=21$ $D=43$ $D=63$ Head Type
Layers $N=2.0$ $N=1.0$ $N=0.5$ $N=0.04$ $N=1.3$ CDF 3 0.63$\pm$0.01
12.01$\pm$0.02 -14.88$\pm$0.37 -11.59$\pm$0.67 157.83$\pm$0.10 Shared-CDF 3
0.52$\pm$0.02 11.12$\pm$0.09 -14.97$\pm$0.17 -11.34$\pm$0.03 158.10$\pm$0.07
Spline Flow 3 0.64$\pm$0.01 11.99$\pm$0.03 -14.05$\pm$0.14 -11.08$\pm$0.04
159.05$\pm$0.07 CDF 5 0.54$\pm$0.01 12.27$\pm$0.01 -13.20$\pm$0.26
-10.67$\pm$0.06 159.41$\pm$0.04 Spline Flow 5 0.65$\pm$0.01 12.20$\pm$0.03
-13.15$\pm$0.19 -11.09$\pm$0.01 159.47$\pm$0.02

##  6 Conclusions

In this paper we have introduced T-NAFs, a new class of neural flows that
exploits transformers as an autoregressive conditioner mechanism that
generates the parameters of invertible transformations. We have showcased the
performance of the model on a variety of real-world density estimation
problems from the classic UCI benchmarks. Overall the results have
demonstrated that T-NAFs offer a powerful combination of efficiency,
scalability, and flexibility, outperforming existing models like B-NAFs and
NAFs in various scenarios. The innovative use of transformer architectures
within T-NAFs not only addresses the scalability and training stability issues
inherent in traditional autoregressive models but also paves the way for more
sophisticated and nuanced modeling of complex data distributions. Future work
can explore further optimizations and applications of T-NAFs, potentially
expanding their utility in broader machine learning and data science contexts.

Limitations There are two main limitations of T-NAF that are worth discussing.
(i) While T-NAF is efficient in terms of parameters, it can be expensive in
terms of FLOPs due to the quadratic cost of the attention mechanism. However,
this problem can be tackled by using recently proposed methods like
FlashAttention Dao et al. 2022 and PagedAttention Kwon et al. 2023. (ii) T-NAF
has a more complex architecture than NAF and B-NAF, therefore there are more
factors that may play a role in the performance of the model.

#### Acknowledgments

Funding in direct support of this work: Massimiliano Patacchiola, Aliaksandra
Shysheya, and Richard E. Turner are supported by an EPSRC Prosperity
Partnership EP/T005386/1 between the EPSRC, Microsoft Research and the
University of Cambridge.

## References

  * Archer & Wang (1993) Norman P Archer and Shouhong Wang.  Application of the back propagation neural network algorithm with monotonicity constraints for two-group classification problems.  _Decision Sciences_ , 1993.
  * Daniels & Velikova (2010) Hennie Daniels and Marina Velikova.  Monotone and partially monotone neural networks.  _IEEE Transactions on Neural Networks_ , 2010.
  * Dao et al. (2022) Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  In _Advances in Neural Information Processing Systems_ , 2022.
  * De Cao et al. (2020) Nicola De Cao, Wilker Aziz, and Ivan Titov.  Block neural autoregressive flow.  In _Uncertainty in Artificial Intelligence_ , 2020.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Dinh et al. (2016) Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp.  _arXiv preprint arXiv:1605.08803_ , 2016.
  * Dolatabadi et al. (2020) Hadi Mohaghegh Dolatabadi, Sarah Erfani, and Christopher Leckie.  Invertible generative modeling using linear rational splines.  In _International Conference on Artificial Intelligence and Statistics_ , 2020.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Durkan et al. (2019) Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios.  Neural spline flows.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Fakoor et al. (2020) Rasool Fakoor, Pratik Chaudhari, Jonas Mueller, and Alexander J Smola.  Trade: Transformers for density estimation.  _arXiv preprint arXiv:2004.02441_ , 2020.
  * Germain et al. (2015) Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle.  Made: Masked autoencoder for distribution estimation.  In _International Conference on Machine Learning_ , 2015.
  * Grathwohl et al. (2018) Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud.  Ffjord: Free-form continuous dynamics for scalable reversible generative models.  _arXiv preprint arXiv:1810.01367_ , 2018.
  * Huang et al. (2018) Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville.  Neural autoregressive flows.  In _International Conference on Machine Learning_ , 2018.
  * Jaini et al. (2019) Priyank Jaini, Kira A. Selby, and Yaoliang Yu.  Sum-of-squares polynomial flow.  In _International Conference on Machine Learning_ , 2019.
  * Kingma & Dhariwal (2018) Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  _Advances in neural information processing systems_ , 2018.
  * Kingma et al. (2016) Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.  Improved variational inference with inverse autoregressive flow.  _Advances in Neural Information Processing Systems_ , 2016.
  * Kobyzev et al. (2019) Ivan Kobyzev, Simon Prince, and Marcus A Brubaker.  Normalizing flows: Introduction and ideas.  _arXiv preprint arXiv:1908.09257_ , 2019.
  * Kwon et al. (2023) Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E Gonzalez, Hao Zhang, and Ion Stoica.  Efficient memory management for large language model serving with pagedattention.  _arXiv preprint arXiv:2309.06180_ , 2023.
  * Martin et al. (2001) David Martin, Charless Fowlkes, Doron Tal, and Jitendra Malik.  A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics.  In _International Conference on Computer Vision_ , 2001.
  * Müller et al. (2018) Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus Gross, and Jan Novák.  Neural importance sampling.  _CoRR_ , 2018.
  * Oliva et al. (2018) Junier Oliva, Avinava Dubey, Manzil Zaheer, Barnabas Poczos, Ruslan Salakhutdinov, Eric Xing, and Jeff Schneider.  Transformation autoregressive networks.  In _International Conference on Machine Learning_ , 2018.
  * Papamakarios et al. (2017) George Papamakarios, Theo Pavlakou, and Iain Murray.  Masked autoregressive flow for density estimation.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Papamakarios et al. (2021) George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference.  _Journal of Machine Learning Research_ , 2021.
  * Rezende & Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International Conference on Machine Learning_ , 2015.
  * Sill (1997) Joseph Sill.  Monotonic networks.  _Advances in Neural Information Processing Systems_ , 1997.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in Neural Information Processing Systems_ , 2017.
  * Wehenkel & Louppe (2019) Antoine Wehenkel and Gilles Louppe.  Unconstrained monotonic neural networks.  In _Advances in Neural Information Processing Systems_ , 2019.