Flexible Language Modeling in Continuous Space with Transformer-based Autoregressive Flows
Authors: Ruixiang ZHANG
Shuangfei Zhai
Jiatao Gu
Yizhe Zhang
Huangjie Zheng
Tianrong Chen
Miguel Ángel Bautista
Joshua M. Susskind
Navdeep Jaitly
Venue: NeurIPS 2025 poster ()
Keywords: Normalizing Flows
Autoregressive Models
Latent Autoregressive Flows
Primary area: deep_learning

Abstract: Autoregressive models have driven remarkable progress in language modeling. Their foundational reliance on discrete tokens, unidirectional context, and single-pass decoding, while central to their success, also inspires the exploration of a design space that could offer new axes of modeling flexibility.
  In this work, we explore an alternative paradigm, shifting language modeling from a discrete token space to a continuous latent space.
  We propose a novel framework that employs transformer-based autoregressive normalizing flows to model these continuous representations.
  This approach unlocks substantial flexibility, enabling the construction of models that can capture global bi-directional context through stacked, alternating-direction autoregressive transformations, support block-wise generation with flexible token patch sizes, and facilitate a hierarchical multi-pass generation process.
  We further propose new mixture-based coupling transformations designed to capture complex dependencies within the latent space shaped by discrete data, and demonstrate theoretical connections to conventional discrete autoregressive models.
  Extensive experiments on language modeling benchmarks demonstrate strong likelihood performance and highlight the flexible modeling capabilities inherent in our framework.

Flexible Language Modeling in Continuous Space with Transformer-based Autoregressive Flows
Ruixiang Zhang, Shuangfei Zhai, Jiatao Gu, Yizhe Zhang, Huangjie Zheng, Tianrong Chen
Miguel Angel Bautista, Josh Susskind, Navdeep Jaitly
Apple
(ruixiangz, szhai, njaitly)eapple.com
Abstract
Autoregressive models have driven remarkable progress in language modeling. Their foundational reliance discrete tokens, unidirectional context, and single- pass decoding, while central to their success, also inspires the exploration of a design space that could offer axes of modeling flexibility. In this work, we explore an alternative paradigm, shifting language modeling from discrete token space to continuous latent space. We propose novel framework FarFlowLM, that employs transformer-bact continuous representations. This approach unlocks substantial flexibility, enabling the construction of models that can capture global bi bi-directional context through stacked, alternating -drrection autoregressive transformations, support block-wise generation with flexible token patch sizes, and facilitate hierarchical multi-pass generation process. We further propose new mixture- based coupling transforma- tions designed to capture complex dependencies within the latent space shaped by discrete data, and demonstrate theoretical connections conventional discrete autoregressive models. Extensive experiments language modeling benchmarks demonstrate strong likelihood performance and highlight the flexible modeling capabilities inherent in our framework.
autoregressive normalizing flows [73] to model these
Introduction
Transformer- based autoregressive models [5, 66] have emerged as the dominant paradigm for lan- guage modeling, achieving remarkable performance by predicting discrete tokens one at time under the cross-entropy objective. By scaling both model size and training data, these models excel at next-token prediction and have become the foundation for modern natural language generation systems.
Their remarkable success provides a strong foundation and inspires further exploration into the design space of autoregressive sequence modelling. While the established paradigm of discrete, typically unidirectional autoregressive generation offers powerful and well-understood framework, we consider whether alternative formulations might offer new dimensions of flexibility and open different avenues for model construction. This work investigates such possibility: What if antoregressive language modeling were reformulated within continuous latent space? Moving continuous representations allows for the preservation of the sequential factorization familiar from autoregressive methods, while potentially unlocking new modeling capabilities.
This adoption of continuous latent spaces is a direction also pursued by other generative frameworks; for instance, diffusion models have demonstrated strong capabilities in this domain. Our investigation, however, distinguishes itself by concentrating on autoregressive normalizing flows for modeling the joint distribution of these continuous sequences. We pursue this specific avenue because normalizing
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

flows offer expressive power for density estimation, and critically, their intrinsic sequential processing aligns them closely with the fundamental mechanisms of discrete autoregressive language models. This alignment provides unique vantage point: it allows us to explore how the core principles of autoregressive modeling can be evolved and potentially augmented when transitioned into a continuous, learnable transformation framework, thereby seeking to extend and enrich this successful paradigm.
Building on this perspective, we propose a novel framework that employs Transformer-based autore- gressive normalizing flows [73] to model these continuous latent representations. This shift from discrete token space to a continuous latent space is key to unlocking substantial modeling flexibility, Our formulation facilitates the construction of models capable of capturing global bi bi-directional context through stacked, wise generation with adaptable token patch sizes and enables a hierarchical multi-pass generation process, allowing for the observation and influence of sequence formation through intermediate stages. These capabilities arise naturally from the continuous and invertible nature of the flow-based transformations.
dirrettinn autoregressive transformations. It also supports block-
To this end, we present TarFlowLM, Transformer-based Auto-Regressive Flow Language Model that uses mixture- coupling transformations to effectively model the complex, multi-modal distributions created when mapping discrete data continuous latent space. Our key theoretical contribution is establishing the equivalence that transforms these mixture distributions into exact, invertible normalizing flow layers. Specifically, we show that a dinees-ammnii Mixture of Gaussians (MoG) distribution can be realized as ID Mixture-CDD flow, and a multi-dimensional MoG distribution as a tinuous approach and standard discrete autoregressive models. Experiments on language modeling benchmarks show that our method achieves strong likelihood performance and, importantly, enables greater modeling flexibility. Our work expands the possibilities of language modeling by extending autoregressive methods into the continuous domain.
flow. We also draw theoretical links between our con-
Background
Normalizing flows offer method for constructing flexible probability distributions over continuous variables [51]. The core idea is to start with a simple base distribution Phose (u) defined on Rd, often standard Gaussian, and transform samples ~ Phase (u) using an invertible and differentiable function f IRd + Rd, known follows potentially much more complex distribution p(z). A key advantage of this approach is that the probability density function p(z) can be computed exactly using the change of variables formula:
diffcomorphism. This process yiclds variable
f-' that
log p( 2) log Phase (z)z) + log |det (z(
Here, (z) is the Jacobian matrix of the transformation f evaluated at z, and det(-)] denotes the absolute value of the determinant. Complex distributions are typically modeled by composing multiple simple transformations f f(()) o f()), where each layer f() is designed to be easily invertible and have a tractable Jacobian determinant. The log- densiiyitt then becomes a sum of log- regedettrminant terms from each layer, plus the log-density of the final transformed variable under the base distribution:
log p(z) log Plase (u) +
Σ 8=1
log det Jyco (h(e=1))
(2)
where h = z. h0
and u (z).
The main challenge in designing normalizing flows lies in choosing transformations f((i) that are both expressive enough to model complex data distributions and computationally efficient, particularly regarding the calculation of the inverse f-1 and the Jacobian determinant det Jy.
A prominent and widely used class is autoregressive flows [39, 50]. These structure the transformation such that u; depends on zj with j < This ensures the Jacobian J, is triangular, making its determinant the product of diagonal entries, computable in (d)) time. A common choice for the element-wist transformation is the affine function: 2; = a; + 8.. In the autoregressive setting, the scale a; > 0 and shift B, for dimension are functions of the preceding dimensions, e.g..

---

