Augmented Coupling Flows

• Augmented coupling flows are a family of methods that modify standard coupling transforms by augmenting the state or conditioning to improve expressivity and achieve richer transport dynamics.
• Techniques such as Augmented Normalizing Flows, Augmented Bridge Matching, and distilled couplings integrate auxiliary variables or teacher-induced latent codes to balance performance and computational trade-offs.
• Applications span density estimation, symmetry-preserving physical modeling, and video compression, demonstrating empirical gains in metrics like FID and faster sampling speeds.

Searching arXiv for the cited papers to ground the article in the relevant literature. arXiv search: (Berthelot et al., 9 Mar 2026) arXiv search: (Huang et al., 2020) Editor's term augmented coupling flows denotes a family of coupling-based transport and generative constructions in which the basic coupling mechanism is strengthened by augmenting the effective state, conditioning variables, coordinate representation, or endpoint pairing. Across the recent literature, augmentation appears in several distinct but related forms: auxiliary latent variables added to a normalizing flow’s state space, source-conditioned bridge dynamics that preserve empirical pairings, extra dimensions that expand the approximation class of coupling flows, symmetry-preserving auxiliary copies for equivariant modeling in Cartesian coordinates, conditional side information in compression-oriented flows, and distilled endpoint assignments for flow matching (Huang et al., 2020, Bortoli et al., 2023, Lyu et al., 2022, Midgley et al., 2023, Chen et al., 2022, Berthelot et al., 9 Mar 2026). The common theme is that the coupling is no longer treated as a fixed implementation primitive; it becomes the object through which expressivity, symmetry, identifiability, or optimization geometry is improved.

1. Coupling flows and the meaning of augmentation

A standard coupling transform splits an input into two parts, leaves one block unchanged, and transforms the other block conditionally. In the formulation summarized by Cubic-Spline Flows, one keeps $\mathbf{x}_{1:d-1}$ fixed, computes parameters $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$, and applies elementwise transforms $y_i=g_{\theta_i}(x_i)$ for $i=d,\dots,D$; the Jacobian is lower triangular, so the determinant is the product of diagonal terms. That computational structure explains why coupling flows are fast in both forward and inverse directions, but also why plain affine couplings can be less expressive than autoregressive flows (Durkan et al., 2019).

The literature uses augmentation in more than one technical sense. In some papers, augmentation literally enlarges the state from $x$ to $(x,e)$ or from $x$ to $(x,a)$. In others, it means carrying forward the source sample $X_0$ inside the learned drift, or adding one extra dimension to obtain stronger approximation results, or replacing independent noise/data pairing with a learned endpoint assignment. This suggests that “augmentation” is best understood as a structural modification of the coupling mechanism rather than a single architectural recipe.

Setting	Augmentation mechanism	Stated role
ANF	Pair each observed sample $x$ with auxiliary $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$0 and learn $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$1	Improve expressivity and bypass some topological constraints
AugBM	Condition drift on $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$2 rather than $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$3 alone	Preserve empirical source–target coupling
Para-CFlows	Lift to one extra dimension and preserve context coordinates	Achieve $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$4 universality over parametric diffeomorphisms
SE(3) equivariant flows	Add auxiliary coordinates $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$5 and alternate updates of $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$6 and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$7	Preserve $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$8 and permutation equivariance
NFM	Replace Gaussian noise endpoint with teacher-produced latent code	Distill a structured coupling for flow matching

A conceptually earlier use of augmentation appears in stochastic network theory, where flow counters are appended to the population state to form an augmented state–flow process $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$9. There the augmentation is redundant in the probabilistic sense, because $y_i=g_{\theta_i}(x_i)$0 is determined by the path of $y_i=g_{\theta_i}(x_i)$1, but it makes possible Markov couplings that order throughputs even when direct population ordering fails (Leskelä, 2014). Although this setting is not generative modeling, it illustrates a recurring principle: augmentation often restores a useful invariant or order relation that is inaccessible in the original state space.

2. Auxiliary-variable normalizing flows

Augmented Normalizing Flows (ANFs) introduce one of the clearest enlarged-state formulations. Instead of fitting a bijection on $y_i=g_{\theta_i}(x_i)$2 alone, the model pairs each observed sample $y_i=g_{\theta_i}(x_i)$3 with an independent auxiliary variable $y_i=g_{\theta_i}(x_i)$4 and learns a joint density

$y_i=g_{\theta_i}(x_i)$5

The architecture is built from coupling-layer-style transforms on $y_i=g_{\theta_i}(x_i)$6, with an “autoencoding” pair of affine couplings:

$y_i=g_{\theta_i}(x_i)$7

$y_i=g_{\theta_i}(x_i)$8

Stacking alternating encode and decode maps yields the full invertible transform $y_i=g_{\theta_i}(x_i)$9 (Huang et al., 2020).

Training uses Augmented Maximum Likelihood Estimation (AMLE),

$i=d,\dots,D$0

which is a variational lower bound on $i=d,\dots,D$1. The gap between the marginal likelihood and the bound is the “augmentation gap,”

$i=d,\dots,D$2

The paper reports that this gap is small, around $i=d,\dots,D$3 bits/dim, and gives benchmark results of $i=d,\dots,D$4 bits/dim on MNIST, $i=d,\dots,D$5 on CIFAR-10, $i=d,\dots,D$6 on ImageNet $i=d,\dots,D$7, $i=d,\dots,D$8 on ImageNet $i=d,\dots,D$9, and $x$0 on CelebA-HQ. On CIFAR-10, ANF attains an Inception Score of $x$1 and FID of $x$2 (Huang et al., 2020).

The significance of ANFs is twofold. First, the auxiliary coordinates let the model move probability mass through extra dimensions, making distributions easier to Gaussianize. Second, the construction explicitly bridges normalizing flows and latent-variable models: the joint flow on $x$3 can be read as an exact-Jacobian counterpart of a variational model, and hierarchical ANFs generalize hierarchical VAEs and auxiliary-variable variational methods. The stated tradeoff is also explicit: relative to standard coupling flows, ANFs incur extra computation because each step updates both $x$4 and $x$5 and alternates encode/decode passes (Huang et al., 2020).

3. Transport and approximation theory

Two different theoretical programs motivate augmented coupling flows. One is transport-theoretic and is developed in ANF. The other is approximation-theoretic and is developed for Para-CFlows.

For ANFs, the main theorem is that augmented coupling updates can approximate a Hamiltonian ODE as a universal transport map. The paper defines a Hamiltonian system

$x$6

with schedules

$x$7

and shows convergence in distribution to $x$8 and $x$9. Under a smoothness/Lipschitz assumption, a sequence of additive ANFs can approximate these trajectories arbitrarily well. The conceptual consequence is that sufficiently deep additive coupling updates on an augmented space can transport arbitrary data to a Gaussian prior while collapsing the auxiliary variable to a point mass (Huang et al., 2020).

Para-CFlows address a different question: which diffeomorphisms can coupling flows approximate, and with what regularity? The paper proves that if the internal function class can approximate certain single-coordinate transforms, then affine coupling layers plus invertible linear transforms achieve $(x,e)$0-universality for compactly supported diffeomorphisms. In the lifted construction, one approximates a target map $(x,e)$1 by embedding into one extra dimension and projecting back:

$(x,e)$2

The same logic extends to parameter-preserving families $(x,e)$3, yielding universality results for parametric coupling flows, or Para-CFlows (Lyu et al., 2022).

Taken together, these results show that augmentation serves at least two rigorously different roles. In ANF it enlarges the dynamical state so that a richer transport can be realized. In Para-CFlows it enlarges the ambient dimension so that coupling architectures can approximate arbitrary compactly supported diffeomorphisms in $(x,e)$4-norm. A plausible implication is that augmentation is not merely an empirical heuristic for flow architectures; it is also a mechanism for changing the representable transport class itself.

4. Coupling preservation and non-Markovian augmentation

Augmented Bridge Matching (AugBM) addresses a different limitation of standard flow and bridge matching: preservation of marginals does not imply preservation of the original empirical pairing. In the paired setting, training data are distributed according to a joint law $(x,e)$5 over $(x,e)$6. Standard bridge matching constructs a bridge process, such as the Brownian bridge

$(x,e)$7

and then learns a Markovian drift depending only on $(x,e)$8. The paper proves that this Markovian projection preserves the original coupling only when the training coupling is itself the static Schrödinger bridge, equivalently entropic optimal transport for Brownian reference dynamics (Bortoli et al., 2023).

AugBM restores the coupling by augmenting the dynamics with the initial sample. The augmented SDE is

$(x,e)$9

and the neural approximation uses a regression objective on $x$0. Because the learned drift depends on $x$1 rather than $x$2 alone, the process is no longer Markovian, but the original joint coupling $x$3 is preserved by construction (Bortoli et al., 2023).