Train
Generation
Cat
sat
mat
Cat
sat
on
mat
V=50257
[T. DJ
Un Un-Potchifyy
kx
Transformer AR Flow
(a) Seq. length T, Patch size K, vocab, size V.
mat
Transformer AR Flow
mat
(b) Forward pass, patch size
Transformer AR Flow
On
(c) Latent embedding channel mixing.
Figure I: Our approach transforms discrete token sequences into continuous latent space modeled by stacked Transformer-basee autoregressive normalizing flows, creating fully invertible pipeline for training and generation. Left: During training (Red top-down arrow), tokens encoded into Gaussian latents by a VAE encoder, transformed through flow layers into standard Gaussian variables. and reconstructed by tied decoder. At inference (Blue bottom-up arrow), samples from a standard Gaussian are inverted through the flows and decoded into text. Right: Key flexibility features: (a) Tokens can be grouped into patches and embedded into latent space using mixture- of of-Gaussians transformations with dynamic vocabulary sizes; (b) Block-wis autoregressive flows alternate directions, enabling context exchange between initially independent tokens; (c) Channel mixing permutes latent dimensions between flows to facilitate cross-token information sharing.
a= (Zzz)) and B; = (zzzs). If the base variables u, are drawn from standard Gaussian (O0, 1). this affine transformation implies that each 2; conditioned on the context Ze; follows Gaussian distribution (22; μ4 (zzz),,, a, ()) To ensure all variables influence each other, multiple autoregressive layers are typically stacked, often with variable permutations between them.
3 Autoregressive Language Modeling in Continuous Latent Space
Notation. We use italies for scalars (x,), bold for vectors (z,), and bold with indices for sequences (X1.7, 21:T).
We propose utilizing the Variational Autoencoder (VAH) framework [38]. The idea is map sequences of discrete tokens X1.7 = [#I...., ØT) into sequences of continuous latent variables Z1:7 (Z,...., ), where each z, E Rdd. We then model the joint distribution p(z1.r), the prior, to capture sequential dependencies. This approach aims preserve the sequential factorization structure familiar from autoregressive models while gaining the benefits of end-to-end differentiability inherent in continuous spaces.
Our model consists of three components (see Fig. 1): an encoder q(Zi.T] Xx:)) and decoder p(X (Z.T) that map between discrete token sequences and continuous latent sequences, and

---

Table 1: Comparison of ID and d-D Mixture of Gaussians (MoG) normalizing flow layers, trans- forming imput z to standard Gaussian u/ u. We use C to denote the conditional context.
1-D Mixture CDF Layer
d-D Mixture Rosenblatt Layer
Input
Scalar
Vector z R
Output
Scalar E
Vector u E IR"
Base PDF (Passe)
(u; 0. 1)
(u 0, Ia)
Input PDF (p(- ())
Puixx 1 (2 C) = E; (T;
Pmix Pmixx-(2(()) Ej Tz m T
Forward Map (f)
((Fmix (2; C))
u, (F, C, (()),,,, ..l), for 1.dd
Fini:: CDF of p(z/C)
(Rosenblatt, see Alg, 1)
Inverse Map "
Pimis (b)
Solve F2(2el... (bu;) for 2, for = 1..d
log det. Jy
log Pmix- (=|C)
log Pouse (u)
log Pig.ain d C) log Phase
u)
an autoregressive prior p(Z1:r) over the latent space. The model parameters (covering encoder, decoder, and prior) are learned by maximizing the Evidence Lower Bound (ELBO) on the data log-likelihood log p(x.T)):
((xIT) Ex
log ((x,7 ZITT)
log p(ZI.T log q(zi.)
(3)
decoder
prior
encoder
3.1
The Bridge: Encoder and Decoder
The encoder and decoder facilitate the transition between the discrete data space tokens and the continuous latent space of vectors 2,-
Encoder q(1 x1:T). We use factorized encoder, q(Z (x.T) II g (zt), where each
token à E (l...., V} is assigned a learnable mean f4g E Rd and variance o for an isotropic Gaussian: g(2//0, = k) = (Miz; (a, Tf = (2x).. Thus, the codebook {Ms, of) kJ k== token Gaussian distribution latent space.
21V maps each
Decoder p(x1:r [i:r). The decoder maps each latent vector back to a distribution over tokens, with again factorized structure p(x1.r| /1.) I p(za/z)). We reuse the encoder's s Gaussian parameters {uk, o k== for the decoder with Bayesian posterior parameterization assuming p(a,=k)g(z, (x.=k)
Nx(2,)) NT(a) In other words, the probabil-
uniform prior: p(z; = k k,2,)) = Σソ p(
ity of decoding token l from 2, is determined by how likely 2; is under the k-th encoder Gaussian, compared to all other tokens.
3.2 Autoregressive Prior Modeling for p(Z).T)
To model the prior p(Z1.T) over latent sequences, we use the standard autoregressive factorization: p(ZI.T) 11 p(zz/zer), where Zer = (Z)...., 1;-1) is the history. The main task is to define the conditional distribution p(z, k.:)). We present forms of parameterizing the conditional p(z//ze)) one with dimension-wise autoregressive factorization Sec. 3.2.1, the other with token-wise autoregressive factorization Sec. 3.2.2. In each formulation, we use mixture- based probability distribution, and show that we can equivalently convert the mixture probability density into an equivalent invertible normalizing flow layer. Finally, we show how stacking such flow layers yields an expressive and flexible prior model p(Z).r).
3.2.1 Dimension-wis Autoregressive 1-D Mixture Modeling
We start from decomposing the conditional p(z, (zc)) using the chain rule across the dimensions of Zz
p(z: (2() = II p(e,, Ze,cs)
(4)
Here, žτ,i the i-th scalar component of 2, Zt = (21,1,..,; 21,51) are its preceding components. We parameterize each scalar conditional density p(z,,zzze <i) mixture of V 1-dimensional

---

TEXT8 (Test BPC
WpenWe (Validation Perplexity .)
Space
Type
Method
BPC
Method
Perplexity
Diffusion
Plaid [27]
1.48
Gaussian Diffusion
- 27.28
Diffusion
BFN [25]
1.41
MAC [59]
1.40
D
AR
Transformer AR [1]
1.23
Transformer AR
17.54
D
Diffusson
SEDD Absorb [45]
S 1.39
SEDD Absorb
24.10
D
Diffusion
MD4 [58]
1.37
MD4 [58]
22.13
D
Diffusion
TCSM [77]
1.25
GenMD4 [58]
21.80
Diffusion
EDLM [70]
1.24
MDLM [56]
S 23.21
C
NF
Latent NE [82]
1.61
NF
CNF [41]
1.45
NF
Argmax Flow [31]
1.45
NF
TarFlowLM Affine
1.54
TarFlowLM Affine
Σ 148.22 21
C
NF
TarFlowLM Mix-1
S 1.37
TarFlowLM Mix-1
S 27.11
NF
TarFlowLM Mix-d
1.30
TarFlowLM Mix-d
22.64
Table Performance comparison across different model types and datasets. C=Continuous space, D=Discrete space, AR=Autoregressivve Nr=Normaiziig Flows.
Gaussian distributions. The parameter bundle [Te, mr,i, oi,] is predicted by causal Transformer- based model from the context (z<t, Zt,e)). The conditional density is then, dropping the subscripts f, here for notational brevity:
plzraleer Zt, <i)
Σ π (2 o'k)) n,, m, o²| TTransformer(ccss Zt, (5)
We denote this 1-D mixture PDF as (t. We discuss the connection to the mixture-of-logisties coupling [29] in Sec. D.
Transforming 1-D Mixture Density to an Invertible Flow. The conditional 1D mixture density Pmiis-(-) (-) from Eq. (5) naturally defines an invertible normalizing flow layer. This layer transforms an input žt,i into Wr such that Wr, follows standard normal distribution A (0, 1). The transformation given by ,t,
(Fmix-1 (2t,i Zcti Zt) where Fmix-1 the CDF of Pmix-1 and 4-1 is the inverse standard normal CDF, This mapping is invertible and its key properties are summarized in Table 1 (see Appendix A for details).
The log absolute Jacobian determinant for this transformation, central to normalizing flows, takes the form log the.s
Sea.s log Pmixx (zt,a Zee 2,) log ((u,;; 0, 1). This expresses the log-density of F,, under the mixture model as the sum of the log-density of ue,i under the standard normal and the log-Jacobian:
Oun
log Pmix-1 1 ZZc, Zt, log (u,,s 0, 0, log dzy,
(6)
In this way, learning the mixture parameters i m, 2²2 is equivalent to learning ID Mixture-CDE flow layer. We summarize this the following proposition, where we omit the condition Zet, Zt, for brevity.
Proposition 1. (Single-dim) Equivalence of mixture distribution and Mixture-CDF Flow.
Let (mix-1(2) Σk= T2;(2; ma, of) be 1D MoG probability density function, and let F Fm) be corresponding cumulative distribution function (CDF). The transformation f:R + R defined by u = f(z) " (z)), where & is the standard normal CDF, an exact normalizing flow between the density Pmix-1 (z) and the standard normal density (1;;0,1)
Proof. The result follows directly from the probability integral transform theorem, see Appendix A for details.
5

---

3.2.2 Token-wise Autoregressive d-D Mixture Modeling
Alternativll,, p(z; Z<e) can be modeled directly d-dimensional distribution. We model p(z; z directly as a d-dimensional distribution using a mixture of Gaussians with shared global codebook & = {(ux, k containing V Gaussian components. The network only predicts the V mixture weights ग. (Zcr) (a --dim vector) from the history Zesi
plz Z<t)
Σ n. (Z< Mk, pkx of t., (Z<r) Transformer(zet)
(7)
This approach requires predicting only V parameters per time step, making it computationally efficient. We denote this d-dimensional mixture PDF as Pmix a (itt).
Connection to Discrete AR Language Models. This formulation is closely linked to discrete autoregressive language models; see Sec. F for details.
Transforming d-D Mixture Density to Invertible Flow. The d-dimensional conditional mixture density d can also be realized as an invertible normalizing flow layer. Unlike the 1D case, where the CDF and its inverse are available in closed form, direct CDF-based transformation is not tractable for general d-dimensional mixtures. Instead, we use a sequential transformation u ga(z; z<t) based on Rosenblatt's 's theorem [55], which maps z E R to E R such that u follows a standard Gaussian (00, Ia). This process, described in Algorithm proceeds dimension by dimension: at each step i, a ID Mixture- CDF transform is applied to 2,, conditioned on the previous components Z.g. The full derivation, including the inverse transformation, is given in Appendix B.
The log absolute Jacohian determinant of this transformation has a form directly analogous to the 1D case (see Appendix B): log (eet .g.. (z)l log Pmix-a Ze (z<c) log (u: 0, La). This relationship. also summarized in Table 1, shows that the log-probability z under the d-dimensional mixture is given by the change of variables formula:
log Puixx [ (z+a)= log (u: Ia) + log |det (z)l
(8)
We summarize this in the following proposition, where we omit the condition Z<t for brevity.
Proposition (Multi-dim) Equivalence mixture distribution and Rostebllossn Flow.
Let Pmix al Ek=1 Tz; m, La be a d-dimensional MoG PDF. Let the transfor- mation g RRd Rdd be defined sequentially for i ,,...,dhy
u1 = g. + - (F(2z()))
where (z(zzzzz)) the CDF of the true conditional probability distribution p(z, (e)) derived from Pmix-d (z) (see Lemma 1). This transformation is an exact normalizing flow between Pmix-a(2 z) and the standard d-dimensional normal distribution N(u O,
Proof. See Appendix B for details.
We also provide proof of the invertibility and differentiability of the proposed Rosenblatt Flow in Proposition 3 in Appendix Sec. B.
3.3 Stacking Flow Layers for Expressive Priors
To construct flexible prior p(Z).T) over the latent sequence, stack multiple invertible flow layers. Each flow f(() transforms its input h (f-1) to output h starting from (0) 21.r and ending with
h i:T- which is modeled by standard Gaussian base distribution Phase (uLT).
(I)
U1.T
The model is trained by maximizing the ELBO. For single sample, the por-token objective is:
ENaaen {allee) ) + log (u, O. I) +
C, ~ - log
( 0, 7) log P (h! (c())) - log
(2)
1=1
(9)

---

Elad of Capping Afrx a. Moiare
Capple
Progressive Layerwize Training
Modan Coplin
321
mixd Rosenblatt Mixture
Model varlant
mix-1 1 CDF Mixture
ME1RO
5:5
Buse
NELBO
23000
50000
110000 Tee Stepe
123000
Figure 2: Ablation on model variants.
Figure 3: Affine vs. Mixture coupling.
where the first term applies when using the tied encoder-decoder (equation 3.1). Here, Pmix, is the mixture density for the input layer l, and c( denotes the conditioning information.
(f)
For each flow layer ((0), the objective is twofold: (a) maximize the log-likelihood of its input under the learned mixture p and (b) encourage the output h,") to match the standard normal distribution. (e)
Log.V(h) (6) 0, I) for each layer and token. This stacked structure enables the model to incrementally transform the latent variables and capture complex dependencies. Different flow types and autoregressive directions can be combined across
This is achieved by maximizing log IggPui Pmix, (h, (c()))
layers for greater expressiveness.
Progressive Layer-wise Training Our stacked normalizing flow prior naturally supports progres- sive layer-wise training. which simplifies optimization for deep flow models. We start by training the first block of flow layers (e.g.. f(1) to ((k))), possibly together with the VAE encoder and decoder, until their ELBO contribution (Eq. equation 9) stabilizes on validation set. Once converged, we freeze their parameters and add the next block of layers (e.g.. (kt++) to (*a)), training only the new layers to further transform the fixed output from the previous block. This process repeats, adding and training new layers while keeping earlier ones fixed, until all L flow layers are trained. Progressive training improves stability, each new layer incrementally normalizes the representation, and helps the model reach better optima in deep architectures.
4 Experiments
We conduct experiments to evaluate our proposed continuous latent space framework for autore- gressive language modeling. The primary objectives are twofold: first, to assess likelihood performance on standard language modeling benchmarks, and second, to demonstrate the flexible modeling capabilities enabled by the continuous and flow-based formulation.
Datasets and Evaluation. We evaluate models on standard language modeling benchmarks, specifically TEXT8 [47] and OPeNWeBTexT [22]. For character-level tasks, we report bits per character (BPC), while for word-level tasks, perplexity (PPL) negative ELBO (NELBO) on the respective test sets, given the VAE-based nature of our model.
Model Configurations. We conduct our experiments using the GPT2-Small architecture, following the setup in [1, 56, 58]. This model has 12 layers, hidden size of 768, and 12 attention heads. All conditional distributions within the flow layers are parameterized by Transformers, as described in Sections 3 and 3.2.2. We primarily explore two model configurations:
Mix-1: This variant employs a dimension-wise autoregressive I-D mixture prior (p(z,, Zeti Zi,4s): Sec. 3.2.1). The model consists of three flow layers with alternating directions (left-to-right and right-to-left), starting from the z - direction. The number of Transformer layers in each flow is [2. 2, 8). Unless otherwise specified in ablation studies. we use V = 64 mixture components for OPENWESTENT and V = 27 for TEXT8. To achieve dimension-wise autoregression, we use the same MI.P-based MAF [50] [41].
Mix-d: This variant uses a token-wise autoregressive d-dimensional mixture prior (p(z, z<r); see Sec. 3.2.2). The model is composed of three flow blocks, trained progressively
7

---

Cerpur scn of Generatin Pertiout with Patch Suee]
Effect of Number of Mietures
36
Mixtures
num mixtures-1
3.4
2
Geseutie Pe Pe-deoty (bpoca))
rum num_nisturew-d
3.2
n
Text BPC
3.0
num mixtures-1
Oun nach baen2 Ral
AB 1TPPPPP..
num 27
28
26
24
10000
20000
30000
60000
50000
Padi Lergt
230
Training Step
Figure 5: Generative PPL evaluation with patch size 2.
Figure 4: Effect of mixture components numbers.
Table 3: Unconditional samples from TarFlowLM.
seening that he is the only person here. He is tall. He is balding, with face that is nei ther sell-shaped, nor is it really human. "An" The man is a mon, isn't curious. "How does
Soul! I'n genna be in love with ny dog!" She said, giving him big hug and then neking hane by now! (Hoping that she' d give hin hug toa...) The two boys cane their parents" bedreon to sleep. After that, they continued to play game chess with their parents. (It was bit of a joke in our world, so I'll just say...) .) The two boys had long chat. A lot of then were talking about the idea of should look like...
she'd cone
he?" The nann the niddle get there?"Nu"lll "ile does. He's not an ordinary person. He's not monstor. He's human. He's huran th extraordinary skills. He's an amazing man, "in other nan had been watching the nan, who called human, in silence for alnost hour. He didn't t look too happy, but he was still a hunan, so it didn't matter. He his body, but the
went
I will nake hin a cute dog! Anyway, they started talking about things like... "Do you to nake ne your dog?" Then they began to tease "Tell ne, how long are gonna stay?"
what
looking that nan, man hinself, which was why he was so excited. He had watched for so long, so long
outlined in Sec. 3.3, The first flow layer uses Mix-d Mixture-Rosenblatt coupling, with the Gaussian mixture codebook tied to the encoder. The remaining two layers use Mix-1 Mixture- CDF coupling. For OPEN WEB TEXT, we set V 50257 in the first flow layer and = 64 in the next two; for TEXT8, we use = 27 in the first layer and V = 2 in the remaining layers.
We use a latent embedding dimension of d 16 for all OPENWEBTEXT experiments and d 5 for all TEXT8 experiments.
Perplexity evaluation We assess the language modeling capabilities of our approach by training models on three standard benchmarks: TEXT8, and OPENWEBTEXT. For each dataset, we evaluate the NELBO of our trained models on their respective validation or test splits. Specifically, for TEXT8, we follow the established character-level setup and data splits, typically using fixed-length text chunks for training. For OPENWEB TEXT, we train models using the common GPT-2 tokenization and context length of 1,024 tokens, reserving a portion of the dataset for validation, upon which NELBO is reported. We report results in Table 2.
Results discussion. Table 2 shows the discrete autoregressive Transformer is best (Text8 1.23 BPC; OWT 17.54 PPL). TarFlowLM is competitive among continuous-spaceffloww baselines: Mix-d (22.64 PPL; 1.30 BPC) outperforms several diffusion/NF methods and approaches the strongest diffusion results; Mix-1 on par with Gaussian diffusion. While below discrete AR on PPL/BPC, FarFlowLM adds bi-directional, block-wise, and hierarchical flexibility.
Importance of Mixture Coupling We provide empirical evidence for the effectiveness of the mixture coupling layers introduced earlier. Specifically, we compare two model variants on the OPENWEBTENT dataset: (1) baseline using only affine coupling layers (6 flow blocks, each with 2-layer causal transformer, totaling 12 transformer layers), and (2) an identical architecture but with Mix-1D (mixture) coupling layers. As shown in Fig. 3, incorporating mixture couplings leads to substantial improvement in negative ELBO, clearly demonstrating their importance for modeling discrete text data.
8

---

AR Flow
weokthier How- Saic.
(j Clare said,
"It not rected for beval, he seld.
probices) any
going that (xmpptestty
"It is not a civil narriage any sore, but we are going to use that as iotivation to make it even more fun and (vaca Kids, said.
The two nothers were esting
through
Mettens
Estcut
Hoigechin,
entire executive the Sunns (physscist
The two brothers) were (s studying
under
Esona
Cobral
Cocheux,
the top executive of the Christian [istener
AR Flow
#jendoftextjj
Bon said.
wealthier
Haw Sald.
"It is not civil probless any (difficuitry but we going to get that sinplicity to make it even ddfffccutt protracted for beyeny" beadiers
civil probless any (Bore,) but are going to get that simplicity to make it even store prot racted for boys,"" he (said.)
dsfferent no bothers citing Hustoreenn
Eongeeppedry
the entire
iexceurávveee
the Paychiat
iccooin
citing (through Muttens
Calcul
Malgechin,
entire
the Sunn physicist
Figure 6: Demonstration of the stacked autoregressive flows for text editing in continuous latent space. Each panel shows the transformation of the decoded text from input (green) to output (red), while also highlighting the difference (edits) made by the flow layer.
Flexible Patch Size: Block-wise Multi-token Generation Our framework offers greater flexibility in sequence modeling by moving beyond the traditional single-token generation approach. Instead, processes "patches" of multiple tokens at once, allowing for simultancous prediction and generation. In contrast, standard discrete autoregressive language models even those trained with multi-token objectives [21, 42]-stii -till generate tokens strictly one at time to maintain autoregressive consistency, since generating multiple tokens together breaks their core assumptions. We address this by using a continuous latent space and stacked normalizing flows to capture dependencies within each token patch. First, stacked autoregressive flows with alternating directions (see Fig. 1, patch size two) enable later layers to propagate and refine joint information, even if earlier blocks treat tokens as conditionally independent (as shown by the shaded area). Second, intra-patch dependencies further modeled by mixing permuting latent embedding dimensions between flow layers (channel mixing, Algorithm 2), allowing information to flow between token positions within a patch. Together, these mechanisms let the model effectively learn intra-patch dependencies. Block-wise generation also changes the computational landscape: shifts some of the workload from sequence length to model depth (number of flow layers), reducing the effective sequence length. As shown in Fig. 7, this can significantly lower the total FLOPs needed to process a sequence (e.g.. 1024 tokens) compared to standard transformers, especially with larger patch sizes, all within generalized autoregressive framework. We also trained a discrete autoregressive language model with a two-token prediction objective and compared its generative perplexity (measured by GPT-2 Large) to model, conditioning on varying prefix lengths. The results show a large gap, highlighting how standard AR LMs break down when the next-single-token assumption is violated. See Fig. for generative PPL with patch size 2. See Table 3 for unconditional samples from FarFlowLM.
Flexible vocabulary size. key flexibility of our model is the ability freely choose the number of mixture components in each stacked autoregressive flow layer, especially when using mixture-based transformations (see Sections 3.2.1 and 3.2.2). This number, denoted V in Eq. equation (ID) and Eq. equation 7 (d-dimensional), serves as an internal "vocabulary" for the flow, and can be set independently for each layer.
This design enables strategies like coarse-to- to-rse-to-fine -ine modeling: earlier layers can use fewer components to capture broad structure, while later layers more refine details. Importantly, the internal mixture size does not need to match the data's 's discrete vocabulary.
Empirically, this flexibility is robust. On TEXT8, as shown in Fig. 4, the model performs consistently well across a wide range of mixture counts (2 to 27), even though the dataset's vocabulary is 27 characters. On OPENWEBTED using 64 or more mixture components in the mix-1d coupling layers

---

Nam
Parkle Crodin San.
Total Forward FLOPs Comparison
A Trasdan thech sm
(hach
Forcaral TR PP (Log Scale)
600
60
18
Mode Configuratinn
hdo
Flx
Figure 7: Left: Latent evolution analysis. Right: Transformer FLOPs comparison.
matches the performance of using the full GPT-2 tokenizer size (50257). In summary, our framework allows the internal "vocabulary" of mixture components be tuned for the task, decoupling from the data's vocabulary and providing a practical lever for model design and efficiency.
Flexible text editing in continuous space. Our framework models language in a continuous latent space using stacked autoregressive normalizing flows, enabling to decode intermediate representations at each flow layer. This provides a unique, step-by-step view of how the model incrementally edits and refines latent states into coherent text-a process that is difficult to realize discrete token spaces, key advantage this approach its natural support for coarse-to-fing generation: early flow layers make broad, global changes to the latent sequence, while later layers focus on increasingly fine-graine ain-igr adjustments. We quantify this progression using several metries, including the mean L2 norm of token embeddings, mean pairwise cosine similarity, Participation Ratio (PR), and centroid movement between layers. Typically, we observe that the mean L2 norm decreases and cosine similarity increases across layers, indicating that representations become more compact and internally coherent. The centroid movement diminishes with depth, reflecting the transition from coarse fine editing. Alternating flow directions (e.g., left-to-right and right-to-left) further enrich the representations by integrating information from both past and future context. This hierarchical, coarse-to-fine refinement is direct benefit of operating in a continuous space, where smooth, learnable edits are possible at every stage. Fig. visualizes per-layer edits on decoded text. For further details and quantitative results, see Appendix G.
Model Ablation We performed ablation studies on OPEN WEB TEXT to evaluate key architectural choices (see Fig. 2). Adding dimension-wis autoregressive affine flows base model yields only minor NELBO improvement. Replacing affine couplings with 1-D mixture-CDF transforms greatly improves NELBO, and using full d-dimensional -esteblaR flow further boosts performance. Progressive layer-wise training fereeiingeeigggeee earlier flow blocks aaciives the best results, demonstrating the benefit of simple curriculum over flow depth. We also examined the effect the number of Gaussian components V in 1-D mixture-CUl CDF couplings on TEXT8 (Fig. 4). We did not use MAF module to isolate the mixture coupling effect. single component (affine flow) gives high BPC, while all other > 1 number of mixtures can work well.
5
Conclusion
This work introduced TarFlowLM, a novel framework that recasts language modeling in continuous latent spaces using Transformer-based autoregressive normalizing flows. Our approach achieves strong likelihood performance while enabling significant flexibility, including bi-directional context, block-wise generation, and hierarchical refinement. These findings highlight the potential of continu- ous latent variable models with mixture-base transformations to advance flexible and expressive sequence generation.
10

---

Limitations
While our framework, TarFlowLM, demonstrates strong likelihood performance and notable modeling flexibility, we acknowledge certain limitations inherent to the current instantiation, particularly con- cerning sampling efficiency. Autoregressive normalizing flows, by their nature, involve a sequential generation process. Each step in the generation depends on the previously generated ones, which can lead to slower sampling speeds compared to models that allow for more parallelized generation, such some non autoregressivve or diffusion-basee based approaches.
In this work, our primary focus was on establishing the viability of transformer--asee autoregressive flows for flexible language modeling in continuous latent spaces, emphasizing likelihood estimation and the exploration of novel modeling capabilities like block-wise generation and hierarchical refinement. Consequently, systematic investigation into optimizing sampling speed or exploring advanced sampling techniques (e.g., parallel decoding strategies, distillation) for this class of flow- based language models was beyond the scope of the current paper.
Improving the sampling efficiency of autoregressive flow models for text generation remains an important and active area of research. We consider this a promising direction for future work, which could involve developing specialized architectural modifications, exploring alternative flow parametrizations, or adapting techniques from other generative modeling paradigms to accelerate the sampling process without compromising the expressive power and flexibility demonstrated this paper.
Broader Impacts
The research presented in this paper, TarFlowLM, introduces a novel framework for language modeling in continuous latent spaces using transformer- based autoregressive normalizing flows. As with many advancements in machine learning and artificial intelligence, particularly in the domain of generative models, this work has the potential for range of societal impacts, both positive and negative.
Potential Positive Societal Impacts.
The enhanced flexibility offered by our approach-sur such as bi-directional context modeling, block-wise generation, and hierarchical refinement-com could lead to significant positive developments.
Improved Creative Tools: Models capable of more nuanced and controllable generation could serve as powerful assistive tools for writers, artists, and designers. helping to brain- storm, draft, and refine creative content. The ability to edit text in continuous latent space and observe intermediate generation steps might offer new paradigms for human-Al collaboration in creative tasks.
Enhanced Controllability in NLP Applications: The flexible architectural components, like adaptable patch sizes and independent mixture component selection per layer, might enable finer-grained control over the generation process. This could be beneficial for appli- cations requiring specific stylistic attributes, content constraints, or conditional generation tasks (e.g.. personalized dialogue systems, summarization with specific focuses).
Advancements in F Fundamental Understanding: By exploring language modeling in continuous latent spaces, this research contributes to a broader understanding of how complex sequential data can be represented and manipulated. Such foundational insights can spur further innovation in machine learning for sequences beyond just text, potentially impacting fields like bioinformatics or time-series analysis.
Potential for More Efficient Modeling: While sampling efficiency is noted limita- tion, the architectural flexibility (e.g., block-wise processing potentially reducing effective sequence length for attention mechanisms) hints avenues for developing more computa- tionally efficient models for certain tasks or sequence lengths, which could make advanced NLP capabilities more accessible.
Potential Negative Societal Impacts and Risks. Advancements in generative language model- ing, including the techniques explored in this paper, also carry inherent risks that warrant careful consideration.
11

---

Economic Disruption: As generative Al tools become more capable, there is potential for economic disruption in professions that involve content creation or communication. While these tools can also augment human capabilities, the long-term societal adjustments need to be considered.
Security Concerns:
The potential for generating human-like text could be used in social engineering attacks, to create more convincing phishing emails, or to automate abusive online presence.
Considerations for Responsible Development. While this work is primarily foundational, focusing on novel modeling paradigm, it is crucial for the research community to engage in ongoing discussions about the responsible development and deployment of such technologies. Future research building upon FarFlowLM should ideally incorporate investigations into:
Techniques for detecting synthetically generated text from such continuous-space flow models.
Methods for identifying and mitigating biases in the learned representations and generated outputs.
Developing frameworks for controllable and ethical AI, ensuring that generative capabilities are aligned with human values and safety.
Our aim is to contribute the open scientific exploration of generative models. We believe that by understanding the capabilities and limitations of new approaches like TarFlowLM, the community is better equipped to foresee potential impacts and work towards harnessing these technologies for societal benefit while mitigating potential harms.
References
[1] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, edi- tors, Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 17981-179999 2021. URL h http///poccee...e.eoe// neur sourips.cer tttpppccppppppeeee 1 2021 /hash// 958c530554478bccdeee 25b7 3d- Abstract. html.
[2] Jens Behrmann, Will Grathwohl, Ricky T. Q Chen, David Duvenaud. and förn-Henrik Jacobsen. Invertible residual networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 573-582. PMLR, 2019. URL http://proceedings. mlr www.mmll....ss..m.mmmmmmmmmme vttess//v7/bbhrrmaa 9a. html.
[3] Samuel R. Bowman, Luke Vilnis, Oriol Vinyals, Andrew Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. In Stefan Riezler and Yoav Goldberg, editors, Proceedings of the 20th SIGNLL Conference Computational Natural Language Learning, pages 10-21, Berlin, Germany, 2016. Association for Computational Linguistics. doi: 10.18653/v1/K16- 1002. URL htttps:/aaclanthollle o
[4] Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https: / //openreview.nee www.we..eeffooooeeeeeeee id= xsqj09Fm.
[5] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florinn Balcan,
12

---

and Hsuan-Tiee Lin, editors, Advances in Neural Information Processing Systems 33: An- nual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL. ptttp:://rroceediie wewpeeedddinnnneiiiiiee 1457c9d6bfcb49674 Seffbaac1-22f44a-bstrreeee html.
[6] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligianni- dis, and Arnaud Doucet. A continuous time framework for discrete denoising models. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Daniclle Belgrave, K. Cho, and A. Oh, editors. Advances in Neural Information Processing Systems 35: Annual Conference on Neural In- formation Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022. URL http: / hwtpppppppsssss////ee www.nnissspppaieeeeeee wwtpccprppppppeceeeeeeeeeeee htles//paper/222/haslll b5b528767aa35f
Cow.eettttcoooeeeneeeeeoeee html.
[7] Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Gener- ative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design, 2024. URL https://arxiv.. org. abs/ 04997.
[8] Nicola De Cao, Wilker Aziz, and Ivan Titov. Block neural autoregressive flow. In Amir Globerson and Ricardo Silva, editors, Proceedings of the Thirty-Fifth Conference on Uncer- lainty in Artificial Intelligence, UAI 2019, Tel Aviv, Israel, July 22-25, 2019, volume 115 of Proceedings of Machine Learning Research, pages 1263-1273. AUAI Press, 2019. URL html.
http:// http://proceedinns mll wttppppss/vvl/ddeeeeeee
[9] Mathieu Chalvidal, Matthew Ricci, Rufin VanRullen, and Thomas Serre. Go with the flow: Adaptive control for neural odes. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenRevvew.net, 2021. URL https: / //openreview. www/feruuuune met/foouumid=giii t4HdDNa.
[10] Tong Che, Yanran Li, Ruixiang Zhang, R Devon Hjelm, Wenjie Li, Yangqiu Song, and Yoshua Bengio. Maximum-likelihon augmented discrete generative adversarial networks. ArXiv preprint, abs/1772.00983, 2017. URL httpp:://rxiv.. wwtp:/iii.oggabb///// 100. 707217798......
[11] Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garneut, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information P Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 6583, 2018. URL https: /ppooeedd/sssoes html.
neurips. httppsscccppppprreeppppp 201 httphhahhhh//fhfee df
[12] Ting Chen, Ruixiang Zhang, and Geoffrey E. Hinton. Analog bits: Generating discrete data using diffusion models with colffcondiionnng. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
httl:/tps:://ot /ttttpps///ppeeevieeee net/pdffid-33 Sesitj9900 JRRQQFW.
[13] Cyprien de Masson d' Autume, Shakir Mohamed, Mihacla Rosca, and Jack W. Rae. Train- ing language gans from scratch. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelz- imer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 4302 4313, 2019. URL https: / /ttp//ppedonns..o/// www.ineusspp.. httpcc/aep///////aaaaaaaaae html.
[14] Yuntian Deng. Anton Bakhtin, Myle Ott, Arthur Szlam, and Mare' Aurelio Ranzato. Residual energy-based models for text generation. In 8th International Conference on Learning Repre- sentations, ICLR 2020. Addis Ababa, Ethiopia. April 26-30, 2020. OpenReview.net, 2020. URL https://openreviis wwweeeewwweeweeeeuueeeeeee 1d5g11GGHHKHHH
[15] Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, Curtis Hawthorne, Rémi Leblond, Will Grathwohl, and Jonas Adler. Continuous diffusion for categor- otgg/bs//21/ 15089.
ical data, 2022. URL
[16] Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. International Conference on Learning Representations workshop Track, 2014.
[17] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France,
13

---

April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https:// // openreview.net/formmm// hett/forumm1d--Hkpbnngl..
[18] Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural odes. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 14. 2019, Vancouver, BC, Canada, pages 3134-3144, 2019. URL https::/proceeeeee neur iirrips.cco cc/ wwtppcppp/pppccceeccccccceeeeeet / 221///2119///// htthahh//bbe99aad4eaFFee 21we222bbbbfff1155 (444449555999111111 1eeaaadddd1111111111111111 c8 html.
[19] Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. ArXiv preprint, hos/2407.15595 2024. URL https://arxiv. http://..rrg/b////4//00 15595.
[20] Mathieu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. MADE: masked autoen- coder for distribution estimation. In Francis R. Bach and David M. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille. France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference P Proceedings, pages 881-889, JMLR.org. 2015. URL http: / / http///roeeennnn...e.. mlr. www.mm..../s..e.e.mmmmmmm wttp///wt//emm.n.... hemmm
[21] Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Pad and Gabriel Syn- naeve. Better & faster large language models via multi-token prediction. ArXiv preprint, abs/2404.19737, 2024. URL https://arxiv... https://arrv..rrr/ab//2222 19737.
[22] Aaron Gokaslan and Vanya Cohen. OpenWebText Corpus. https://Skylionnee
thub. Accessed: 2024-05-14. An open-source recreation of the
ha/OpenWebTexttorppss 2019. WebText dataset.
[23] Ian J. Goodfellow, Jean Pouget- Mehdi Mirza, Bing Xu, David Warde-Farley, Sher- jil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In Zoubin Ghahramani, Max Welling, Corinna Cortes, Neil D. Lawrence, and Kilian Q. Weinberger, editors, Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, wttpsssesaiiicscccppeee/e//ee cc/pappr/2014 Canada, pages 2672-26888 2680, 2014. URL Abstract. html.
hash/5ca3e9b1225 1
[24] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: free-form continuous dynamics for scalable reversible generative models. In 7th International Conference on Learning Representattions, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL httpp::///penrreeee wwtp:/iwwwwe//foonm/eeee netp:formmiddrrxxknneee
[25] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks. ArXiv preprint, abs/2308.07037, 2023. URL https: http://rrriv...gg////////// 07037.
[26] Jiatao Gu, James Bradbury, Caiming Xiong, Victor O. K. Li, and Richard Socher. Non- autoregressive neural machine translation. In 6th International Conference on Learning Rep- resentations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018 URL htttttt///peenreiiieee www.eeewweeeeeeeeeeeeeeeeeeeeeee httt/forumidd=B111
[27] Ishaan Guirajani and Tatsunori B. Hashimoto. Likelihood based diffusion language models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http: / /ttp::///aaaa papprs.pppessppp/ppp nips. cc/ wwppppppp_..peeeeeeeeeeeeee watppaaaiaeaappaaaaaaaaaaeeellle (illss/paer/2200 hash/35b5c1 7 75e1 1755139bfff55 htw.:/225561770fff---attt-rrrr
html.
[28] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Puttern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016. pages 770-778. IEEE Computer Society, 2016. doi: 10. 10.10C/R....
URL https: / /doii org/18 /0. 11 1109/CVPP 2016.90.
[29] Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow-based generative models with variational dequantization and architecture design. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML. 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 2722-2730. PMLR, 2019. URL http: pttpp/edocccc...cp/ee/p/p mlr, wttpsss//////h/// html.
14

---

[30] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Con- ference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL. https://proceedinng wwwiggsnneeriiii neurips. wttpcccpee/2////eeeeee/e //20// #cbbcfcc85Marf0067711bb 119aaa1001c94bbbbbbatttttt html.
[31] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. In Marc' Aurelio Ran- zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NewrIPS 2021, December 6-14, 2021, virtual, pages 12454-12465, 2021. URL https: http::::/pprceeeeeee http///oooeennn....e neur cc/paper/2222 h82//hashh Abstreectttttttttttt.m... html.
[32] Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron C. Courville. Neural au- toregressive flows. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML. 2018, Stockholmsmüssann Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 2083-2092. PMLR, 2018. URL http: httptt::/pprcee http:///rreeeeng.....c mlr caasss//v//huuu 18d. .ttll.
[33] Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simun
and Computation, 18(3):1059-
1076, 1989.
[34] Minguk Kang, Jun- Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In IEEE/CVF Conference on
Computer Vision and Pattern Recognition, CVPR 2023, Vancouver BC, Canada, June 17- URL
24, 2023, pages 10124-10133 10134. IEEE, 2023. doi: 10. 0W09CCVPR5222..222..0008 https://doi... org/10. /0/ 1109/CVPR52729 2023. 2023.00976.
[35] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In IEEE Conference Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pages 4401-4411 Computer Vision Foundation / IEEE, 2019. doi: 10. http:/VVR.2019.0000 content_CVPR 11/PR__019/htmm1mm /Karras /ttp A_Style-Based Baaeyyle-BBaedd_eeeerattttrr
URL http: http:::///pecacess... thecvf. com/
tecture_for.
Generative_Adversseit http://iwe..eeeerrrraaaeeeerrrro wwtpoorkss_cc__ee___ 201PR_2219__p__ -aepppppp_____ html.
[36] Dongjun Kim, Byeonghu Na, Se Jung Kwon, Dongsoo Lee, Wanmo Kang, and Il-Chul Moon. Maximum likelihood training of implicit nonlinear diffusion model. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Ad- vances in Neural Information Processing Systems 35: Annual Conference on Neural Infor- mation Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022. URL http: wttp://ppppppppp.pp...pppp wttpppp.ccppppp..pppppp filess/aper///222 html.
[37] Diederik P. Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 con- volutions. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicold Cesa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, Decem- ber 3-8, 2018, Montréal, Canada, pages (0236-10245, 2018. URL https: http:eeei////oreiinngge/. neurips.. cc/ cc/paper////e///c/e//eeeeee 201 2018///1 1atash//ahhaaa
d1 (111db323622000333 cc (777777777911 3200-Abstract. html.
[38] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun, editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014. Conference Track Proceedings, 2014. URL
http:/ /ttp:////xxv.///// www.oogg/////oeeeooooo 1312. 6114.
[39] 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, 29, 2016.
[40] Xiang Li, John Thickstun, Ishaan Guirajani, Percy Liang, and Tatsunori B. Hashimoto. Diffusion-Im improves controllable text generation. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances Neural In- formation Processing Systems 35: Annual Conference Neural Information Process-
15

---

ing Systems 2022. NeurlPS 2022, New Orleans, LA. USA, November 28 Decem- ber 9, 2022, 2022. URL http://papers. papppppppreeppppppppppaapeepeeee htppppssc/epppeeeeeeeeeeeeeeeere ftw.ss/paper///222/aaaa
html.
[41] Phillip Lippe and Efstratios Gavves. Categorical normalizing flows via continuous transforma- tions. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL -GLNZeVDuik.
[42] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. ArXiv preprint, abs/2412.19437. 2024. URL atttps://arxive www.xiivv.orr....o..... org httgg/bb///21// 19437.
[43] Guan-Horng Liu, Tianrong Chen, and Evangelos A. Theodorou. Second-order neural ODE optimizer. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Anmual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages paper/
2021. URL https: / pttp///roccens...//e/ neurips.cc/ (etaa772d0306666bb33FFo Abs Abstract. .tmll.
[44] Sulin Liu, Juno Nam, Andrew Campbell, Hannes Stärk, Yilun Xu, Tommi Jaakkola, and Rafael Hoanez-Bombarelli Think while you generate: Diserete diffusion with planned denoising. ArXiv preprint, abs/2410.06264, 2024. URL htttps://arxtt wwwivi.org//bbboe org/abs/ /taps///.....m 06264.
[45] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimat- ing the ratios of the data distribution. In Forty-first International Conference on Machine Learning, ICMI, 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL
https://openre
met/forrmmid==CicRRRI
[46] Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Maximum likelihood training for score-based diffusion odes by high order denoising score matching. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 14429-14460. PMLR, 2022. URL https: / http://ppooeeeenn...ee//oo mml. ptt...pree//v11//12222 html.
[47] Matt Mahoney. text8 Corpus. http: htt::/matttat http://ometmmmmm...m/mmmmmo wwt. www./////eet/ttttetaeeeaae html, 2011. Ac- cessed: 2024-05-14. Original site for the text8 dataset, part of the Large Text Compression Benchmark.
[48] Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score match- ing: Generalized score matching for discrete data. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural In- formation Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 Decem- ber 9, 2022, 2022. URL http://papers: wwwp.ccpppppp.....eectpppppppppe wwtpppppeeeille/pppeeeeeeeeeeeee R1Ma336607e8994554eebblf.
html.
[49] Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan
Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data.
ArXiv preprint, abs/2406. 03736, 2024. URL https:
arrgabs/2406. 03736.
[50] George Papamakarios, Iain Murray, and Theo Pavlakou. Masked autoregressive flow for density estimation. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Anmual Conference on Neu- ral Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 2338-2347, 2017. URL https: http::///ooceeeeeeee neur p2ccpppppr/217/haaaa html.
[51] George Papamakarios, Eric T. Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. J. Mach. Learn. Res., 22:57:1-57:64, 2021. URL http://jmlr wrg. http//ppprs/////////p//
10-1999.....000 html.
[52] Massimiliano Patacchiola, Aliaksandra Shysheya, Katja Hofmann, and Richard E Turner. Transformer neural autoregressive flows. ArXiv preprint, abs/2401.01s 00885,, 2024. URL https:
//arxiv.org/abs. argg 01855.
16

---

[53]
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. In The Twelfth International Conference Learning Representations, 2024.
[54] Lev Semenovich Pontryagin. Mathematical theory of optimal processes. Routledge, 2018.
[55] Murray Rosenblatt, Remarks on a multivariate transformation. The annals of mathematical statistics, 23(3):470-472, 1952.
[56] Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T. Chiu, Alexander M. Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. ArXiv preprint, abs/2406.07524,, 2024, URL abs/2406. 07524.
[57] Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aäron van den Oord. Step-unrolled denoising autoencoders for text generation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL
[58] Jiaxin Shi, Kehang Han, Zhe Wang. Arnaud Doucet, and Michalis K. Titsias. Simplified and generalized masked diffusion for discrete data. ArXiv preprint, 0os/2406.04329) 2024. URL
https://arxiv. wwwpaaiwww....raaaaeeaeaeeooeere abs/abb///99 04329..
[59] Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference on any-order autore- gressive models the right way. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA. November 28- December 9, 2022, 2022. URL http://t//// /ttppr///ppppppppp/eeeeeeeeeessp nips cc/paper_fii _files/ htte///////aaaa1/232a2aaee 23 0ew.a001199000000000000000033330 -wwp::/tttddaaattt-t---t------a- CelbssrrcttCoofeeeeeee
html.
[60] Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep un- supervised learning using nonequilibrium thermodynamics. In Francis R. Bach and David M. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning. ICML 2015, Lille, France, 6-11 I July 2015, volume 37 of JMLR Workshop and Conference Pro- ceedings, pages 2256-2265. JoLR. 2015. URL http::ttpp////e http:::/prrocceeingg.. wwa.ggg.m...rrrr.. milrppress/v37 sohl-dickstein15 5. html.
[61] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Mare Aurelio Ranzato, Alina Beygelzimer. Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Process- ing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 1415-1428, 2021. URL h /ppoce///o neurips. cttpsccccpppppee 2021
2021 / /ttp////....TT.cc/77c0666666226.
html.
[62] Yang Song, Jascha
Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representationss ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL httpttt////oeeereeieeee net/ pet/forummidd=xXIGG11
[63] Jianlin Su and Guang Wu. f-vaes: Improve vaes with conditional flows. ArXiv preprint, http://aawxv...orabsbbass/
abs/1 809.05861, 2018. URL
[64] Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous- time discrete diffusion models. In The Eleventh International Conference on Learning Rep- resentations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL settpdfF2ddBBWWWWW/Y
https://openreviier
[65] Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, V. N. Vishwanathan, and Roman Garnett, editors, Advances Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 1600-6315., 2017. URL httpss////oooceeeiiingeg neurips. cc/ 1aaapper/211//ash/aaaa hhtash//hhaaaaaaa11// 1eaaaaeeaaaaa00eee0000000000---0
html.
[66] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you nnips, 2017. Advances in neural information processing systems, 30, 2017.
17

---

[67] Clément Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. In The Eleventh Interna- tional Conference on Learning Reprexentations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023.
OpenReview.net, 2023. URI. https: /ttp:/:////ppeeeieeeee net/pdf wwtp:nttt//pfddduaAa---a--- -dd-aAA-Nu866w.
[68] Guanghan Wang, Fei Yang. Yuchen Li, et al. Remasking discrete diffusion models with inference-time scaling. arXiv preprint, abs/2503.00307, 2025. URL https://arxiv.org abs/2503 Nbs/2503.00307.
[69] Maciej Wiatrak, Stefano V Albrecht, and Andrew Nystrom. Stabilizzing generative adversarial networks: A survey. ArXiv preprint, abs/1910.00927, 2019. URL 1910.00927.
[70] Minkai Xu, Tomas Geffner, Karsten Kreis, Weili Nie, Yilun Xu, Jure Leskovec, Stefano Ermon, and Arash Vahdat. Energy-based diffusion language models for text generation. ArXiv preprint, abs/2410.21357,, 2024. URL https: http://aawiv..ooa////////// .23357.
[71] Minkai Xu, Zhaoyue Huang, Jiaming Song, Stefano Ermon, et al. Energy-based diffusion language models for text generation. arXiv preprint, abs/2410.21357, 2024. URL https:
/ //arxiv.org' www.:///g/bb///// 2410 1.4.0... 21357.
[72] Lantao Yu, Weinan Zhang, Jun Wang. and Yong Yu. Seqgan: Sequence generative adversarial nets with policy gradient. In Satinder P. Singh and Shaul Markovitch, editors, Proceedings of the Thirty-Firee AAAI Conference Artificial Intelligence, February 4-9, 2017, San Francisco. California, USA, pages 2852-2858. AAAI Press, 2017. URL http: //aaii. org/ /ss///s/iii/ /ndex. CTD/AAATT/AAAT1// wwwpp/aaa/////////aaa/p /view/vewwweee/e 14344.
[73] Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing flows are capable generative models. ArXiv preprint, abs/2412.06329, 2024. URL ht
/ www.i..org/bb//242222
[74] Ruixiang Zhang, Masanori Koyama, and Katsuhiko Ishiguro. Learning structured latent factors from dependent data:a generative model framework from cntommtiion-thhorcee perspective. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 11141-11152. PMLR, 2020. URL htt:: pttt:://proceedingss mlr pww.ss/ell 19/ zhang20m. html.
[75] Ruixiang Zhang, Shuangfei Zhai, Etai Littwin, and Joshua M. Susskind. Learning representation from neural fisher kernel with low-rank approximation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https: wwtphht:t///pereeeeeeee /opprrevieneeeeeeeeeeeeeeee wwwiiie..eeeeeoeoeee met/Forum2id=J1rhaNcc..
[76] Ruixiang Zhang, Tong Che, Boris Ivanovic, Renhao Wang. Marco Pavone, Yoshua Bengio, and Liam Paull. Robust and controllable object-centric learning through energy- models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https::/ http:/enwweiiieeeeeeede/ee...... wcNtbEtcGIC.
[77] Ruixiang Zhang, Shuangfei Zhai, Yizhe Zhang, James Thornoon, Zijing Ou, Joshua Susskind, and Navdeep Jaitly. Target concrete score matching: holistic framework for discrete diffusion. ArXiv preprini, abs/2504. 16431, 2025. URL nttpps://arxtt wwtpxiiv..rr///oooooooooo org/ //tg//b///// 2 16431.
[78] Zijun Zhang, Ruixiang Zhang, Zongpeng Li, Yoshua Bengio, and Liam Paull. Perceptual generative autoencoders. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 11298-11306 PMLR, 2020. URL http::/ press/v119/2hang20ab html.
mlr.
[79] Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. ArXiv preprint, abs/2302.057377, 2023. URL https://arxxv.orrg abs/2302.05731
[80] Zangwei Zheng, Xiangyu Peng, Tianji Yang, Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all, 2024. URL https:://// //g: thub. wwwp:/ppcaaieeccooeeeeeeeeeeeeee
18

---

[81] Juntang Zhuang, Nicha C. Dvorekk, Sekhar Tatikonda, and James S. Duncan. MALI: A memory efficient and reverse accurate integrator for neural odes. 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL httpss///openeeeiieee wwt.:eernunnnee feet/formmiid=blSSjie Fej
[82] Zachary M. Ziegler and Alexander M. Rush. Latent normalizing flows for discrete sequences. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA. volume 97 of Proceedings of Machine Learning Research, pages 7673-7682. PMLR, 2019. URL http: pttp///oeceen......////////../.p mlr wwwpss//9//ziiegerr1999 9a, html. html.
19

---

Appendix
Table of Contents
A Details for 1-D Mixture-CDI Flow Layer
20
A.1 Proof of Proposition 1
22
B Details for d-D Mixture-Rosenblatt Flow Layer
22
Proof Proposition
25
B.2 Proof of Proposition
................................
27
C Detailed Derivation of a Flow Layer's Learning Objective
28
D 1-D Mixture-of-Gausian and Mixture-of-L.ogistiess Coupling
30
Additional Experimental Details
33
F Connection Discrete Autoregressive Language Models
34
G Latent Space Evolution Metrics
37
H FLOPs Calculation Details
39
Related Works
41
A Details for 1-D Mixture-CDF Flow Layer
This section details the derivation for the ID Mixture of Gaussians CDF (Cumulative Distribution Function) normalizing flow layer, as introduced in Section 3.2.1 of the main text. The conditional mixture density for a single scalar component 24, given its context (Z<e. Zt,<i) is defined in Eq. equa- tion 5 as:
M22 21,t1) Σ π (2,,) [[), (<l))
where the parameter bundle (T, m, o²| (containing V mixture weights, means, and variances respectively) is predicted by Transformer model from the context (Zzr, Zt,<s). We denote this ID mixture probability density function (PDF) as Pmix 1 (2 Rec, 2t,<;). This conditional PDF can be used construct an invertible transformation suitable for normalizing flows.
The CDF of Pmix- (2e,al zet ci) denoted F (Ze,s; Zet. 2, cx), is a weighted sum of the CDFs of the individual Gaussian components:
2r, m[k]
Fmix.1 (2 2r
π
[k]
[c]
where ((-) is the CDF of the standard normal distribution (o0,, 1), and [kk] = Vo is the standard deviation of the k-th component. The mixture parameters [k], m [k], [kk] are functions of (Z<t, Zi)
Using the probability integral transform, if random variable Z,,; is drawn from Pmnix (Z,,izz jict, 2t,4)), then Fmix (Z,,; net, 2t,cs) is uniformly distributed on [0, 1]. To map 21, to variable u4, that follows standard normal distribution (0, 1), compose the mixture CDF with the inverse standard normal CDF, x- --)(-);
Vt Φ (Fmix-1 (24,)-1(24,,,2 Zets (t))
(10)
20

---

This transformation is invertible. The variable 24,4 can be recovered from ue using 2t,i Fmis (f(u,,,): tiee, k,,,),, where - is the inverse CDF (quantile function) of the mixture. F F generally does not have a closed- foreedfoom expression and may require numerical methods for its evalua- tion.
For this transformation serve normalizing flow layer, its Jacobian determinant is needed. In this ID case, this is the absolute value of the derivative ōz,,"
0us. We compute this using the chain rule.
Let Yedr
F (21,6 Pict, 2t,ci)- Then ut,i + (b= and
Oua, d4 (Wear)
O (Zr,5) Zect Zt,ei)
Zzg dyeat
Ōzz,i
Mear Fan
We evaluate each term:
1. The derivative of the inverse standard normal CDF,
[Manl) If ussi = ( dyar
then Year = =(uu.)- Differentiating Year = (u,,,) with respect to u4,i yields dyat duc,
N (u,,, 1), where N (u,,;; 0, 1) is the PDF of the standard normal distribution. Therefore, day. dyar dgar
(Vea)
(0,1)
2. The derivative of the mixture CDF with respect to 2r,i,
F (Rri 2i (L..c))
82a,2
FF (2,35 Zett Zi.c
i
2 m [c]
[ki]
Ōzi,
Ōzi,
[k]
m[k]]
Σ
T [kii
:0,
o
[k]
K
Recognizing that N
[2, (a.--m/k), [x
00 :0
is the PDF (2 [k), a ( the sum be-
comes:
Fmix. (ze, 2i)
(kk) = Pmix. 1 /acc 2s)
02g,
Thus, the derivative the mixture CDF is its PDF.
Combining these results, the derivative of the transformation is:
dug, Pmix-1 (2), /Zzc, Zi)
1zg,
Nu,,, 0, 1)
The log absolute Jacobian determinant is therefore:
Jugs
log
log Pmix 1 (zt,a Mec. m.)e) log N N(u,s; 0)
(11)
This equation relates the log-density of 2τ, under the mixture model Pmix-1 to the log-density of ug under the standard normal base distribution and the log-Jacobian term. Rearranging Eq. equation 11 according to the change of variables formula (log p(z) log p(u) + log det Jyl):
Our,
log Pmix-1 Rice, Zi, = log N (Mr, ;: 0, 1) + log 0z,
This confirms that learning the mixture parameters (r, m, 2 is equivalent to learning this ID Mixture- CDF flow layer.
When such a transformation is used as layer l in stack of flows, øt,à corresponds to the i-th scalar component the input to this layer, 11 (e-1) The output ut,, corresponds to the i-th scalar component (.(() The parameters for the mixture density Pmix-1 (h (f-1) ((o((())ccoo(()) for layer are determined by neural network conditioned on cond!") which is the layer-specific conditioning information. The term log Pmix-1 (f(-- (cond()) in the log-determinant calculation (Eq. equation 11, adapted for layer f) represents the log- probabbbayl of the layer's input under the specific 1D mixture model implemented by that layer.
21

---

A.1 Proof of Proposition 1
Proposition 1. Let P=ix-- Ex= (2; mg, o be ID MoG probability density function (PDF), and let Fmix-1(2) be its corresponding cumulative distribution function (CDF). The trans- formation R + E defined by u = f()) -- ((Fmx-((22)), where & is the standard normal CDE, is an exact normalizing flow between the density pm) and the standard normal density N(;0,, 1).
Proof. The proof demonstrates the equivalence from both forward and inverse directions. A key property is that (2) is strictly increasing, since its derivative, Fiix-- Pmix-1 (2), is a sum of strictly positive Gaussian densities and is therefore strictly positive. Thus, Fiix-11 is bijection from R to (0, 1).
Forward Direction (Z + U). Assume random variable ~ Pmix-1(2) We show that the transformed variable U ((2) follows standard normal distribution by computing its CDF:
P(U u = (( (Fmix-1 s u
P(Fmix-1 17) - (())
(since + is strictly increasing)
Fmix-1 (F 1
(by definition, P(Z S z') -Fuix-(2z)))
(()).
The CDF of U is (u)), which is the CDF of the standard normal distribution. Hence, U ~ NN(0 1).
Inverse Direction (U + Z). Assume random variable U (0, We show that Z (U) = F (T{ is distributed according to P by computing its CDF:
P(Z 2) ( ( (Q))) (0L)) 2)
(since Fiix-1 1 is strictly increasing). By the probability integral transform, the random variable o(U) is uniformly distributed on [0, 1].
P() S F 1
Therefore, P P(T(U) S v) for any v [0, 1]. Since Fmix-1 (2) E (0, 1), have:
P(z S 2) Fm).-1(2).
The CDF of Z is F Differentiating with respect to z yields the PDF of Z as pz(2) dg 1 = Pmix (2), which completes the proof.
B Details for d-D Mixture-Rosenblatt Flow Layer
We use C to denote the conditioning context variable. A d-dimensional conditional mixture density, such Puix (z|C) = ΣV Ek=1 & (C))(zz m2 (c(C,, (C) be realized as exact and invertible normalizing flow layer. This construction is used for modeling complex conditional distributions within the flow framework. For example, can represent the prior's conditional distribution p(z, (Z<))
as defined in Eq. equation 7, or general conditional density P Punx, (h! ate layer& l in stack of flows (see Section 3.3).
(f-1) ccodd()) for an intermedi-
In this description, E Rd denotes the input variable to this flow layer. If this layer is the first transformation applied to latent variable 2, from the VAE, then corresponds to z, (which also h in the notation of stacked flows). If it is an intermediate layer é in the stack, z corresponds to h (e-1). The term C represents the conditioning information, such as Zet for the prior, or layer-specific context co The parameters my and og of the Gaussian components can either be predicted based on C. or, as in Eq. equation they can be part of shared codebook (c.g.. {At, o2 where only the mixture weights T (C) are predicted.
The goal is to define an invertible transformation ga Rd + RRd denoted u = (z; C), such that the output variable u E Rd follows standard d-dimensional Gaussian distribution, i.e., u ~ NO, Ia). This transformation is constructed by processing the components of z = (2), =²) sequentially, from i to using a method related to the Rosenblatt transformation. The forward pass detailed in Algorithm I.
22

---

Forward Transformation z + u: The transformation from z to u is built sequentially, processing one component time. At each step u, is computed from 2; using C and the previously processed components Zzzi = (21, ( Mixture component weights are maintained and updated throughout this process.
The procedure begins with initialization: for the first component 21, the initial mixture weights are set by the parameters of Pmix Pmix(z(().
( (C) (C) for 1: =::=...... I,....VV. V. a
(12)
For each dimension == 1,.... d, the following steps are performed: The forward transformation proceeds as follows. For each dimension ,,....d: d:
Conditional Density of 2₁. The conditional probability density of 2 given and Zei ID mixture of Gaussians, denoted ( CC, 2.x)- This is computed by marginalizing over the mixture components using the current weights 0k (8)
0Kk ①
((r me, RC())
(13)
Here, (C) is the i-th element of the mean vector ms(C) for component k, and (C(C) is the (isotropic) variance for component k. By construction, ( [z, equivalent p(2: Zci).
Conditional CDF of 24. The corresponding ID conditional cumulative distribution function (CDF) for 24 is
F,(2(,,,,e) I Puixx i (EIC, Zei)dE
( ak (C, (* 2: #=-me:(C)
(14)
(c(C)
k=1
where ((-) is the CDF of the standard normal distribution and s; (C)
Transformation to ug The i-th component u, of the transformed variable is obtained by applying the probability integral transform to 2; using F; followed by the inverse CDF of the standard normal:
u4 = (F,(2,C,,
(15)
2; is drawn from i [C, zes), then [2, z<i) is uniformly distributed on [o, 1]. and thus u; is standard normal. The sequential conditioning ensures that Ui...., tla are mutually independent and standard normally distributed.
Update Mixture Weights (if < d). After processing 2; and obtaining if is not the last dimension, the mixture weights for the next dimension, oi (i+1). are updated using Bayes' rule to incorporate the information from
(6+1) 25i) = p(component k(C) 0
(cmpponnttntt k; C, zci) -((component kt/,, p(2g C,cce)
(e, ()( 0 (i)
(C, Zes Pmnx, (2,(C, 2zz)
(16)
where Zei = (21,...,,;))
2.). These updated weights are then used for the next dimension.
Inverse Transformation u + z: The transformation is invertible. To compute the inverse, z 8 C); we proceed as follows, iterating from = to d:
Initialiaatinn (for 1): Set the initial mixture weights a (C) in Eq, equation 12.
d): For each dimension given u; and the current weights (8) (C. z21) (where Z have been determined in previous steps), solve for 2, by finding the value 0
Sequential Inversion (for 1,...,d):
that satisfies
Flzz(C, zzs) 4(u))
(17)
23

---

where F; is defined in Eq, equation 14. Since F; is strietly monotonic function (as the CDF of continuous distribution with positive density), a unique solution for 2; exists. This inversion typically requires numerical ID root-finding algorithm.
After determining 2,, if < d, update the mixture weights to 0 (i+1) (C, zgi) using Eq. equation 16, now with the newly found 24. This process is repeated sequentially for each dimension until all components of z recovered.
Jacobian Determinant: The Jacobian matrix Jg. (z) of the forward transformation = ga (z( C) has entries (Jg.) ): ãz, (and C), the Jacobian matrix is lower-triangular. Its determinant the product of its diagonal entries: det Jns (z) II 05.
Bug, Due to the sequential construction, where uy depends only on ž1, x
au
To find 0u. differentiate Eq. equation 15, its form ((u)) Fz (.z).. Differentiating both sides with respect 2, (treating C and Z< constant for this partial derivative):
(()(,,) ((F
(18)
02:
Using the chain rule on the left side yields (u,) dus where o(.) is the PDF of (0,, 1). The right U=."
side the derivative of CDF with respect its variable, which its PDF: 0z ( z<i). Thus,
Ou: imix,, (z: nes)
(us)) Oz:
(19)
Since p(u)) > 0, we can write:
Pmix. (=g xcce)
(20)
02; p(ui)
The log absolute Jacobian determinant is then:
Su.
log |det JRn (z)l =
Σ
log
Σ (log Pmix, (zi z2s) - loog of (u:)
(21)
By the chain rule of probability, the log-density of the input variable z under the conditional mixture Pmix (z|C) is:
d
log log Punix (z|C) Pnnix(/(C)
[I (I22IC,,
II
Σ log logp(=;(C,2c)
(22)
In our construction, p(z,(C, m2ee) is precisely Pmix, (=,(C, from Eq. equation 13. So,
- Pmix, zz;) = log Pmix (z|C). The log-density of the transformed variable u under the standard d-dimensional Gaussian base distribution Phase (u) = A (uu 0, Ia) is log Prase (u)
log 0, 1) Σi= log Φ(u)). Substituting these into Eq. equation 21, we arrive at the
expression for the fog-determinant in normalizing flows:
log |dct Jg. (z) = log Pinix (z(C) - log (u: 0, Ia)
(23)
This confirms that the sequential transformation ga correctly implements the desired conditional mixture density Pmix (z(() (C) as normalizing flow layer. When such a layer is used, for example as layer ((() in a stack, it transforms an input h. (f-1) (which plays the role of z in this derivation) to an output h (which plays the role of u). The parameters defining the mixture (initial weights
Th (condd ()), means my (cond(")) ), and variances & (cond) are predicted by neural network based on conditioning context cond!
24

---

B.1 Proof of Proposition
Here we provide formal proofs for the normalizing flow layers based on the Mixture of Gaussians (MoG) CDF, as described in Appendices A and B. We establish that these transformations exact, invertible, and correctly model the target mixture densities.
The core of the d-dimensional conditional probability distribution p(a, kzs)) for each dimension. The following lemma formally derives this distribution and shows that it takes the form of ID Mixture of Gaussians, which is analytically tractable.
flow lies in sequentially computing the true
Lemma 1 (Conditional Distribution of an Isotropic MoG), Let I& a random vector whose probability density function is mixture of V isotropic Gaussians:
p(z) = TNN(2; mx LaT.)
k=1
O. Then for any dimension E (1,...,df, ,d], the conditional probability density of its component given the preceding components Zei (Zi,..., 2i-1) is also 1D Mixture of Gaussians, given by:
where Σx Tk 1 1,
Z<)) 2<)) = Σ (0)
ak
k=
(4) where ma the i-th element of mg, and the mixture weights dg (zez) are the posterior probabilities of component membership given Zez:
( 0 (Zez) 2 P(K = k/zce)
(1) (For the base case i 1, Z<1 is empty and ag
Tk).
Proof. We derive the conditional distribution p(= <<)) introducing latent categorical variable
K V} that indicates which Gaussian component is drawn from. The hierarchical model
is:
1. Sample component index K = k with probability P(K = k)
2. Sample from the chosen component: p(z) k) =N(z: my La..).
The target conditional density can found by marginalizing over this latent variable using the law of total probability:
V p(z: nes) = M(sliees K =)PPPK =lZzz))
(24)
k=1
We now analyze the two terms in the summation separately.
1. The Component Density: p(z/zze,, K= under the condition that know the sample originates from component k. For given component z follows a single multivariate Gaussian distribution N N(mE, diagonal, which implies that all dimensions of z are mutually independent #when conditioned K = b* Therefore, the value of is independent of the preceding values Zzc
-K=b) k:). This term is the conditional density of 2; given Zgi-
The covariance matrix L is
p( Zei K k) p(= k).
The distribution p(=/K- 1c) is the i-th marginal of the multivariate Gaussian N(z; my. La).). which is ID Gaussian:
p(aKK k) = -( o.)
25

---

2. The Mixture Weights: P(K = k klzes). This term represents the posterior probability of being in
(+) component after observing the first dimensions. We denote this by a (ze:) and compute using Bayes' rule:
P(K = klzz2) p(zzi|K = k)p(K=MM
pizci K = k)P(K = Σソ=1 - j)P(K=
p(zci)
The required terms are:
P(K = k) Ty (the prior probability of selecting component k).
p(z<a K = k) is the marginal density of the first i - dimensions of the k-th Gaussian. Due to the isotropic structure, this (Zzz mi o where ma,cs contains the first i - elements of my.
Substituting these into the formula gives the expression for the weights as stated in the lemma.
Conclusion. Substituting the derived component density and mixture weights back into Eq. equa- tion 24, we obtain:
p(a) =
Σ
(zes)
(22 o.
Mixture Weight
ID Gaussian Component
This confirms that the true conditional distribution is a ID Mixture of Gaussians with analytically computable parameters.
With the analytical form of the conditional distribution established, we can now state and prove the proposition for the d-dimensional flow in a more direct manner.
Proposition 2. Let Pmix (z) = =K- TeN(2; ma: Ta be d-dimensional isotropic MoG PDE
Let the transformation g Rdd - R2d be defined by the sequential application for 1,....d: dd:
u, =
where Fi(zi(zes) is the cumulative distribution function (CDF) of the ID MoG conditional density p(zs)lie)) as derived in Lemma 1. This transformation is an exact normalizing flow belween Proix and the standard d-dimensional normal distribution N(u; 0, Ia).
Proof. The proof relies on the chain rule of probability, Pmix (z) = H p(zz (ecz), and the result of Lemma 1.
Forward Direction (Z - U). Assume a random vector Z ~ (m). We show by induction that the components of the transformed vector U (Z) are independent and identically distributed as (O 1). By Lemma 1, (Z Zez) is the true conditional CDF of Zí. By the conditional probability integral transform, the random variable W; (2, Zex) is uniformly distributed on [0. and is independent of the conditioning random vector Zej Consequently, U² = + -(W,) is distributed as (0, 1) and is also independent of Zei Since (U1,... U,-1) is invertible function of Z U; is also independent of (U1,... "U,,...,,--. 1, By induction, all components of U are i.i.d. (0 1), U (0, I,).
Inverse Direction (U + Z). Assume a random vector ~~ (0, La). The inverse transformation
Z --g-((U) is computed by sequentially solving (Z, ((,)) for Z, This procedure is equivalent to ancestral sampling from the factorized distribution: for each i 1...., d, one samples Z, from the true conditional distribution p(2, Zzs) via inverse transform sampling. The joint density of the resulting vector is therefore I p( 22 (2s) (z)
Jacobian Determinant. The Jacobian of g is lower-triangular, as u, depends only on Zgi- Its determinant is the product of the diagonal entries du,
Treating Z constant, we have:
dz.
Ous (F(((i() 0z, dz,
F((2ala p(=s(zes)
Ó(u)) o(u))
26

---

where & is the standard normal PDF and we used the fact that the derivative of a CDF is its PDF. The log-determinant is:
Ou.
log dee.t Jz(2))
Σ
log
Σ
(logr(zikze)) - log ((u)).
0z;
Using the chain rule of probability for the first term and the definition of the multivariate normal for the second term, this simplifies to:
log det Jg(z)] = log Pmix log (u: 0, Lz).
This confirms that the change of variables formula is satisfied exactly. The recursive Bayesian update algorithm described in Appendix B is the practical implementation of the analytical derivation in Lemma 1.
B.2
Proof of Proposition 3
We now prove that the multidimensional transformation global diffeomorphism, which is a stronger condition than mere invertibility and differentiabiiltyy, ensuring the transformation is smooth and stable.
Proposition 3. Let Pmix (z) = - (2: mg. La be a d-dimensional MoG PDF. The Rosen- blatt transformation g IRd + R.d. defined sequentially by
U; = gz(z) = Φ- (Fa(z.. for pori=1...,,d
where Fg is the CDF of the true conditional density p(zi(zes)) is a global - disffeeoomoo.,hssm
Proof. Our proof strategy is to first show that g is a local diffeomorphism everywhere on Rd and then establish that it is also global bijection. These two properties together imply that is a global diffeomorphism.
Local Diffeomorphism. A map is local C di-difffcooomhiism at point if it is smooth (C*) in neighborhood of that point and its Jacobian determinant is non-zero that point.
(a) Smoothness of g: The map g is smooth if each component function u, = ge(z) is smooth. The function 9, - o F; is composition of functions. The conditional density p(2; (Zc)) is ratio of marginal densities of the MoG. Since marginals a MoG are themselves finite sums of Gaussian PDFs, they are smooth functions. As the denominator is strictly positive, p(z; ( is a smooth function of all its arguments, (2;, &ei). By differentiation under the integral sign, its CDF Fa (zc)) also smooth. Since 4-1 is smooth on its domain, the composition g, smooth. Thus, the overall map is smooth on Idd
(b) Non- Vanishing Jacobian Determinant: The Jacobian matrix Jg(z) is lower trianglia due to the sequential construction of g. Its determinant the product of its diagonal entries:
det Jg( = II Su:
d
p( p(za(zes)
II 1=1
1=1 Oz,
d(u))
Since both the conditional density p(=; (2ss) and the standard normal PDF ((u) are strictly positive everywhere, each diagonal term is strictly positive. Consequently, det Jg (z) > 0 for all z E Rd
By the Inverse Function Theorem, since g is smooth and its Jacobian determinant is non-zero everywhere, g local C deediiffeomorphiim at every point in Rdd
Global Bijectivity. We show that g is bijection by constructing a unique inverse z - -u) for any u E Rd. The inverse is found by solving b(u)) = F, z<a) for each 2; sequentially. For any i = L..., d, and for any fixed, previously determined context z the function (( is strictly increasing bijection from R - (0, 1). Thus, unique solution 2; F: (#(u,) (z exists each step. This inductive construction yields a unique z for any u, proving that g is global bijection.
A map that is a local diffeomorphism at every point and is also a global bijection is global diffeomorphism. Having established both properties for g, we conclude that diffeomorphism from Rd to ddd This confirms that the transformation is smooth, globally invertible with a smooth inverse, and possesses a well-defined, non-zero Jacobian determinant, making it robust and thcoretically sound building block for normalizing flows.
global C
D
27

---

C Detailed Derivation of a Flow Layer's Learning Objective
This appendix provides detailed, step-by-step derivation of the learning objective for single layer within the normalizing flow prior. The aim is to elucidate how the optimization process shapes the behavior of each transformation in the flow.
We consider generic invertible transformation f!() Rd + R representing a single layer at time step & in the flow. Let x = h,t-1) be the input to this layer and y = h, be its output, so y - (x; cond!"). The function x!² parameterized, and its parameters (and thus its behavior) are learned based on a conditioning context cond!" For brevity in the derivations that follow, we will often drop the explicit indices (t, E) and the conditioning context, writing f(x), (x) etc..
with the understanding that these are specific to a layer and its context.
The Log-Determinant of the Jacobian density of x to the probability density of y:
The change of variables formula relates the probability
px(x)= = py (f()) |det (x))
(25)
where px is the density of x, py is the density of y, and J (x) is the Jacobian matrix of f evaluated at X.
Normalizing flow layers like those discussed in Sections 3.2.1 and 3.2.2 (e.g., Mixture-of- or transformations based on d-dimensional mixtures) are constructed such that they define conditional model density for their input. Let Pmodet (x cond) be this parameterized density that the layer learns for its input x. The transformation f designed such that if x were drawn from Pmodel (x cond), its output y f(x would be drawn from a simpler, fixed target density, Ptarper (y) (typically a standard Gaussian). Substituting Px (x) Pmodal [x] cond) and py (y) Punga (y) into Eq, equation 25, we
get:
Pmodet (x(od) Pearga (x)) |det (x))
(26)
Rearranging and taking the logarithm gives the expression for the log absolute Jacobian determinant:
log |det. (x)) = llog Pmoda (x condd) - log Puarea (x))
(27)
This equation matches the form of equation in the main text. It shows that the log dettemeemmnnnt
term, crucial for computing the density of 21:2', is composed of the likeiiilloood of the input under the layer's 's learned model Pmodel and the log- likeliiiohood of the output y = f(x) under the fixed target density Prarget-
The Expected Log- Log-Jacobian in the ELBO The parameters of the entire model, including those defining Pmodet (x cond) for each layer, are learned by maximizing the Evidence Lower Bound (ELBO), as shown in Eq. equation 3. The ELBO includes the term ( . [log p(2).))) The terms and log determinant terms from
log-prior log p(Z):T) is given by a sum of each layer:
T
L
Σ
log Poase (u,) + Σ log det J 1.00 (hg (
log p(zi.) =
(28)
specific layer f p(0) This distribution Tin (xn(x) is induced by the input data X1.7, the encoder gf ( (x)r), and all preceding flow transformations fI)) f()-1), The ELBO maximization objective, with respect to the contribution
Let (x) denote the actual distribution of the input h (f-1)
of this layer's Jacobian to the prior, involves maximizing:
Juyer = (x). [log |det J (x)]
(29)
Substituting Eq. equation 27:
Juyer = (x)) [log Pmodel (cond) - log Prargrr (f(x))
(30)
= [x) (x) [log Pmodel (x cond)] (x)) [log Ptarger (f (x))
(31)
Let (our(y) ) be the actual distribution of the output y = (xx when x ~ gin(x), The second term in Eq. equation 31 can be rewritten using this definition:
Eng (x) log Prarger (x)) = E [log Purger (y)
(32)
So, Eq. equation 31 becomes:
Juyer= (x)( [log Pmodel (x (cond)) Egne(y) [log Prarget (y)
(33)
28

---

Relating to KL Divergence and Differential Entropy entropy (q) and Kullback-Leiller (KL) divergence (Llp))))
We use the definitions of differential
/ g(z) log (())dz
(q) (x) log q()) =
(34)
q(z)
/ ((() log g(z)
(L(gllp) = (() log
dz
(35)
p(z)_
p(z)
From the definition of KL divergence, we have KL(allp) = Bg [z)) -
=
- (q) - Rearranging this gives a useful identity for an expected log-likelihoud
term:
Bytz) [log p(z)] = - - (L(glpp)
(36)
Applying Eq. equation 36 to the two terms in Eq. equation 33:
E (x) [log (x (x cond)]
H H(mn) - KT (itoollmmode)
(37)
( [log Prarget (y)] - -H (gour) - KL {L(qqur |Parget)
(38)
Substituting these back into the expression for Juyerr:
Inyer H-H(un, (yin) KL (Pmodee)) - | H (4oat) KL( Prurges)
(39)
H (4out) H (gm) KL KLLgin||Pmosto P/mmoee) + (L Prargst)
(40)
This matches the main text.
The KL Divergence Identity for Normalizing Flows A key property normalizing flows the invariance of KL divergence under invertible transformations, specifically relating the KL divergence at the imput of a layer to the KL divergence at its output. We aim to show:
KL ([0odellPmmde)))) (L l|targer)
(41)
Recall that (x) is the actual density of the layer's input and (y) is the actual density of the layer's output, where y = fxx). These are related by the change of variables formula:
qin (x) Yoon (x)) |det (x)x
(42)
Similarly, Pmodel (xl cond) is the density parameterized by the layer for its input, and (ya is the fixed target density for the output. The transformation f constructed such that these are also related by the change of variables formula:
Pmodel (x(ond)) = Ptarger (f (x)) /det Jy(x))
(43)
Now, consider the ratio
Qni
Prnsdd x [ond)
qin (x)
Tom (f(x) /det Jx(xx]
Jout (f (f())
(44)
Pmodel X (cond)
Pranger (f()) ddet Jy(x)l
Dharger (f(x)
The KL divergence KL qua Pmode)) is defined as:
/ (lu(x)) log
(x)
(L Pmodea)
=
Pmodel (x cond) dx
(45)
Substitute Eq. equation 44 into the logarithm:
Yout (f(xx)
KL (L (mlooe)))))) = qin (x) log
Ptarget (f (xx)) dx
(46)
Now, perform change of variables in the integration from to y f(x). We have x x= f- and dx = J (y)l dy. The integral becomes:
KL (Pmodda) = qia "(x)) log Jout ( (y)
det J (x) )-) dy
(47)
From the change of variables for Oi (Eq, equation 42, rearranged for f-)), we know that Hu (y)) det Jy .y)) qoall). Substituting this into the integral:
KL (Pmoodaa))) -f
Gour (ouu(t)) (v) log Youa (y)
Ptarer dy
(48)
The right-hand side is, by definition, (T.(gous Puurget)) Thus, the identity in Eq. equation 41 is established.
29

---

Final Expression for Expected Log-Jacobian Substituting the KL divergence identity (Eq. equa- tion 41) into the expanded expression for Juyer (Eq. equation 40);
Jaaye = H (qout) (qm) - KL (9in (ullPPmdd)) + KL Qhn Pmodel))
(49)
H (owi)
H (qm)
(50)
This confirms the results in the main text: the expected log-determinant of the Jacobian for an invertible transformation is the change in differential entropy from its input distribution to its output
distribution.
Interpretation of the Learning Objective for a Single Layer The overall ELBO maximization drives the learning of parameters 0, which include the parameters defining Prodde (x cond) for each layer. The contribution of a single layer's Jacobian to the ELBO is Juyer = H (fom) ) - H (gm)).
However, the parameters of the current layer ʃ directly influence (x/oed( To maxi- mize Juyerr consider its expression before the KI. terms cancel: Juaer = H (Yous) (aa) (La + KLL You [Paagge)). Using the identity (L Pmodd))) = KL( Ptarga), this is also: Jayer = gx)) [log Pmoda coond]] - (x)((r) [log (y)]
The optimization process adjusts the parameters of f (which define Pmnade) to maximize this quantity. Maximizing (x,,() [log Pmodet (x (con))] equivalent to minimizing (4m (lmadde)) (since -H (gm) does not depend on the parameters of the current layer f. but rather on the data and preceding layers/encoder). When (L(gin Pmodet) minimized, Pmodel (x(ond) becomes a good approximation of the actual input distribution qin (x).
Due to the identity (T ((oddeelleee)) (T Prageet; minimizing KL ((mmeell)ee) simultaneously implies minimizing (Lou Prarget). This means that the actual output distribution qom is driven to match the fixed target distribution (y))
In summary, each flow layer learns to:
1. Model its input distribution:
The layer's parameters (defining Pmodel (x (codd) are adjusted so that Pmodd (x cond) accurately represents the distribution qin (in(x) of the data it receives. This corresponds to minimizing (T Pmodet))
2. Transform its input to target distribution: As a consequence of (1) and the construc- tion of f, the layer transforms its input into an output y whose distribution Tom (y) closely matches the predefined target distribution (y)) This corresponds to minimizing KL [Paarget).
This step-by-step process, repeated through the stack of flow layers, allows the overall normalizing flow to transform complex data distribution into simple base distribution (e.g., standard Gaussian), while enabling exact density calculation.
D
1-D Mixture-of-Gaussian and Nixture-of-Logiittes Coupling
The analysis in the main text shows that an AR MoG probability distribution can be interpreted as an implicit normalizing flow layer. This layer utilizes the cumulative distribution function (CDF) of the Gaussian mixture map between the data space u, and a base space (implicitly standard normal). Similar transformations, built explicitly around mixture CDFs, are employed in various normalizing flow architectures. This section provides a detailed comparison between the transformation implicitly defined by our model and common explicit parameterizations, establishing their functional equiva- lence. We will consistently analyze the transformations in the ddtaatto-lattett direction for clarity, and then briefly discuss the inverse (latent-to-data) direction.
For conciseness within this section, we may drop the time index t and dimension index i from variables and parameters when the context is clear (c.g., using x for a single data dimension, e or w for the corresponding latent dimension, and Wks Akk Jg for the mixture parameters predicted based on history). We maintain the standard normal PDF (p-) and CDF (-), and the standard logistic CDF (r) = 1/ + e e"r) and its inverse, the logit function o (p) log(p/ p)).
The Core Component:
Gaussian Mixture CDE Both the implicit and explicit layers rely on the same core component: the CDF of the Gaussian mixture distribution. For a single data dimension
30

---

the mixture CDF Fiix (2), conditioned on history (parameters Tk. Pk: Jk), is:
I
pie
Fiix (z)
40
(51)
O
The derivative of this CDF with respect is the probability density function (PDF) of the mixture:
V
m(x) df (2)
Σ
p4
1
(52)
dr
Ok
k=1
This mixture PDF m(xr) plays a key role the Jacobian determinant calculations.
Layer Definitions and Jacobians (Data-to-Latenl Direction). We now define two primary variants of the mixture CDF transformation layer. both mapping the data variable x to latent variable, but differing in the choice of outer non-linearity and corresponding base distribution.
MoG Variant: ProbitMiittueeCDF This layer corresponds the transformation 9t,; identified in Section Sec. 3.2.2. It maps the data a to a latent variable z, assumed to follow a standard normal distribution (0, 1).
(Fiiix (x))
(53)
Its Jacobian determinant, 0z/Or, is calculated using the chain rule and properties of d-1 and Fmix
Oz m((x)
(54)
dr (2)
The log-Jacobian determinant is therefore:
Oz
log
log mx) log (z()
(55)
dr
MoL Variant: LfgitMixtureCDF Explicit flow layers often use the logit function o as the outer non-linearity, mapping the probability Fmix (x) latent variable w, typically assumed follow
standard logistic distribution. We consider the core transformation without additional affine terms for now:
= (F
(56)
Its Jacobian determinant, Ow/Ox, is calculated using the chain rule. The derivative of o (p) is 1/(p(1-p) The PDF of the standard logistic distribution is Progistic (w) = o'(w) = (w) (w)). Since Faux (z) )- (w), we have Faix (1
(x)x())) = Progistic w).
Ow do (p)
d.Fmix (x)
m(x)
Fmix (x)x())(1 Fmix(14 (x)) -((x)
(57)
dz
dp
dr
Plogistic (w)
p= Pmii (2)
The log-Jacobian determinant is therefore:
Ow
log
log m (x: log Plogistic (w)
(58)
dr
Comparison in the Data-to-Latent Direction.
Table 4 summarizes the two variants when viewed
the data-to-latent direction.
Table 4: Comparison of Mixture CDF Flow Variants (Data æ + Latent).
Layer Variant
Transformation
Latent Base Dist.
Log-Jacobian log (0(latent)
④ 1111 ((Fix LogitMixtureCDF (core) w = o (Fmix (x))
ProbitMixtureCDF
N 1) Logistic(O.1)
log m(x) log sp(2)
log (ar) log Plogistic (w)
Both transformations share the same initial structure involving (m) and yield a log-Jacobian determinant of the form log mn(x) log Phuse Patpphusaaattt (latent), where Phase is the PDF of the corresponding base distribution (Gaussian or Logistic). The difference lies solely in the choice of the outer bijection (4-1 vs σ and the associated base distribution.
31

---

Connecting the Variants via Re-parameteriiation. The Probit and Logit variants are functionally equivalent because the standard normal and standard logistic distributions can be transformed into one another via simple, smooth, invertible function. Define the map IR IR that converts standard logistic variable w into standard normal variable
g(w) = + (aw))
(59)
The inverse map g - R IR converts standard normal variable into standard logistic variable wt
ツ (=) - -
(60)
These functions act translators between the two latent spaces.
can demonstrate the equivalence by composing one layer type with the appropriate function g or
LogitMixtureeCDF followed by g: If we take the output from the LogitMixtureCDDF layer (Eq. equation 56) and apply g. we get:
g(w) = g (x)) = $ (a(o (z)))) = (F (61)
This result exactly the output = of the ProbitMixtureCDDF layer (Eq. equation 53). ProbitMixtureCDE followed by g If we take the output from the ProbitMixtureCDF layer (Eq. equation 53) and apply g we get:
9 (2)
(Φ (#))
(0( (Fiix - (Fiixx
(62)
This result is exactly the output of the LogitMixtureCDF layer (Eq. equation 56).
This shows that the two layer types are related by composition with the fixed, invertible function g. In the context of normalizing flows, composing with such a function simply constitutes another valid flow layer, effectively re-parameterizing the latent space without changing the transformation's ' overall expressive power on the data x.
Likelihood Equivalence. The equivalence is further confirmed by examining the data log log-likelihood log p(x) induced by each transformation. Using the change of variables formula for the data-to-latent direction T (x): p(r) log Poase (2) + log det JT (x)].
ProbitMixtureCDD:
0z
log p(x) log Peanssian (z) + log
0x
log ((2) )+ (log m (I)
log (p()) log m(x)
(63)
LogitMixtureeCDF
dw
log p(z) log Plogistic (w) log
Or
log Plogistic (w) (log mxx)
log Plogistic (w)) log mx))
(64)
Both formulations yield p(z) = log mxx), the log-PDF of the Gaussian mixture model defining the transformation. The contribution from the base density's log-PDF always exactly cancels the second term in the layer's log-Jacobian determinant. This confirms that both variants define the exact same probability distribution over the data given the same mixture parameters (тT, Mix ox). Consequently, the gradients with respect these parameters during training will be identical.
Handling Affine Transformations. Explicit layers like the GaussMixCDF snippet often include an affine transformation y e"w + b applied after the core logit transformation = ( This affine transformation is itself an invertible flow layer with a simple constant log-Jacobian determinant:
(Fmix ((Fxx(())) (z)).
Oy Dw
log
log |e
(65)
This term simply adds a to the total deg-determinant of the composed flow (a - - transformations do not change the core functionality related to the mixture CDF and can be treated as separate layers (like Batch Normalization or ActNorm layers) or fused for implementation. Their presence or absence does not affect the fundamental equivalence discussed above.
y). Such affine
32

---

Latent-to-Data Direction. For completeness, we consider the inverse (generative) direction, map- ping the latent variable back to the data 2.
MoG Variant: ProbiMixture Inverse This corresponds to the transformation h:
I = h(z) = Fm mix (+))
(66)
The log- log-Jacobian is the negative of the forward log- Jacobian (Eq. equation 55):
dr
log
(log (()) log (0()) = log (0z) - log mm(r))
(67)
dz
Mol. Variant:
LogitMixtureCDF Inverse Inverting w = o (F mx)) gives Fmix(x) (w:), so:
Fixxx (a(w))
(68)
The log-Jacobian is the negative of the forward log-Jacobian (Eq. equation 58):
Ör = (log m=)) log Progistic (w)) = log Plogistie (w) - log (x)
log
(69)
Ow
Table 5 summarizes the inverse transformations.
Table 5: Comparison of Mixture CDF Flow Variants (Latent - Data x.)
Layer Variant
Transformation
Latent Base Dist.
Log-Jacobian log |0 (1atent)
ProbitMixtureCDF Inv.
Fm (#))
(O 1) Logistic(O.1)
log ((z) - log m(x)
LogitMixturcCDF Inv. F mix (ww))
Plogisic (w) log m(r)
The transformation implicitly defined by our AR-MDN base prior is functionally equivalent to explicit Gaussian Mixture CDF flow layers found in the literature. Both leverage the mixture CDF Fmix to perform complex. data-dependent monotonic transformation. Differences in the choice of outer bijection (probit vs. logit) and associated base distributions (Gaussian vs. logistic), as well as optional affine components, amount to implementation choices e-paraeterrizaiinss connected by the simple diffeomorphism g. They represent the same class of transformations and induce the same probability density on the data for identical mixture parameters. Our VAE approach benefits from this expressive power implicitly by optimizing the mixture density p(u,/uz) (.))) = m(u) directly, avoiding the potentially costly computation mis --1, or o during training.
E
Additional Experimental Details
This section provides supplementary experimental details and results that support the findings presented in the main body the paper.
Generative evaluation. For MAUVE, generate 5,000 samples for each configuration. We use GPT-2 Large as the embedding model, and set the MAUVE scaling hyperparameter to For generative perplexity, also GPT-2 Large the external language model to evaluate generated texts. We additionally report the average per-sequence entropy as a simple diversity metric.
Fig. 8 offers more complete view of the training and validation loss curves, expanding the comparison shown in Fig. of the main text. These curves illustrate the learning progression for different model configurations. The results underscore the benefit of incorporating mixture- based coupling layers, which exhibit improved convergence and final performance (lower negative ELBO) compared to configurations relying solely on standard affine coupling layers. This observation reinforces the point made in the main text regarding the importance of mixture couplings for effectively modeling discrete text data within our continuous latent space framework.
To further clarify the mechanics of our proposed model variants, Alg. details the forward transfor- mation process for the d-dimensional mixture flow layer. This transformation is a core component of the 'Mix-d' model configuration (discussed in Section 3.2.2), enabling the direct modeling of token-wise dependencies in d-dimensional latent space using mixture of Gaussians. The algorithm
33

---

Melhod
Timesteps
Entropy
Data
14.7
1.0
5.44
Autoregrestive
1024
12.1
0.76
5.22
SUNDAE
200
34.7
5.2
Ssd-LM
D3PM Absorb
1024
$42.3
SEDD
250/12/10244
110.1 107 2/ 104-7 103.2
0.007 / 00.00 70.008 00.00
563/
5.22 5.61
MDLM
254512/10242048
55.8 53.0 55,3/555
0.023 /.031 / /0,042 0.003
5.49/548
/.46 5,46
La singledim CDF (Ours)
20.7
D:47
5.58
TarPlow1.M multi-dim Rosemblatt (Ours)
Table 6: Generative evaluation comparing generation quality and diversity. We report generative
perplexity (Gen PPL: lower is better), MAUVE (higher is better), and average sequence entropy (higher indicates more diversity). Some baseline numbers are from [68, 71].
Alline Coupling Mixture Coupling
Afline Coupling
Mixture Coupling
9
8
(EL8D)
7
6
3
Training Step
x10
Figure 8: Comparison of training and validation loss curves across different model configurations. Our mixture-based coupling layers demonstrate improved convergence compared standard affine coupling.
outlines how input vector z is mapped output vector that follows standard Gaussian distribution, through sequence of conditional ID mixture-CDF transformations.
Alg. 2 describes the channel mixing and unmixing operations. These procedures are employed in the context of block-wise multi-token generation, as discussed in the "Flexible Patch Size" paragraph of the experiments section. By permuting latent dimensions between flow layers, channel mixing facilitates information exchange across different token positions within patch, allowing the model to capture intra-patch dependencies more effectively as data propagates through the stack of flow transformations. The ChannelUnmix procedure is the exact inverse, ensuring the overall transformation remains bijective.
F
Connection to Discrete Autoregressive Language Models
specific configuration our continuous latent space framework reveals a direct relationship with conventional discrete autoregressive language models. This connection arises when we consider the model structure outlined in Section 3 under particular assumptions for the prior p(Zier) and the encoder/decoder..
Specifically, we focus on the token-wise autoregressive prior for p(ZI.T) as factorized in Section 3.2, where each conditional p(zz(z.rr) is modeled as a d-dimensional mixture of Gaussians. We adopt the formulation from Eq. equation 7, where a shared codebook of Gaussian components is used. The key assumptions for this connection are: Tied Prior and Encoder Components: The
prior (o 2 Jk= used in the prior p( Zu Z<ce) are identical to the VAE V Gaussian components {140000
34

---

Algorithm 1 Forward Transformation for d-Dimensional Mixture Flow Layer: ga(z; C)
1: Input: Data z = (z1,..., zd) E R context C; mixture parameters (T,(C), my (C), (
2: Output: Transformed variable = (u,,..., 2a) E R
3: Initialize weights: a5" & TT(C) for j l...... K.
4: for = ., do
5: m1, f i-th component of my (C); 8j
(C(C)
6:
pi(zi(, - Ei= a
(2z $7)
7:
C, zci) 2==1 aj
KK
(0)
53
8:
4 4 & 1++0-11F1 (FF((, zcii)
9:
if d then
10:
aftl))
1 N(zamm,no))
for = 1, 1...., K.
11: Return
Algorithm 2 Channel Mixing and Unmixing Operations
1: procedure CHANNELMIX(()
2: D t size of the last dimension X
3: Xeven elements of X even indices along the last dimension
4: X 4- elements of X at odd indices along the last dimension
5: X' <- Concatenate( (Xeeaateeeeeee..e evven Xoda) along the last dimension.
6:
return X'
7:
procedure CHANNELUNMIX(X
X' is tensor, e.g., E RRBxTxPDm or R
8:
D <- size of the last dimension X'
9:
Assert D is an even number.
Dhalf f- D/2
11:
Xissttbalt f- the first Dhatr elements of X' along the last dimension
12:
X Shaf + the last Dhaif elements of X' along the last dimension
13:
Initialize Xour with the same shape as X' Place elements of X iistthall into the even-indexed positions of Xour along the last dimension.
15:
Place elements of X second_ball into the odd-indexed positions of X along the last dimension.
16:
return X
encoder's Gaussian components ( defined in Section 3.1 (Eq. equation 3.1). Thus, V equals the vocabulary size, fo prior pk. and (o mrioryz of The prior conditional then becomes:
p(z Zet) Σ Te ph, of Σ (Zz)))))1)N
(70)
where TT, (Z.r)) the mixture weights predicted by Transformer based the history Zec. Tied Encoder-Decodee The decoder p(x, 22)) uses the tied Bayesian formulation from Eq. equation 3.1:
p(z, klzz) Nxa(z))
(71)
Ny(,)
No Additional Flow Transformations for Prior: We consider the case where the prior p(Zi.))
Il p(z, (Z()))) with p(z, (.ce) defined by Eq. equation 70 is used directly in the ELBO (Fq. equation 3). This corresponds to setting
in the context of stacked flow layers (Section 3.3), meaning 21.7 is not further transformed into u1.7 via additional flow layers for this analysis.
Under these conditions, the ELBO (Eq. equation 3) is maximized. Let's analyze the terms in the ELBO related to a single token 2t, given its history X The relevant part of the ELBO expectation, for sample 21:7 ((-)x1:)), can be written per token t as:
C. log,((z: (2)) log p(zz (z))) log g (z))
(72)
35

---

Substituting the definitions:
te log (21)
+ log
Σ
Te. (2 (
log Nz, (2+)
(73)
k=1
V
= No, - -log
EN
+log
Σ
T Tr(ees))
Nog (21) (74)
log
Σ n (Z21 log
Σ Ny((z)
Na((e) Σ T: (xct))
log
Σ3 N2((z)
log
Σ (zer)) = le 21))
(75)
A=
The full ELBO involves an expectation Ex (XLer) Ee Left
Proposition 4. (Connection to Pooss-Entropy Discrete AR LM). Under the rying conditions specified above (tied prior/encoder components, tied decoder, and L 0 for additional flows), in the limit where the encoder/prior Gaussian components become infinitely narrow (i.e., o + 0 for all k = ...... V), minimizing the negative of the ELBO terms related to prior prediction and reconstruction for token Vt approaches minimizing - log T. (Mmyy.... Mas (rr.. This is equivalent to the negative log-likee likkeliiiood (cross-entropy) objective for discrete autoregressive model predicting loken Fa given the (deterministic embedding of) previous tokens.
Proof. Consider the term L from Eq. equation 75. The latent variables ze are sampled from ((z (x)) = N2, (2.) = (Zz;; May o2 I). As of + 0 for all k:
For a given true token 24, the sample z, ~ Nox (Z:) will converge in probability to its mean:
Mxx
2. Consequently, the history Zet (Zz, 2 will converge to (Hars---- -
3. The decoder probability p(arz klzz)) (Eq, equation 71) evaluated 2, pzz will behave
as follows: If μn are distinct, then for Z₁ ~ pz, N2. (z:) will be large, while (z(z,) for
if Iz.
j * Dt will be very small. Thus, p(x, klzz pm)
if #
Substituting these limits into Eq. equation 75, the sum Σ n., (z ( [k)p(2; ; = klz,) becomes:
Ti(Mo):-.. Me ( 1 1 Σ (M
Mrx [k] = T (Hz). Mr [ar)] (76)
Therefore, in this limit, the ELBO contribution L: approaches log ग. (Mayy --... Mz [aa]] Maxi- mizing the ELBO thus involves maximizing this term, which is equivalent to minimizing its negative:
log TT, (1 M [xz]] This is precisely the cross-entropy loss for predicting the true token I; given the sequence of means of the previous true tokens as context.
This connection illustrates that our continuous latent variable framework, under specific simplifying assumptions (most importantly, very low-variance, well-separated encoder components and direct use of the tied mixture prior), can recover the objective of standard discrete autoregressive models. offers perspective on how continuous space modeling can generalize or relate to established discrete paradigms. When of > 0, the objective involves "soft" version of this cross-entropy, where the target p(æ, klz,) is not one-hot.
36

---

G Latent Space Evolution Metrics
Our framework's 's formulation of language modeling within a continuous latent space, processed by stacked autoregressive normalizing flows, allows for step-by-step observation of text formation. At any intermediate stage l of the L flow transformations, the continuous latent sequence h h can be decoded using the decoder to yield corresponding textual output. This ability to materialize text from intermediate representations offers a view into how the model refines an initial latent state towards final coherent sequence. The continuous nature of the latent variables naturally enables this step-wise refinement process. Each flow layer performs a smooth, differentiable transformation across the entire sequence representation, effectively functioning as sophisticated editing operation. This type of fine-graing greeegraddd, fully Iearnable editing process is challenging to implement in discrete token spaces, where making intermediate adjustments typically involves solving complex discrete optimization problems.
To further characterize quantitatively the transformations, we analyze statistics of the latent codes h h
after each flow layer. These include: (1) the mean L2 norm of token embeddings h0:;; (2) the mean pairwise cosine similarity among these embeddings within a sequence; (3) the mean Participation Ratio (PR) of the set of token embeddings (hg (b(')T and (4) the mean centroid movement, measuring
the Euclidean distance between the average sequence representation mean (h (+) and that of the 1:T)
preceding layer, mean(h(??"))
Observations of these metries reveal structured evolution. A consistent decrease in the mean L2 norm across flow layers suggests that the transformations guide the latent representations towards a more defined and compact manifold. Simultancously, an increasing mean pairwise cosine similarity indicates that token embeddings within sequence become more semantically clustered, enhancing internal coherence. The utilization of embedding dimensions can be particularly revealing when considering alternating flow directions. For example, an initial Left-to- Right (L.2) flow might establish foundational representation with a certain value (e.g.. 7). A subsequent Right-to-Lefi (R2L) flow could then see a drop. This reduction might occur as the R2L pass focuses on integrating suffix-based or future context, temporarily constraining the representations to lower-dimensional manifold that captures these specific right-side dependencies. The final L,2R flow, benefiting from the synthesis of both left-anchored and night-anchored information, might then exhibit significant increase, indicating that it expands the representational capacity to integrate these now bi-directionally informed features into richer, more expressive state. A diminishing mean centroid movement with each subsequent layer further points to a coarse-to-fine adjustment, naturally enabled by the stacked continuous transformations: initial flow layers induce larger global changes, while later layers perform more subtle, fine-tuning modifications. This hierarchical refinement is a direct benefit of operating in continuous space where gradual adjustments are possible, unlike the often all-or- all-or-nothing choices in discrete generation.
This section provides detailed definitions and computation methods for the metrics employed to analyze the evolution of continuous latent representations (1)), across stacked autoregressive flow layers. We consider a batch of B sequences, where each sequence consists of T tokens. The d-dimensional latent embedding of the t-th token in the b-th sequence, after processing by the f-th
flow transformation, is denoted h E RRd, The initial state of these embeddings, prior to any flow
transformations (e.g., sampled noise or an encoder's output), is designated h The subsequent
bb.t states, corresponding to the outputs of the L flow layers, are b,,:.... .....,,,.t. Our statistical analysis primarily focuses on characterizing these I, flow layer outputs.
Mean 1.2 Norm The 1.2 norm (Fuclidean norm) of a token embedding h measures its magnitude in the d-dimensional space:
d
12 =
(77)
37

---

For each sequence at flow layer e, the average L2 norm of its token embeddings is calculated as:
T
NO
(78)
The Mean L2 Norm reported for flow layer is the average of these per-sequence values over all B sequences in the batch:
B
Mean L2 Norm
N
(79)
B
bm
The standard deviation of the set { J6-1 across the batch also computed. This metric indicates whether the flow transformations tend to expand, contract, or preserve the overall scale of the token embeddings.
(e) B (
Mean Pairwise Cosine Similarity (Intra-Sequencc) Cosine similarity quantifies the angular relationship between two embedding vectors. For any two token embeddings h and h within the same sequence flow layer l, their cosine similarity is:
(0) hg. 2
cos sim
(0) 6,4 bott
(80)
( |
112
(E) 6.12
lla + Ecos
where Ecos is small constant (e.g.. 10-9) added to the denominator for numerical stability. For each sequence b at flow layer f. (assuming > 1), the average cosine similarity is computed over all unique pairs of distinct tokens (11, (2) where 1 < <2:
T-1
T
6h T(T 1 (2:4+11
Σ
Σ
cos sim(h sim(h
(81)
If < 1, this average is typically considered undefined assigned a default value (e.g., 1f0 as there are no distinct pairs) The Mean Pairwise Cosine Similarity for flow layer é is the average of these per-sequence values over the batch:
B
Mean Pairwise Cosine Sim. (4)
1
(0)
(82)
B
b=1
The standard deviation of ( is also reported. This metric provides insight into the internal coherence or degree of representational similarity among tokens within sequence.
Mean Participation Ratio (PR) (Intra-Sequence) The Participation Ratio (PR) estimates the effective number of dimensions utilized by collection of embeddings. For a given sequence at flow layer É, let H())
(1) E T xd be the matrix where each row token embedding (2) This
analysis requires > 1 and > 0. First, the embeddings are centered by subtracting their mean
(2) 6, = + b, from each h, to obtain the centered matrix H. (2) b.centered The d x d sample
covariance matrix of these T centered d-dimensional embeddings is:
(
T-
b.oentened)
b,centered
(83)
Let {A,...,, be the non-negane non-negative eigenvalues (()
The PR for sequence at layer & is defined
12
2₂)2
PR!
(84)
ΣS +
EPR
where EpR is small constant to prevent division by zero all eigenvalues are zero. The PR ranges from 1 (if all centered embeddings are collinear, indicating usage of one effective dimension) to d (if
38

---

the variance is distributed isotropically across all d dimensions). The Mean Participation Ratio for flow layer & is the average over the batch:
B
Mean (R(
PR
(85)
B
The standard deviation of {PR, Jh is also reported. This metric indicates the breadth of the dimensional subspace actively occupied by the token representations within sequence.
(
Mean Centroid Movement (Inter-Step) The centroid of sequence at a specific original step E {O..... : (where s denotes the initial state, and L...., L denote the outputs of the successive flow layers) is the mean of its token embeddings at that step:
T
1
h
b,t
(86)
T
The centroid movement for sequence b induced by the f-th flow layer (which transforms the represen- tation from original step - to original step E) is the Euclidean distance between the centroids of these two consecutive states:
M
111
14 112, for l = 1,....
(87)
The Mcan Centroid Movement for flow layer é is the average of these per-sequence movement magnitudes over the batch:
B
Mean Centroid Movement
1
M
(88)
B
The standard deviation of 1 that each flow layer imparts the average representation of sequence.
Ib=1 also reported. This metric measures the magnitude of change
H FLOPs Calculation Details
This appendix provides a detailed breakdown of how the Floating Point Operations (FLOPs) for the forward pass of an entire sequence are calculated for both the regular Transformer model and the Flow Transformer model featuring the MAF-MLP head, as presented in Fig. 7. We assume multiply-accumule operation constitutes FLOPs. FLOPs from bias terms, normalization layers, and activation functions (like GELU or ELU in the main Transformer body or within the MAF-MLP module) are generally not counted, as their contribution is typically much smaller than that of matrix multiplications.
Regular Transformer Model The calculation for a standard autoregressive Transformer model as follows. Let L be the number of Transformer layers, be the model dimension (embedding dimension, d_model). ) be the input sequence length, and be the vocabulary size. We assume an MLP (Feed-Forward Network) ratio of 4, meaning the inner dimension of the FFN is 4d. We also assume that the sum of attention head dimensions dastn equals d.
The FL.OPs for each component for processing an entire sequence of length S are:
First, for Input Embeddings (Token + Positional), each token's embedding lookup and addition of positional encoding is approximated. For S' tokens, this is Cembed 4 S.
Second, for Attention QKV Projections, for each of L layers, input of shape (S, d) is projected to queries, keys, and values. This can be viewed as a multiplication by weight matrix of effective shape (d, 3d). FLOPs per layer: 2 x S x d 3d = 6 d? Total for L layers: CoKс = 6 L dl S.
Third, for Attention Logits (Query-Ker Key Dot Products), in each of layers, each of S query vectors (dimension d) computes dot products with S key vectors (dimension d). FLOPs for one query vector to attend to S key vectors: 5 (2d) 2 S. For all query vectors a layer: (2d = 5
Total for L layers: CoK 2 2
39

---

Fourth, for Attention Output Projection, for each of L layers, the attention output (shape (S, d)) is projected by a weight matrix of shape (d, d). FLOPs per layer: 2 x S x d 2 S d. Total for layers: Cproj 2 I. L S.
Fifth, for Feed-Forward Network (FFN/MLP)) each FFN block has two linear layers. The first linear layer (d + 4d): 2 x S x x 4d = S df. The second linear layer (4d + d): 2 x S 4d x = 85 d. Total FFN FLOPs per layer: 16 S.S. dL. Total for L layers: OFF 16 de
Finally, for Ouiput Linear Head (De-embeddingg, it maps final hidden states (shape (5, d)) to vocabulary logits (shape V)) using a weight matrix of shape (d, V). FLOPs: Chead 2 x $ x 2 V S.
Summing these components, the total forward FL.OPs for the entire sequence are:
Cregular
fwdd Red//squence = Cembed + CoKV + CoK + Proj + OFF + Chead
S S 16 S 2dVS 24 24
4d5 S
+4dS+2dVs
This can be factored by St
Creegular
Fud//squeence
=
S 24 + + 2dV
For the calculations related to Fig. 7, the parameters L (number of layers) and d (model dimension) correspond to the GPT-2 configurations specified in the main text. The sequence length set to 1024, and the vocabulary size V is 50257.
Flow Transformer Model (with MAF-MLP Head) This model uses the same Transformer body as the regular model but replaces the standard linear output head with the custom MAF-MLP module. The input sequence structure is also modified based on patch size P.
Let L be the number of Transformer layers in the main body, der be the model dimension of the main Transformer body (i.e., n_embd from GPT-2 configs), Shase be the original sequence length (e.g.. 1024), P be the patch size, and Sspecial |Stase/ PJ be the effective sequence length for the Transformer body, calculated using integer division. If this value becomes O (c.g.. if P >> Shase ), it is clamped for calculation purposes.
The FLOPs for the Transformer body (excluding any output head) are:
Cbady houdyy/uo-heall head = 24 L dis Supecial + 2 L dhe S2 Sapecial + 4 4 Sapecial
To this, we add the FLOPSs from the MAF-MLP module, calculated per token and then multiplied by Ssperial-
For the MAF-MLP module, let Gj be the module's input channels, determined by 16P, num mii__mittrres__ = 64, Cout per_in be the module's ' output channels per input channel, defined as parameter, fixed 128 for
3 num mixtures = 3 64 192, Dmoddle (the module's 's hidden_siss the plot) be the module's internal hidden dimension, and E be the module's computed embed_sii _ize, calculated as minxmax(1. int moddul x 10..0/160//// (cm 1))), 96). where int(. denotes integer casting. (We ensure
1 > 0 as P 2 = Gn 2 16).
The FLOPs for the MAF-MLP module (Cmodule) per token, considering only its linear layers,
First, the in_ ttttttt_faaaaa features layer projects an input of effective dimension 3(Gn - 1) to E(Cin 1). FLOPs: C 2 x (Cin 1) x E(Ci. - 1) = 6 E (Cin - 1)2.
Second, the features_ to hi dden layer projects an input dimension der + Elan 1) (concate- nation of transformer output features of dimension dur and module's 'ss internal in_ featurest to an output dimension of (D (mmoduule/22 22)c FLOPs: C2-
[du + E(cn - 1)] (Dmodule/2 x Gm)
Dmodule Gu [dut + E(Cin 1)].
Third, the hidden_to_ooe layer projects an input of dimension (Dmodule/ 2m) to Cout_ppe_innin FLOPs: C3 2 x (Dmodule/2 cin) (iont_per_inCnn) D Doodule eu Cout_per_in-
40

---

The total FLOPs per token for the module is moodule = C1 + C2 + O3:
module 6 E (Cin- 1)2 + D module Oin [det + E(In 1)] + D module En Sout_per_ia
The total forward FLOPs for the entire sequence in the Flow Transformer model are:
Cepecial fwd/soq
Cbody/ooee /nod + Ssporial x Cmodule
Sspecial
(24 [24dd + 2 2 Sspecial + 4d1) + Cmodule
The specific parameter values used generate the data for Fig. are summarized below. For the regular Transformer, the base model parameters (L, d, S 1024, V 50257) are taken from standard GPT-2 configurations (e.g., gpt2, gpt2-medium um, gpt2-large. gpt2-x1). For the Flow Transformer (Special Transformer), the main body parameters T, and der also correspond to these GPT-2 configurations. The parameters specific to its MAF-MLP head and sequence processing depend on the patch size P E {1, 2, 4} as follows: Sspecial = [1024/P] (clamped to if O); Cm 16P: Dmodule (the module's s hidden_size) is 128; num montmmmixueessstt = 64 resulting in Cout fout_pee__in and E min(max(1, [ Moodule 9/616 (cmn 1)) where int(.) denotes integer casting. These values directly correspond to those listed in Table and Table 2 of the user sper-provided values document that informed these calculations.
192:
I Related Works
Coupling-based Normalizing Flows A substantial body of work has focused on designing expres- sive invertible transformations with tractable Jacobian determinants for normalizing flows. NICE [16] pioncered the use of additive coupling layers, which made the Jacobian determinant computation straightforward. RealNVP [17] built upon this by introducing scaling and shifting operations, thereby increasing model flexibility. Glow [37] further improved these models by incorporating invertible 1 x convolutions, which led to better performance in image generation. Flow++ [29] enhanced expressiveness by integrating attention mechanisms. iResNet [2] showed that standard ResNet ar- chitectures [28] could be made invertible by adding normalization steps. Additionally, normalizing flows have been instrumental in improving Variational Autoencoders (VAEs) [38] by enabling more flexible posteriors [63, 78], and in diffusion models [62] by introducing flexible nonlinear drift and diffusion terms [36]. However, these approaches often require carefully designed and restrictive architectures, which can hinder scalability.
Continuous Normalizing Flows Neural Ordinary Differential Equations [1 reinterpret the ResNet architecture as deterministic ODE in the continuous-timg limit, which naturally extends to the concept of Continuous Normalizing Flows. In this setting, invertibility is guaranteed, and the Jacobian determinant reduces to the trace of the Jacobian. The adjoint method, based on Pontryagin Maximum Principle [54], enables efficient gradient computation with constant memory requirements. FFJORD [24] made Jacobian trace estimation more efficient by using Hutchinson's 's estimator [33]. Despite these advances, such models can be numerically unstable during training and sampling, as discussed [43, 81]. Their expressiveness can be further increased by introducing auxiliary variables [9, 18).
Autoregressive Normalizing Flows There has also been considerable progress in combining normalizing flows with autoregressive models. IAF [39] proposed dimension-ww wise affine transfor- mations conditioned on previous dimensions for variational inference, while MAF [50] utilized the MADE [20] architecture to realize invertible mappings via autoregressive transformations. Neural autoregressive flow [32] replaced the affine transformation MAF with monotonic neural network per dimension, increasing expressiveness at the expense of analytical invertibility. T-NAF [52] extended NAF by employing a single autoregressive Transformer. Block Neural Autoregressive Flow [8] instead fits an end-to-end autoregressive monotonic neural network, as opposed to NAF's dimension- approach, but also loses analytical invertibility.
Probability Flow in Diffusion Diffusion models [30, 60, 62) generate samples by simulating stochastic differential equations. Song et al. [62] introduced deterministic ODE formulation, known scoreflow [61], counterpart to the stochastic process. This scoreflow can be viewed
41

---

special case of continuous normalizing flows, where the learned score and base drift are combined into a new drift term. However, Lu al. [46] showed that the standard training objective in diffusion models, which is based on a first-order score approximation, does not maximize the likelihood for the scoreflow.
Diffusion models, other generative models Diffusion models [30, 62] have recently emerged as powerful generative models, achieving impressive results. Stable Diffusion [53] and OpenSora [80] have demonstrated the ability of diffusion models to generate extremely high-dimensional data. Other prominent generative models include Variational Autoencoders (VAEs) [38] and Generative Adversarial Networks (GANs) [23]. VQ-VAE [65] addresses posterior collapse and achieves strong generative performance, serving as a key component later latent diffusion models [53]. In the GAN domain, works such as [4, 34, 35] have demonstrated the ability to generate high-resolution images with relatively low inference cost, though training GANs remains challenging due to stability issues [69].
Continuous Diffusion for Discrete Data A common strategy for modeling discrete data is to operate in a continuous embedding space, adapting continuous diffusion techniques. This enables the of powerful continuous models, but requires mapping back the discrete domain. Diffusion-LM [40] applied continuous diffusion to word embeddings, allowing for controllable text generation via gradient-based guidance. Plaid [27] focused on likelihood-based text modeling, jointly optimizing embeddings and model parameters using the VLB, categorical reparameterization, an output prior, a learned conditional likelihood p(ar) (2)), and conditinni., CDCD [15] used a probability flow ODE on embeddings, employing score interpolation to jointly train embeddings and a denoising Transformer with a cross-entropy loss and time warping. Bit Diffusion [12] represented discrete data as continuous "analog bits." incorporating self-conditioning and asymmetric time intervals. While these methods are effective, they rely on continuous relaxations or embeddings, motivating the development of models that operate directly on discrete spaces. Many of these works also explore non-autoregressive, parallel generation approaches [3, 10, 13, 14, 26, 31, 40, 57, 72, 74, 75, 76], in contrast to sequential autoregressive models.
Discrete Diffusion Models Another line of research focuses on diffusion processes that inher- ently defined on discrete state spaces, typically using Markov chains. Building on early work [31, 60]. D3PM [1] generalized discrete diffusion by employing various structured transition matrices (such as uniform, absorbing, and Gaussian-like) and training with a hybrid "LB/cross-entropp loss. Campbell et al. [6] extended this Continuous- -ime Markov Chains (CTMCs), deriving continuous-time ELBO and proposing efficient sampling methods like tau-leaping and predictor-comeetor schemes, leveraging factorization for scalability.
Rather than simulating Markov chains directly, some works define score-like quantities for discrete diffusion. The concrete score, given by the ratio of marginal probabilities pa (y)/p, (x), serves discrete analogue to the continuous score [45, 48). SEDD [45] trained models using a score entropy objective (T.DSE) based on this ratio, relating it to the ELBO and employing Tweedie r-leaping for sampling. Sun et al. [64] introduced categorical ratio matching in CTMC framework, learning singleton conditionals p. (x)))) with a tractable loss and an analytical reverse sampler. Building on this, Ou et al. [49] showed that for absorbing diffusion, the concrete score decomposes into a time-independent conditional and a time-dependent scalar, simplifying the model (RADD) and leading to the Denoising Cross Pnoss-En (DCE) loss.
Masked (absorbing) diffusion, which replaces tokens with a special [MASK] token during the forward process, has proven highly effective. MDLM [56] introduced a substitution-bass parameterization (SUBS) and derived a simplified Rao-Blackwellizee ELBO equivalent to weighted Masked Language Modeling (MLM) losses, enabling generative training of encoder-only models. Shi et al. [58] (MD4) further unified this framework, deriving a simple ELBO with SNR invariance properties akin to continuous diffusion and generalizing to state-dependent masking schedules.
Further advances have refined the parameterization and mechanisms of discrete diffusion. Reparame- terized Discrete diffusion Models (RDM) [79] identified a route-and-denoise mechanism, reducing the objective to cross-entropy on noisy tokens and enabling adaptive routing during sampling. Liu et al. [44] proposed Discrete Diffusion with Planned Denoising (DDPD), which factorizes the
42

---

reverse process into a planner (predicting corruption) and a denoiser, allowing for adaptive sampling via the Gillespie algorithm guided by the planner.
Discrete Flow Matching provides another generalization. Gat et al. [19] defined probability paths interpolating between discrete distributions and derived corresponding probability velocities, analo- gous to continuous flow matching, yielding unified sampling theory. [7] formulated discrete flows using CTMCs, learning scores cross-entropy and enabling flexible inference by adjusting the rate matrix family at test time without retraining, also supporting multimodal generation. Discrete diffusion concepts have also been extended to structured data, such as graphs, as in DiGress [67]. which uses specialized noise transitions, auxiliary features, and classifier guidance.
43

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA]-
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1 (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked While "[Yes] " is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providded a proper justification is given (e.g., "error bars expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
the checklist one of the factors in their evaluation.
reported because it would be too computationally
"[No] " or "[NA] is not grounds for rejection. While the questions are phrased acknowledge that the true answer often more nuanced, so please just use your best, judgment and write justification elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction clearly state the contributions, which are devel- oping a continuous latent space language model using autoregressive flows, proposing new transformations, and evaluating its performance and flexibility. These align with the paper's content as presented.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
44

---

Answer: [Yes]
Justification:
The paper discusses limitations in Section 4.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
The paper presents theoretical underpinnings for its methods, particularly for the normalizing flow transformations and mixture models. It refers to appendices (e.g.. Appendix A, B, C, D) for detailed derivations and proofs.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
45

---

Justification:
The paper provides details on datasets, model configurations (GPT2-Sma Small based, specific parameters for flow layers and mixture components), and evaluation metrics in Section 4 and appendices (e.g., Appendix G for FL.OPs..
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: The provided text of the paper does not explicitly state that the code or specific model implementations will made publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc for more details.
public/ public/guides// Solieee//oooeuubiisesees
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /niiss. cccppcccpcccccccccccccc desidee///oeeuueeeeeeeeeee
ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
46

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e-g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section 4 details datasets, model configurations (including architectural choices, latent dimensions, mixture components), and evaluation metrics. Data splits standard. Some hyperparameters like optimizer type explicitly mentioned but other key parameters are.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
The main experimental results presented in tables (e.g., Table 1) and figures within the provided text do not include error bars or report statistical significance tests. Appendix F mentions standard deviations are computed, but these are presented with the main results.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
47

---

Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [No]
Justification: While Appendix G provides FLOPs calculations, the paper does specify the type of compute hardware (e.g.. GPU models, CPU rypes), memory configurations, or wall-clock execution times for the experiments conducted.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: //neur//eiii ips http:ccbcc//eeicccoiiiieee
Answer: [Yes]
Justification: Based on the provided text, the research focuses algorithmic development for language modeling and appears to align with standard ethical research practices.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e-g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper discusses both potential positive societal impacts and negative societal impacts of the work performed in Section 4.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
48

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks. mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
The paper proposes a novel modeling framework. does not discuss the release of specific pretrained models or datasets that would necessitate detailing safeguards for high-risk misuse at this stage.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [No]
The paper cites original sources for datasets (c.g.. TEXT8, OPENWEBTEXT) and related architectural components (c.g., GPT-2). However, the specific licenses and terms of use for these assets are not explicitly stated within the provided paper text.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi thcode. com/ com/datase datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
49

---

Justification;
The paper introduces new framework (TarFlowLM) and associated method- ologies. The paper itself serves as the primary documentation. No separate, releasable new assets (like a specific codebase or dataset) are explicitly mentioned as being provided with their own documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The research described in the paper focuses on language modeling using existing datasets and does not appear to involve crowdsourcing or direct research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code Ethics, workers involved data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: As the research does appear to involve human subjects, IRB approval or discussion of participant risks is not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
50

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
The research focuses developing novel language modeling architecture itself. It does not appear to use an existing LLM as an important, original, or non-standard component of its core methodology.
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer our LLM policy (https: //eusr////u.p....ee /tterrppp..cccooffeeeee//////// for what should or should not be described.
51


Reviews:
[{"review_id": "A0RVfqIhXg", "reviewer": "Reviewer_J8wg", "summary": "The paper introduces  TarFlowLM a novel framework to model language in a continuous latent space using Transformer-based autoregressive normalizing flows. The method encodes sequences of discrete tokens with a VAE into Gaussian latent vectors, stacks alternating-direction flow layers built from 1-d CDF and n-d Rosenblatt mixture couplings to transform latents to a standard normal, mixes latent channels to share information across token “patches,” supports block-wise generation with flexible patch sizes.", "questions": "- I am not sure the motivations of the paper are explained in a sufficiently clear way. Specifically, I don't fully understand the difference between a standard transformer model -- where discrete tokens are mapped to continuous embeddings and the their interactions is modelled in such continuous space -- and the current framework. Could the author help me understanding the core motivation of their framework and how it differs from standard autoregressive approaches?\n- Given a token context, how is inference performed?", "limitations": "Limitations are discussed at page 10 in a separate section. I would suggest the authors to move them in the appendix or incorporate them in the main text within the 9 page limits.\n\nI would also encourage the authors to comment on the results in Table 2. While the paper presents a novel method for language modelling in continuous spaces, at this stage it seems to still underperform compared to standard discrete space autoregressive baselines.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strenghts**\n- I found the ideas behind the disentanglement between the \"internal vocabulary\" size and the actual one interesting. Also, the framework allows for editing at various scales by intervening at different stages.\n\n**Weaknesses**\n- The paper is poorly written with several typos (e.g. lines 87, 116, 118. 169, 197. 244, ...) and some important parts are missing. For example, there is no clear mention on how inference/generation is performed other than Figure 1, which does not clearly explain how to deal with an existing context.\n- Some figures are referenced, but never discussed. An example is Table 2 that reports performance comparisons against other architectures. In addition, the results showed in this Table show that TarFlowLM is outperformed by the standard autoregressive transformer baseline.\n- Some figures are included but I was not able to find any reference to them: see for example Fig. 5, Fig. 6.\n- The appendix section E (Additional Experimental Results) presents no additional experimental results. Figure 8 appears to be a continuation for more steps of Figure 3 which already shows a clear difference between affine coupling and mixture coupling.\n- No discussion/results on the generation speed difference against other non transformer AR baselines.\n- Related works are relegated to the Appendix at page 35.", "was_revised": true, "final_justification": "I thank the authors for their responses. The rebuttal clarified some of my concerns but I still believe the paper would need another round of revision given its current form. I have increased my score from 2 to 3.", "rebuttal": "We thank the reviewer for your valuable time and constructive feedback. We appreciate the positive recognition of our framework's flexibility regarding \"internal vocabulary\" size and its multi-scale editing capabilities.\n\n---\n\n# Q1\n\n> \"I don't fully understand the difference between a standard transformer model -- where discrete tokens are mapped to continuous embeddings and the their interactions is modelled in such continuous space -- and the current framework. Could the author help me understanding the core motivation of their framework and how it differs from standard autoregressive approaches?\"\n\nThank you for raising this crucial question, as it highlights the fundamental distinction of our work\n\n1.  **Standard AR Language Models (e.g., GPT):**\n    *   **Objective:** To model the conditional probability distribution $P(x_t | x_{<t})$ over *discrete* tokens.\n    *   **Output:** At each timestep $t$, the Transformer outputs *logits* that define a **discrete Categorical probability mass function (PMF)** over the vocabulary. Sampling involves drawing a single discrete token from this PMF.\n    *   **Invertibility:** Not inherently invertible end-to-end. There's no smooth, invertible mapping from a sequence of discrete tokens back to a standard latent space that allows for exact likelihood computation and multi-pass transformations.\n    *   **Nature of Modeling:** A sequence of *discrete classification* problems.\n\n2.  **TarFlowLM:**\n    *   **Objective:** To model the *joint continuous probability density* $p(z_{1:T})$ over an entire sequence of continuous latent vectors $z_t \\in \\mathbb{R}^d$. This is achieved by transforming $z_{1:T}$ through a series of invertible normalizing flow layers until it becomes a simple standard Gaussian ${u}_{1:T}$.\n    *   **Output:** The model learns **continuous probability density functions (PDFs)** (specifically, Mixtures of Gaussians) at each step of the flow, which are then used for exact density estimation. Generation involves sampling from a continuous base distribution and inverting the flow.\n    *   **Invertibility:**  fully invertible. This is a defining characteristic of normalizing flows, allowing for exact likelihood computation (ELBO) and flexible multi-pass operations.\n    *   **Nature of Modeling:** A problem of *continuous density estimation* and transformation.\n\nThis shift from \"next-token classification\" to \"full-sequence continuous density estimation\" via invertible flows is what unlocks our framework's novel capabilities, which are inherently difficult or impossible for standard AR models:\n\n*   **Global, Bi-directional Context:** Standard AR models are strictly causal (left-to-right prediction). Our model can stack flow layers that apply transformations in alternating directions (e.g., left-to-right and right-to-left). It can refine information across the whole sequence, including initially \"future\" context, within the same generation pass.\n*   **True Block-wise / Multi-token Generation:** Our model can generate and model dependencies for a \"patch\" of $K$ tokens simultaneously. This is because we operate on continuous latent representations, and the internal flow layers can learn to mix information *within* that patch. In contrast, while standard AR models might be trained with multi-token prediction objectives, they still generate strictly one token at a time during inference to maintain autoregressive consistency. Our approach allows for parallel processing within a patch.\n*   **Hierarchical, Multi-pass Refinement and Editing:** The invertibility of our flows allows us to decode the latent representations at *any* intermediate layer. This provides a unique, step-by-step view of how the model incrementally \"edits\" and refines initial latent states into coherent text (as shown in Figure 6). \n*   **Flexible \"Internal Vocabulary\" Size:** As noted by the reviewer, our mixture-based transformations allow us to choose the number of mixture components independently for each flow layer, decoupling it from the data's discrete vocabulary.\n\nIn essence, **TarFlowLM is not just another Transformer architecture variant**; it's an exploration of a fundamentally different paradigm for language modeling rooted in continuous density estimation with invertible transformations, which we believe opens new avenues for flexible and controllable generation.\n\n---\n\n> In addition, the results showed in this Table show that TarFlowLM is outperformed by the standard autoregressive transformer baseline.\n\nThis highlights a fundamental distinction in modeling paradigms and evaluation metrics.\n\n1.  Methodological Difference: ELBO vs. Exact Log-Likelihood:\n    *   Our framework is an VAE which models the data by maximizing the ELBO which by definition is a *bound* on the true log-likelihood.\n    *   A standard discrete AR model, on the other hand, directly computes and maximizes the **exact log-likelihood**\n    *   Therefore, a gap between the two is not only expected but mathematically inherent to the VAE framework.\n\n2.  The Appropriate Comparison Baselines:\n    Given this distinction, the most relevant and informative baselines are other generative models that share our core paradigm: they first map discrete data to a continuous latent space and then learn a generative model in that space. This includes:\n    *   **Continuous-space models** like continuous diffusion (Plaid, BFN) and other normalizing flows for text (Latent NF, Argmax Flow).\n    *   **State-of-the-art discrete diffusion models**, which also operate in a latent space, albeit a discrete one, state-of-the-art for non-autoregressive language models.\n\n3.  State-of-the-Art Performance in the Correct Context:\n    When viewed through this appropriate lens, our results in **Table 2** demonstrate the strength and competitiveness of our approach.\n    *   Our TarFlowLM (Mix-d) variant achieves 1.30 BPC on \\texteight and 22.64 Perplexity on OpenWebText.\n    *   This performance outperforms all other continuous-space models (both diffusion and flow-based) listed in the table.\n    *   Furthermore, our result is highly competitive with and surpasses most state-of-the-art discrete diffusion models (e.g., SEDD Absorb, MD4), positioning our method at the forefront of modern non-autoregressive and continuous-latent-space generative modeling for text.\n\nIn summary, while a numerical gap to the exact-likelihood discrete AR model is an expected consequence of our VAE framework, our work demonstrates state-of-the-art performance when compared to its true methodological peers. \n\n---\n\n### **Q2: Inference/Generation Procedure.**\n\n\nWe have prepared a comprehensive description of the generation process, including how it handles existing context, which we will add to the Appendix in the revised version.\n\n1.  For the patch/token to be generated, we first sample \"noise\" vector ${u}_t$ $\\mathcal{N}({0}, {I}_d)$.\n2.  Inverse Flow Transformation: This sampled ${u}_t$ is then passed *backward* through the stacked normalizing flow layers in reverse order, i.e., $L \\to 1$. At each layer $\\ell$, the transformation is inverted: ${h}_t^{(\\ell-1)} = [f^{(\\ell)}]^{-1}({h}_t^{(\\ell)} | \\mathcal{C}_t^{(\\ell)})$. The context $\\mathcal{C}_t^{(\\ell)}$ for this inversion is provided by the Transformer conditioned on the previously generated latent vectors $z_{<t}$ (which are already fully inverted through all flow layers).\n3.  Latent to Token Decoding: The final output from the inverse flow process for the current patch/token is $z_t = {h}_t^{(0)}$. This continuous latent vector $z_t$ is then fed into the tied decoder $p(x_t | z_t)$. \n\nThis sequential generation, where each step involves inverting a flow, is a standard procedure for autoregressive normalizing flows, where the context is encoded and provided through the conditioning in each AR flow $[f^{(\\ell)}]^{-1}({h}_t^{(\\ell)} | \\mathcal{C}_t^{(\\ell)})$ as shown above. We will ensure this is clearly explained in the appendix.\n\n> Generation Speed\n\nWe acknowledge this is a critical aspect, and we specifically highlighted it as a limitation in Section (\"Limitations\"), lines 296-300.\n\nAt the same time, we also introduced \"Flexible Patch Size: Block-wise Multi-token Generation\", which demonstrates a computational advantage.\n\nFurther optimization of sampling speed through advanced techniques (e.g., parallel decoding for flows) is indeed an important direction for future work, as noted in our limitations.\n\n---\n\n# Q3: Figure References\n\n\nWe appreciate the detailed feedback.\n\n*   Table 2 Discussion and Performance:\n    *   Table 2 *is* discussed in the main text, specifically in lines 220-226 within the \"Experiments\" section.\n\n*   Unreferenced Figures (Fig. 5, Fig. 6):\n    *   Figure 5: This figure is discussed in lines 270-277 under \"Model Ablation.\" We will ensure it is explicitly referenced as \"Figure 5 (Model Variants Comparison).\"\n    *   Figure 6: This figure is discussed in lines 263-269 under \"Flexible text editing in continuous space.\" We will ensure it is explicitly referenced as \"Figure 6 (Text Editing Illustration).\"\n\n*   Related Works in Appendix:\n    *   We agree and will move it to the main text. This was a decision driven by strict page limits during the initial submission.\n\n---\n\n# Q4: Typos.\n\nThank you for pointing out the typos. We have corrected them in the revised version.\n\n\n# Q5: Limitations Section\n\nWe agree. We will move the \"Limitations\" section to the appendix.\n\n---\n\nWe genuinely thank the reviewer for their diligent review and insightful comments. We believe addressing these points significantly strengthens the paper, especially by clarifying the novel contributions and the performance context."}, {"review_id": "9ojNUM8xxX", "reviewer": "Reviewer_pnYG", "summary": "This paper presents TarFlowLM, which is a Transformer-based models that shifts the process from discrete token space to a continuous latent space. TarFlowLM instead of predicting tokens one by one, encodes discrete tokens into continuous vector representations and then generate the output by sampling from a standard Gaussian and decoding it into text. More specifically, in training, it encodes tokens into Gaussian latents, transforms by flow layers into standard Gaussian variables, and finally reconstructs them with the use of a tied decoder. Experiments are conducted on language modeling tasks and show TarFlowLM's good performance.", "questions": "Please see Weaknesses for the questions.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths:\n---\n- The paper is generally well-written, well-motivated, and easy to follow.\n- The presentation of the progressive layer-wise training scheme is very interesting and practical. It can be very effective method for specifically optimization of deep flow models with improving stability and performance\n-  One of the key strengths of the paper is its theoretical results that establish equivalence between a special instance of the approach with the standard discrete autoregressive models. This provides better insight on the generalization of the model and can guarantee a minimum expectation of the performance.   \n-  I found it particularly interesting and convincing that the proposed mixture-based coupling transformations are more effective than standard affine coupling layers for modeling the complex distributions that arise from mapping discrete text data to a continuous space.\n- In general, I found this direction very important and compelling as a replacement discrete autoregressive paradigm. \n\n## Weaknesses:\n--- \n-  One of the critical weaknesses of the work is the lack of comparison on common benchmarks on language modeling. Could you please clarify why there are no experiments on tasks such as question answering (e.g., HellaSwag, etc.)? I suggest adding more results on downstream evaluations such as translations and/or language modeling, question answering, etc.\n- Based on the experiments, it is not clear how the context length can affect the performance of the method. There is also a lack of discussion about long-context ability of TarFlowLM. These days, having longer-context is one of the most aspect of scaling models, is there any advantage for TarFlowLM compared to Traditional autoregressive approaches? \n- The `Transformer FLOPs comparison` presented in Figure 7 is not convincing enough. Could you please elaborate more on why not reporting the wall-clock time? The presented FLOPs as the difference in #FLOPs in token embedding phase, might cause misleading results. I suggest to further support the efficiency claim, report the wall-clock time of these methods.   \n\n- **Minor**: Missing citation in line 197,", "was_revised": true, "final_justification": "I thank the authors for their responses. Although they provided detailed responses on the rationale for their initial experimental choices, still the limitations exist. However, still I believe this is a valuable contribution and so I will maintain my positive score (but didn't increase it due to the mentioned limitations).", "rebuttal": "We sincerely thank the reviewer for your positive and constructive feedback. We are encouraged that you found our paper \"well-written, well-motivated,\" and the research direction \"very important and compelling.\" We appreciate the insightful comments, which have helped us identify areas for clarification and improvement.\n\nWe address the reviewer's specific concerns below:\n\n**1. On the Lack of Downstream Task Evaluation (e.g., HellaSwag)**\n\nWe agree with the reviewer that evaluating on a broader range of downstream tasks would be an excellent way to demonstrate the practical utility of our model.\n\nOur primary goal in this work was to introduce and validate a novel modeling paradigm: autoregressive language modeling in a continuous latent space using novel mixture-based normalizing flows. As this is a foundational investigation into a new modelling class, our focus was on rigorously establishing its viability and core capabilities.\n\nTo this end, we chose to align our evaluation with the established protocols in the nascent field of non-standard generative language models (e.g., discrete diffusion models like MDLM and MD4, and continuous diffusion models like Plaid). These works have predominantly focused on likelihood-based benchmarks (perplexity/BPC) to demonstrate the fundamental capacity of the new model family to capture the data distribution. Same as these prior work who are also exploring new paradigms for langauge modelling, we chose to use same set of evaluation protocols for our work too.\n\nWe believe our strong likelihood results provide a solid foundation, confirming that `TarFlowLM` is a competitive approach within this domain. We are excited about the future direction of scaling these models and evaluating them on a full suite of downstream tasks, and we will clarify the scope of our current work in the paper.\n\n**2. On Long-Context Ability**\n\nThe reviewer raises a very important and timely point about long-context capabilities. We acknowledge that an in-depth exploration of long-context performance was not a focus of this initial study, and we appreciate the suggestion to discuss this aspect.\n\nWe believe `TarFlowLM` offers a promising architectural direction for this challenge. A key feature of our framework is the ability to model **global, bi-directional context** through stacked, alternating-direction autoregressive transformations. This mechanism could be particularly advantageous for capturing long-range dependencies that are challenging for strictly unidirectional models, as information can propagate in both forward and backward directions through the layers of the flow.\n\nFurthermore, the **block-wise generation** capability, which shortens the effective sequence length, may offer computational benefits for processing long contexts. We see this as a key direction for future research and will add a discussion of these potential advantages and future work to the paper.\n\n**3. On the Transformer FLOPs Comparison and Wall-Clock Time**\n\nWe thank the reviewer for this insightful question regarding the practical efficiency of our model.\n\nOur initial choice to report FLOPs was to provide a hardware-agnostic measure of computational cost, which is a common practice when evaluating model scalability during the training phase. The goal of Figure 7 was to illustrate how block-wise processing (i.e., using a larger patch size) fundamentally changes the computational trade-off by reducing the quadratic cost associated with the sequence length in the attention mechanism.\n\nMotivated by the reviewer's feedback, we conducted a wall-clock time comparison during this rebuttal period using our `Mix-1` model on the OpenWebText dataset. Our findings are nuanced and provide a more complete picture:\n*   For a small patch size (K=1, 2), the training iteration time is comparable to a standard AR model. We hypothesize this is because modern GPU kernels like FlashAttention are highly optimized for long sequences, mitigating the theoretical benefit at this scale.\n*   However, as the patch size increases (K≥4), the reduced effective sequence length leads to a significant wall-clock time reduction of ~30% per training iteration.\n\nThis confirms that the architectural flexibility of `TarFlowLM` can translate to real-world efficiency gains, especially as the patch size is increased. We believe this is a compelling result. We will add a detailed discussion of these new findings, including the wall-clock time analysis, to the appendix in the final version to provide a more comprehensive view of the model's efficiency.\n\n**4. Minor Point: Missing Citation**\n\nThank you for catching the missing citation. We have corrected this in our revised draft.\n\n---\n\nWe hope these clarifications and additional results address the reviewer's concerns. We are grateful for the constructive feedback, which will undoubtedly strengthen the final version of our paper."}, {"review_id": "7rdi8MKoIx", "reviewer": "Reviewer_GUST", "summary": "This work develops a language model that performs language generation in a continuous latent space. They project discrete tokens into learnable continuous representations that are then modeled with autoregressive normalizing flows (using the transformer architecture). By stacking multiple unidirectional flow layers, the model can incorporate bidirectional context into generation. The authors introduce mixture-based coupling transformations which significantly outperform affine coupling layers. They evaluate their approach on two language modeling benchmarks (TEXT8 and OpenWebText), reporting validation likelihoods.", "questions": "What is the generative performance of the model? In terms of generative perplexity, MAUVE score, etc?\n\nHow does changing the patch size affect the likelihood and generative performance of the model?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n\nThis is an interesting approach to language generation that I am not aware of being explored recently. This kind of approach represents an interesting alternative to autoregressive and discrete diffusion models. The overall approach is well-formulated and theoretically sound.\n\nThe proposed mixture-based coupling transformations significantly improve upon the more common affine coupling transformation, although they are not the first to introduce such a coupling transformation. The discussion of the equivalence of their Mixture-of-Gaussian coupling transformation and the previously proposed Mixture-of-Logistics coupling transformation in the appendix is an interesting theoretical contribution in its own right.\n\nThe performance of their approach is strong compared to recent popular discrete diffusion models (e.g. SEDD, MDLM).\n\nThe flexibility of their bidirectional approach with respect to patching is a nice advantage of this framework, enabling the model to operate at multiple resolutions.\n\n**Weaknesses**\n\nIt is usually the case with these kinds of work that there is a persistent gap with the simpler autoregressive approach and that is also the case with this work. I don’t consider this a strong weakness as novel approaches are of course welcome. \n\nThe biggest limitation is the lack of metrics being reported for actual language generation. Reporting likelihoods is great, but quantifying actual generation performance is absolutely critical. Actual generation metrics such as generative perplexity, MAUVE score, diversity metrics, etc. would significantly strengthen the work. I recognize that one plot is included compared against sampling two tokens at a time from an autoregressive model. I understand the point being made by that plot, but comparing against the actual generative behavior of the baselines is critical.\n\nThere is a fair amount of discussion of the patching, but limited results are presented related to that design choice. It would be interesting to actually see that ablated to see if reducing the patch size is feasible in practice or whether that significantly degrades quality.\n\nThere were some parts of the writing that can be improved. For example, missing capitalization on line 83, line 965 is not grammatical.", "was_revised": true, "final_justification": "My primary concerns revolved around the lack of rigorous evaluation of language generation behavior and the lack of an ablation studying the influence of the patch size on the performance of these frameworks.\n* The authors presented language generation evaluations during the rebuttal period that are quite strong. Their framework compares favorably against discrete diffusion baselines and only lags behind the autoregressive baseline. As stated in the original review, I don't view lagging behind such an established approach as a large negative for novel approaches. \n* The patching ablation study nicely provided insight into the interaction between the patch size an number of flow layers. It would be nice to see the impact on the patch size/flow layers on language generation metrics as well, but the presented ablation is reasonable.", "rebuttal": "We are very grateful to the reviewer for their time and for providing such thoughtful and constructive feedback. We are encouraged that the reviewer found our approach to be a \"novel, interesting, and theoretically sound\" alternative to existing paradigms. The suggestions have been invaluable in helping us strengthen the paper.\n\nBelow, we address each of the reviewer's points in detail.\n\n### **1. On the Performance Gap with Discrete Autoregressive Models**\n\n> *“It is usually the case with these kinds of work that there is a persistent gap with the simpler autoregressive approach and that is also the case with this work. I don’t consider this a strong weakness as novel approaches are of course welcome.”*\n\nWe thank the reviewer for this insightful comment and for acknowledging the value of exploring novel approaches. We fully agree that a performance gap in likelihood still exists when compared to state-of-the-art discrete autoregressive (AR) models.\n\nOur work's primary goal is to establish the viability and unique advantages of a new design space for language modeling within continuous latent spaces. As the reviewer noted in the strengths, our model demonstrates superior likelihood performance compared to other continuous-space models (e.g., diffusion, other flows) and also surpasses recent discrete diffusion baselines (Table 2). We believe this positions our framework as a strong and promising new direction, and we are optimistic that future work can further close the gap with traditional AR models while retaining the unique flexibility our method provides.\n\n### **2. On Generative Quality Evaluation (Generative PPL, MAUVE, and Diversity)**\n\n> *“The biggest limitation is the lack of metrics being reported for actual language generation. Reporting likelihoods is great, but quantifying actual generation performance is absolutely critical. Actual generation metrics such as generative perplexity, MAUVE score, diversity metrics, etc. would significantly strengthen the work.”*\n\nThis is a critical point, and we thank the reviewer for this excellent suggestion. To address this, we have conducted a new set of experiments to rigorously evaluate the generative quality of our models on OpenWebText.\n\nThe table below presents these new results, including **generative perplexity (Gen PPL)**, **MAUVE**, and **sequence entropy**, comparing our two main model variants against strong discrete diffusion and autoregressive baselines. The results are highly encouraging and demonstrate the strong generative capabilities of our approach.\n\n| Method               | Timesteps              | Gen PPL ↓           | MAUVE↑           | Entropy↑             |\n|---------------------|------------------------|----------------|-----------------|-----------------------|\n| **Data**            | -                      | 14.7            | 1.0               | 5.44                   |\n| **Autoregressive**  | 1024                   | 12.1            | 0.76               | 5.22                   |\n| SUNDAE              | 200                    | 34.7            | -               | 5.2                   |\n| Ssd-LM              | >10000                 | 99.2            | -               | 4.8                   |\n| D3PM Absorb     | 1024                   | 842.3           | -               | 7.6                   |\n| SEDD                | 256/512/1024/2048      | 110.1 / 107.2 / 104.7 / 103.2 | 0.007 / 0.008 / 0.008 / 0.008 | 5.63 / 5.62 / 5.62 / 5.61 |\n| MDLM                | 256/512/1024/2048      | 55.8 / 53.0 / 51.3 / 51.3 | 0.023 / 0.031 / 0.042 / 0.037 | 5.49 / 5.48 / 5.46 / 5.46 |\n| ReMDM               | 256/512/1024/2048      | 30.5 / 21.1 / 28.6 / 22.8 | 0.216 / 0.350 / 0.403 / 0.610 | 5.34 / 5.21 / 5.38 / 5.30 |\n| **TarFlowLM single-dim CDF Flow (Ours)**  | -      | 20.7 | 0.47 | 5.58 |\n| **TarFlowLM multi-dim Rosenblatt Flow (Ours)** | - | 15.1 | 0.68 | 5.13 |\n\nBaseline numbers are from [1] and [2]. For MAUVE, we generated 5,000 samples and used GPT-2 Large as the reference model. Gen PPL was also computed using GPT-2 Large.\nWe will add this table and a detailed discussion of these new results to the experiments section of our paper.\n\n[1] Xu, Minkai, et al. \"Energy-based diffusion language models for text generation.\" arXiv:2410.21357 (2024).\n[2] Wang, Guanghan, et al. \"Remasking discrete diffusion models with inference-time scaling.\" arXiv:2503.00307 (2025).\n\n### **3. Ablation Study on Flexible Patch Sizes**\n\n> *“There is a fair amount of discussion of the patching, but limited results are presented related to that design choice. It would be interesting to actually see that ablated to see if reducing the patch size is feasible in practice or whether that significantly degrades quality.”*\n\nThe reviewer raises an excellent point about substantiating our claims of flexibility. To investigate the effect of patch size, we performed a new ablation study on the text8 dataset, evaluating the likelihood (BPC) for our `Mix-1` model.\n\nFirst, we trained models with a fixed depth (3 flow layers) but varied the patch size:\n\n| Patch Size | BPC ↓ |\n| :--- | :---: |\n| 1 | 1.37 |\n| 2 | 1.40 |\n| 3 | 1.67 |\n| 4 | 1.89 |\n\nAs shown, with a fixed model depth, increasing the patch size degrades performance. This is intuitive, as larger patches require the model to capture more complex dependencies between tokens within the patch, which are only modeled by subsequent flow layers and channel mixing.\n\nOur framework is designed to handle this by increasing model depth, allowing stacked flow layers to progressively model these dependencies. To confirm this, we fixed the patch size to 4 and varied the number of flow layers:\n\n| AR Flow Layers (Patch Size=4) | BPC ↓ |\n| :--- | :---: |\n| 3 | 1.89 |\n| 4 | 1.81 |\n| 8 | 1.68 |\n\nThese results confirm our hypothesis: increasing the number of flow layers significantly improves performance for a larger patch size, closing much of the gap with smaller patches. This demonstrates a clear and practical trade-off between patch size and model depth, a key feature of our model's flexibility. We will add this informative ablation study to the appendix.\n\n### **4. Writing Improvements**\n\n> *“There were some parts of the writing that can be improved. For example, missing capitalization on line 83, line 965 is not grammatical.”*\n\nThank you for pointing out these issues. We have carefully proofread the manuscript and will correct these and other typos in the final version to improve clarity and presentation.\n\n***\n\nOnce again, we thank the reviewer for their valuable and constructive feedback, which has helped us significantly improve the paper. We hope our responses and the new experiments have addressed all concerns."}, {"review_id": "dtRSmViGU4", "reviewer": "Reviewer_Yc2p", "summary": "the paper proposed a novel framework for language modelling using autoregressive normalizing flows instead of classical transformer architecture.\nAs such the language modelling occurs in a continues stacked multi-variate probability space , rather than relying on classical log-prob sampling. \nThe authors conduct comprehensive language modelling experiments / comparison.", "questions": "please address weaknesses.\nwilling to increase score if code is made public.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "- S1 originality - it is the first paper i've seen that successfully applies normalizing flows as a language modelling framework to that extent.\n- S2 the paper is well explained and can be followed. very thorough work, well explained limitations and discussions etc\n\nMy major issue:\n- W1 issue i have is that code is not made available - as such results are not reproducible at all. Will this be done at some point?\n\nMinors:\n- W2 i find it a bit odd seeing that undiscussed Figures and Tables in main corpus\n- (W3 wall-times might be poor but strategies like patching are already introduced to cope that - still not sure how far they got)\n- (W4 no downstream / prior-ed generative examples)", "was_revised": true, "final_justification": "the authors justified they'd publish the code.\ngiven the novelty, and potential impact to the research community, i find this very much publish-worthy, albeit the paper in general is missing 'some standards', that however may not yet be as relevant in the current state of this research direction.\nas such i'd grade it not 'a flawless strong accept' but indeed a solid accept.", "rebuttal": "We are very grateful for the reviewer's positive evaluation and constructive feedback. We are delighted that the reviewer recognized the originality of our work (S1) and appreciated the paper's clarity and thoroughness (S2). We address the minor weaknesses below.\n\n# W2\n\nWe appreciate the reviewer's attention for detail. The content of these figures is discussed in the main text, but we will improve the explicit cross-referencing in the new version. Specifically, Figure 8 is discussed in Lines 281-287, and Figure 9 is discussed in Lines 263-277.\nWe will add explicit references to connect the text more directly to improve readability.\n\n# W3\n\nPlease see our response to pnYG Q3 for a detailed discussion.\n\n# W4\n\nPlease see our response to pnYG Q1 for a detailed discussion.\n\n# W1\n\nWe will make the training code available after going through internal review process. As a start we include the core MoG coupling transform here below:\n\n```python\nclass MixtureOfGaussianTransform(nn.Module):\n  def __init__(self, *, c_in = -1, num_mixtures = 10):\n    super().__init__()\n    self.num_mixtures = num_mixtures\n    self.c_in = c_in\n    if c_in > 0:\n      self.scaling_factor = nn.Parameter(torch.zeros(self.c_in))\n      self.mixture_scaling_factor = nn.Parameter(\n        torch.zeros(self.c_in, self.num_mixtures)\n      )\n\n  @staticmethod\n  def _compute_mixture_cdf(\n    x,\n    log_pi,\n    mu,\n    sigma,\n  ):\n    x_orig_dtype = x.dtype\n    if x.dtype != torch.float64:\n      x = x.to(torch.float64)\n    if log_pi.dtype != torch.float64:\n      log_pi = log_pi.to(torch.float64)\n    if mu.dtype != torch.float64:\n      mu = mu.to(torch.float64)\n    if sigma.dtype != torch.float64:\n      sigma = sigma.to(torch.float64)\n\n    pi = torch.softmax(log_pi, dim=-1)\n    sigma = sigma.clamp(min=1e-7)  # Safety clamp\n    z = (x.unsqueeze(-1) - mu) / sigma\n    cdf_standard = 0.5 * (1.0 + torch.erf(z / math.sqrt(2.0)))\n    cdf_mixture = (pi * cdf_standard).sum(dim=-1)\n    # Clamp and return in original dtype (or double if already double)\n    return cdf_mixture.clamp(0.0, 1.0).to(x_orig_dtype)\n\n  @staticmethod\n  def _compute_mixture_pdf(\n    x,\n    log_pi,\n    mu,\n    sigma,\n  ):\n    x_orig_dtype = x.dtype\n    if x.dtype != torch.float64:\n      x = x.to(torch.float64)\n    if log_pi.dtype != torch.float64:\n      log_pi = log_pi.to(torch.float64)\n    if mu.dtype != torch.float64:\n      mu = mu.to(torch.float64)\n    if sigma.dtype != torch.float64:\n      sigma = sigma.to(torch.float64)\n\n    pi = torch.softmax(log_pi, dim=-1)\n    sigma = sigma.clamp(min=1e-7)  # Safety clamp\n    z = (x.unsqueeze(-1) - mu) / sigma\n    log_pdf_standard = -0.5 * (z**2 + math.log(2 * math.pi))\n    pdf_standard = torch.exp(log_pdf_standard)\n    pdf_component = pdf_standard / sigma\n    pdf_mixture = (pi * pdf_component).sum(dim=-1)\n    # Clamp and return in original dtype (or double if already double)\n    return pdf_mixture.clamp(min=1e-9).to(x_orig_dtype)\n\n  @staticmethod\n  def _invert_mixture_cdf(\n    p,\n    log_pi,\n    mu,\n    sigma,\n    max_iter = 25,\n    tol = 1e-6,\n    eps = 1e-6,\n  ):\n    p_orig_dtype = p.dtype\n    if p.dtype != torch.float64:\n      p = p.to(torch.float64)\n    if log_pi.dtype != torch.float64:\n      log_pi = log_pi.to(torch.float64)\n    if mu.dtype != torch.float64:\n      mu = mu.to(torch.float64)\n    if sigma.dtype != torch.float64:\n      sigma = sigma.to(torch.float64)\n\n    pi = torch.softmax(log_pi, dim=-1)\n    sigma = sigma.clamp(min=1e-7)  # Original safety clamp\n\n    mu_mix = (pi * mu).sum(-1)\n    var_mix = (pi * (sigma**2 + mu**2)).sum(-1) - mu_mix**2\n    sigma_mix = var_mix.clamp(min=1e-5).sqrt()  # Original clamp value\n\n    p_clamped = p.clamp(eps, 1.0 - eps)  # Use provided eps\n    z_approx = math.sqrt(2.0) * torch.erfinv(2 * p_clamped - 1)\n    x = mu_mix + sigma_mix * z_approx\n\n    bracket_width_factor = 5.0\n    initial_bracket_width = bracket_width_factor * sigma_mix\n    initial_bracket_width = initial_bracket_width.clamp(min=1e-3)\n    a = mu_mix - initial_bracket_width\n    b = mu_mix + initial_bracket_width\n\n    newton_damping_width = 5.0 * sigma_mix\n    newton_damping_width = newton_damping_width.clamp(min=1e-3)\n\n    for _ in range(max_iter):\n      x = x.clamp(min=a, max=b)\n\n      F_x = MixtureOfGaussianTransform._compute_mixture_cdf(\n        x,\n        log_pi,\n        mu,\n        sigma,\n      )\n      m_x = MixtureOfGaussianTransform._compute_mixture_pdf(\n        x,\n        log_pi,\n        mu,\n        sigma,\n      )\n      F_x = F_x.to(torch.float64)\n      m_x = m_x.to(torch.float64)\n      err = F_x - p  # p is already double\n\n      if err.abs().max() < tol:\n        break\n\n      m_x_safe = m_x.clamp(min=1e-9)\n      delta = err / m_x_safe\n      delta = delta.clamp(min=-newton_damping_width, max=newton_damping_width)\n      x_newton = x - delta\n\n      F_newton = MixtureOfGaussianTransform._compute_mixture_cdf(\n        x_newton,\n        log_pi,\n        mu,\n        sigma,  # Pass double tensors\n      )\n      F_newton = F_newton.to(torch.float64)\n      err_newton = F_newton - p  # p is already double\n      improved = err_newton.abs() <= err.abs()\n\n      x_bisection = 0.5 * (a + b)\n      x_next = torch.where(improved, x_newton, x_bisection)\n\n      F_current = MixtureOfGaussianTransform._compute_mixture_cdf(\n        x_next,\n        log_pi,\n        mu,\n        sigma,  # Pass double tensors\n      )\n      F_current = F_current.to(torch.float64)\n      a = torch.where(F_current < p, x_next, a)\n      b = torch.where(F_current >= p, x_next, b)\n      x = x_next\n\n    return x.to(p_orig_dtype)\n\n  @staticmethod\n  def fn_extract_params(\n    nn_out,\n    num_mixtures,\n  ):\n    param_num = num_mixtures * 3\n    nn_out = nn_out.reshape(\n      nn_out.shape[:-1] + (nn_out.shape[-1] // param_num, param_num)\n    )\n    log_pi = nn_out[..., 0:num_mixtures]\n    mu = nn_out[..., num_mixtures : 2 * num_mixtures]\n    log_var = nn_out[..., 2 * num_mixtures : 3 * num_mixtures]\n    # sigma = torch.nn.functional.softplus(log_sigma)\n    return NestedMap(\n      log_pi=log_pi,\n      mu=mu,\n      log_var=log_var,\n    )\n\n  @capture_metrics\n  @staticmethod\n  def fn_x_to_z_transform(\n    x,\n    *,\n    fn_model_fprop = None,\n    num_mixtures,\n    scaling_factor,\n    mixture_scaling_factor,\n    log_pi = None,\n    mu = None,\n    sigma = None,\n    should_flip = False,\n    eps = 1e-6,\n  ):\n    x_orig_dtype = x.dtype\n    x = x.to(torch.float64)\n\n    flip_dim = 1\n    if x.dim() > 2:\n      flip_dim = 1\n\n    if should_flip:\n      if x.dim() <= flip_dim:\n        raise ValueError(\n          \"\"\n        )\n      x = torch.flip(x, dims=[flip_dim])\n\n    if log_pi is None or mu is None or sigma is None:\n      if fn_model_fprop is not None:\n        nn_out = fn_model_fprop(\n          x.to(x_orig_dtype)\n        )  # Pass original dtype if model expects it\n        params = MixtureOfGaussianTransform.fn_extract_params(\n          nn_out, num_mixtures\n        )\n        log_pi = params.log_pi.to(torch.float64)\n        mu = params.mu.to(torch.float64)\n        log_var = params.log_var.double()\n        sigma = (0.5 * log_var).exp()\n      else:\n        raise ValueError(\"\")\n    else:\n      log_pi = log_pi.to(torch.float64)\n      mu = mu.to(torch.float64)\n      sigma = sigma.to(torch.float64)\n\n    p = MixtureOfGaussianTransform._compute_mixture_cdf(x, log_pi, mu, sigma)\n    p = p.to(torch.float64)\n    p_clamped = p.clamp(eps, 1.0 - eps)\n    z = math.sqrt(2.0) * torch.erfinv(2 * p_clamped - 1)\n\n\n    pdf_mix = MixtureOfGaussianTransform._compute_mixture_pdf(\n      x, log_pi, mu, sigma\n    )\n    pdf_mix = pdf_mix.to(torch.float64)\n    log_pdf_std_normal = -0.5 * (z**2 + math.log(2 * math.pi))\n    ldj = (\n      torch.log(pdf_mix.clamp(min=1e-9)) - log_pdf_std_normal\n    )\n\n    if should_flip:\n      if x.dim() > flip_dim:\n        z = torch.flip(z, dims=[flip_dim])\n        ldj = torch.flip(ldj, dims=[flip_dim])\n\n    return z.to(x_orig_dtype), ldj.to(x_orig_dtype)\n```\n---\n\nWe hope these clarifications fully address the reviewer's concerns. We thank the reviewer again for their time and valuable feedback."}]

Author rebuttal:

Meta-review: ***Summary***

This paper introduces TarFlowLM, a language model built on a transformer backbone which does language modeling using autoregressive normalizing flows.  Much like a diffusion model, TarFlowLM can capture bidirectional dependencies in the text.  Although each layer is autoregressive, the layers alternate in direction between left-to-right, and right-to-left.  Unlike discrete diffusion models, TarFlowLM operates entirely in continuous space, due to the fact that it is a normalizing flow.

Background: normalizing flows are a technique for establishing a bidirectional map from one probability distribution to another.  The underlying theory requires an invertible differentiable function, and a Jacobian with a well-defined determinant, which restricts the set of neural network architectures which can be used.  (The authors define and compare two different mechanisms for establishing these properties.)  Normalizing flows have fallen out of popularity in recent years due to the success of diffusion models for images, and autoregressive models for text, but they have interesting mathematical properties that other classes of models lack.  

The authors compare TarFlowLM against both traditional AR models, and recent diffusion models; the technique performs well when compared against diffusion baselines.  

***Meta-review***

The reviews of this paper were somewhat mixed, with ratings of (3, 4, 5, 5).   The three positive reviewers all agreed that the paper was very original, a valuable theoretical contribution, and well-written overall.  There was also agreement that the main weakness of the paper was the lack of benchmarks beyond perplexity or likelihood.  

The one negative reviewer (J8wg) did not seem to be very familiar with the theory of normalizing flows.  That is entirely unsurprising, since normalizing flows are a somewhat niche technique that is not widely used.  The reviewer expressed confusion as to the difference between normalizing flows and a standard AR transformer (side note: both training and inference are utterly different), and were concerned that TarFlowLM doesn't beat the AR baseline.  As reviewer GUST mentions, however, bidirectional LM models have historically failed to match AR baselines, which is why SOTA models are still AR.  However, performance is strong when compared to text diffusion, which is the most important point of comparison.

I agree with the three positive reviewers.  This paper should be accepted, and I am willing to champion it.  My main argument in favor of acceptance is novelty.  As far as I know, (and as reviewers GUST and Yc2p point out), nobody has ever built a transformer language model out of normalizing flows before.  This paper is thus somewhat similar to the very first papers on text diffusion -- it is applying an entirely new technique to language modeling.  Early papers on text diffusion had very weak results, but those initial papers were still groundbreaking, because they established a starting point that later work could build on, and text diffusion is now a popular and active area of research.  Although it has some weaknesses, this paper actually establishes a much stronger starting point than early work on text diffusion, IMO.  

Normalizing flows have some very nice mathematical properties, so I expect to see further work in this space.  See, e.g., JetFormer: An Autoregressive Generative Model of Raw Images and Text, Michael Tschannen et. al. https://arxiv.org/abs/2411.19722.  

Selected quotes from reviewers:

**Reviewer J8wg (rating 3, confidence 3)**: 
* "I don't fully understand the difference between a standard transformer model... and the current framework."
* "Given a token context, how is inference performed?"
* " the results showed in this Table show that TarFlowLM is outperformed by the standard autoregressive transformer baseline."
* "No discussion/results on the generation speed difference against other non transformer AR baselines."

**Reviewer pnYG (rating 4)**: 
* "The paper is generally well-written, well-motivated, and easy to follow."
* "One of the key strengths of the paper is its theoretical results..."
* "One of the critical weaknesses of the work is the lack of comparison on common benchmarks on language modeling. [instead of just likelihood/perplexity.]"

**Reviewer GUST (rating 5)**: 
* "This is an interesting approach to language generation that I am not aware of being explored recently."
* "It is usually the case with these kinds of work that there is a persistent gap with the simpler autoregressive approach... I don’t consider this a strong weakness."
* "The performance of their approach is strong compared to recent popular discrete diffusion models (e.g. SEDD, MDLM)."
* "The biggest limitation is the lack of metrics being reported for actual language generation. Reporting likelihoods is great, but..."

**Reviewer Yc2p (rating 5)**: 
* "Originality - it is the first paper I've seen that successfully applies normalizing flows as a language modelling framework to that extent."
* "The paper is well explained and can be followed. very thorough work, well explained limitations and discussions etc."
Decision: Accept (poster)