The paper places this in direct relation to diffusion models, deterministic flow matching, and Denoising Diffusion Bridge Models. When $x$4, flow matching is recovered as the deterministic limit. With $x$5, the bridge structure tightens the fixed-point set, but only the augmented formulation generically preserves empirical pairings. In image-translation tasks such as edges2shoes, edges2handbags, and ImageNet colorization, the reported outcome is that standard bridge matching or I$x$6SB can match marginals while losing source–target correspondence, whereas AugBM preserves the coupling and yields substantially better coupling FID, especially at high numbers of function evaluations (Bortoli et al., 2023).

This section clarifies an important misconception. Better marginal matching is not the same as better coupling recovery. AugBM shows that, for paired transport problems, the coupling information may need to be carried explicitly as part of the state.

5. Distilled couplings in flow matching

“The Coupling Within: Flow Matching via Distilled Normalizing Flows” reframes the coupling in flow matching as the central object that determines the learned velocity field, the variance of the regression target, and the straightness of the resulting ODE trajectories. Standard Flow Matching (FM) uses an independent coupling: each data point $x$7 is paired with an independent Gaussian $x$8, giving

$x$9

The student is trained by the usual regression loss

$(x,a)$0

and sampled by backward integration from $(x,a)$1 (Berthelot et al., 9 Mar 2026).

Normalized Flow Matching (NFM) replaces the Gaussian endpoint with a teacher-produced latent code from a pretrained normalizing flow teacher. After perturbing the data by $(x,a)$2, the latent is

$(x,a)$3

with $(x,a)$4 chosen so that teacher codes are approximately unit-variance. The student is trained by

$(x,a)$5

The resulting coupling is model-induced rather than transport-optimal: it reflects the teacher NF’s learned bijection rather than a globally optimal earth-mover assignment. The paper describes this as a “quasi-deterministic bijection,” because the teacher is invertible and trained by maximum likelihood, but finite capacity and the small perturbation $(x,a)$6 prevent exact determinism (Berthelot et al., 9 Mar 2026).

The empirical results are unusually strong at low inference budgets. On ImageNet64 at $(x,a)$7 NFE, FID is reported as $(x,a)$8 for FM, $(x,a)$9 for SD-FM, and $X_0$0 for NFM; at $X_0$1 NFE the values are $X_0$2, $X_0$3, and $X_0$4; at $X_0$5 NFE they are $X_0$6, $X_0$7, and $X_0$8. With a better solver configuration, the $X_0$9-NFE ImageNet64 result reaches FID $x$0 for NFM versus $x$1 for the TarFlow teacher itself. The latency table gives TarFlow at $x$2 units versus NFM(31) at $x$3, corresponding to a $x$4 speedup; NFM(15) reaches $x$5, and NFM(7) $x$6. On ImageNet256, the paper reports $x$7 for FM versus $x$8 for NFM at the best $x$9-NFE point, and at $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$00 NFE gives $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$01 for FM versus $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$02 for NFM. The curvature metric also improves: on ImageNet64 with Heun($\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$03) at $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$04 NFE, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$05 is $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$06 for FM, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$07 for SD-FM, and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$08 for NFM; with Euler($\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$09) at $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$10 NFE, the values are $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$11, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$12, and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$13 (Berthelot et al., 9 Mar 2026).

Several further observations constrain how NFM should be interpreted. The teacher’s latent geometry is not locally faithful: on ImageNet64 with $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$14, the same-image/different-noise latent distance is $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$15, while different-image/same-noise is $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$16. Nevertheless, the pseudo-Gaussian latent space improves FM optimization. Teacher quality also matters: a teacher with NLL $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$17 yields teacher/student FIDs of $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$18, whereas a weaker teacher with NLL $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$19 gives $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$20; $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$21 gives the best reported teacher/student pair, while $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$22 degrades both. The authors explicitly suggest that NFM and semidiscrete OT couplings may be complementary, with OT mapping true Gaussian noise to the teacher’s pseudo-Gaussian latent space (Berthelot et al., 9 Mar 2026).

6. Symmetry-aware and application-specific augmented couplings

In physical modeling, augmentation is used to make coupling flows compatible with symmetries that ordinary coordinate splits destroy. The SE(3) equivariant augmented coupling flow operates on centered Cartesian coordinates $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$23 and introduces auxiliary variables $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$24 that transform under the same $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$25 action as the observations. The key lift is the shift-CoM map

$\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$26

after which the model alternates coupling updates of $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$27 conditioned on $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$28 and of $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$29 conditioned on $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$30. Within each core transform, an equivariant network predicts reference vectors and invariant parameters, the coordinates are projected into a learned invariant basis, an ordinary transform such as an affine map or monotonic rational-quadratic spline is applied, and the result is mapped back to Cartesian space. This preserves fast sampling and tractable density evaluation while maintaining $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$31 and permutation equivariance (Midgley et al., 2023).

The reported performance supports the claim that the symmetry-preserving augmentation is practical. On DW4, LJ13, and QM9-positional, the augmented flow is competitive with equivariant CNFs, flow matching, and diffusion models. Sampling and density evaluation take about $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$32 seconds on an A100 GPU, whereas the CNF and diffusion baselines take about $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$33 and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$34 seconds on average for sampling. On alanine dipeptide, the paper reports KLDs around $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$35 to $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$36 with NLL around $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$37, and shows reweighting can improve the Ramachandran KLD from $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$38 to $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$39 at $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$40K. The same architecture can also be trained from unnormalized energies using FAB, which the paper presents as a consequence of fast exact sampling and density evaluation (Midgley et al., 2023).

In learned video compression, B-CANF adapts augmented normalizing flows to conditional coding of B-frames. The codec models a target frame $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$41 conditioned on a motion-compensated prediction $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$42 using conditional additive coupling transforms, hierarchical hyperpriors, frame-type adaptive coding, and B*-frames. For the first transform,

$\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$43

$\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$44

The loss is

$\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$45

with $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$46, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$47, and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$48 in the experiments. In the ablation study, the full model with conditional motion and conditional inter-frame coding reports BD-rate values of $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$49, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$50, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$51, and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$52 on UVG, MCL-JCV, HEVC-B, and CLIC’22, whereas switching both codecs to residual coding yields $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$53, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$54, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$55, and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$56. A 1-step B-CANF gives $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$57, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$58, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$59, and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$60, while the 2-step model gives $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$61, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$62, $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$63, and $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$64 (Chen et al., 2022).

These two examples show that augmentation is not restricted to generic density estimation. In one case it is the mechanism by which exact equivariance becomes compatible with coupling-based sampling and likelihood computation. In the other it turns a flow backbone into a conditional rate–distortion system with hierarchical temporal structure.

7. Tradeoffs, misconceptions, and open combinations

The literature does not present augmentation as cost-free. ANFs gain expressivity but require extra computation for the auxiliary pathway and alternating encode/decode passes (Huang et al., 2020). AugBM preserves the training coupling by conditioning on $\boldsymbol{\theta}=\text{NN}(\mathbf{x}_{1:d-1})$65, but the process is non-Markovian and can make optimization harder and increase variance, especially when the coupling is highly entropic (Bortoli et al., 2023). The SE(3) equivariant augmented flow is much faster than CNFs and diffusion models, yet the authors state that it is still more expensive than internal-coordinate flows and can be numerically delicate (Midgley et al., 2023). In Cubic-Spline Flows, the cubic variant improves expressivity while preserving exact one-pass inversion, but the quadratic baseline is reported as somewhat faster and less numerically fragile (Durkan et al., 2019). In NFM, the distilled coupling is only as good as the teacher’s latent modeling quality, and the teacher-induced map is explicitly described as likely suboptimal in a geometric sense (Berthelot et al., 9 Mar 2026).

A recurring misconception is that augmentation always means adding latent dimensions. The surveyed papers show otherwise. AugBM augments the learned velocity field with the initial sample rather than enlarging the ambient space. NFM augments the endpoint pairing used in flow matching by replacing independent Gaussian noise with a teacher-induced latent code. The queueing-network work augments the state with flow counters to recover order relations on throughputs rather than to increase representational capacity (Leskelä, 2014).

Another misconception is that augmentation necessarily enforces faithful local geometry. NFM explicitly reports that TarFlow’s latent neighbors are not input-space neighbors, yet the latent space still improves flow matching optimization (Berthelot et al., 9 Mar 2026). This suggests that the operational value of augmentation often lies in the learned transport geometry it induces for optimization or invariance, not in local isometry.

Several papers also point toward compositional future directions. NFM explicitly proposes combining teacher-induced pseudo-Gaussian latents with semidiscrete OT couplings (Berthelot et al., 9 Mar 2026). Para-CFlows show that augmentation by one extra dimension can be used to strengthen approximation guarantees for whole families of diffeomorphisms (Lyu et al., 2022). ANFs, conditional augmented flows, and symmetry-aware augmented flows collectively indicate that auxiliary variables can be used not only to improve expressivity, but also to encode task structure, side information, or symmetry constraints (Huang et al., 2020, Chen et al., 2022, Midgley et al., 2023).

Within this broader view, augmented coupling flows are best understood not as a single model class but as a design principle: modify the coupling mechanism so that the resulting transport better matches the geometry, invariances, conditioning structure, or pairing information of the problem.