What it actually does
Start from a simple distribution you can sample and evaluate — a standard Gaussian — and push it through an invertible function. The change-of-variables formula says exactly what happens to the density: it is the base density at the preimage, corrected by how much the transformation stretched or compressed space there. So if you can build a flexible invertible function whose local volume change you can compute cheaply, you get an exact likelihood for free. The entire field is a search for such functions, and every architecture in it is a different answer to the same question: how do you keep a Jacobian determinant tractable without making the map trivial?
- The data distribution is continuous, or has been dequantized so that a continuous density is meaningful
- The data lies in a space of the same dimension as the latent, since the map is invertible and cannot compress
- The target distribution is reachable from the base distribution by a composition of the chosen bijections
- Likelihood is a meaningful measure of typicality for this data — false in high dimensions, where density and typicality diverge
- Preprocessing is identical between training and evaluation, since likelihood is not invariant to it
Likelihood
Read the two terms as separate jobs. The first asks whether the latent code lands somewhere the base distribution considers likely; the second corrects for the fact that the transformation changed how much space a region occupies. Without the correction the model could achieve arbitrarily high likelihood by compressing all its inputs into a tiny region of latent space, which is the degenerate solution the determinant exists to prevent. Two properties follow that distinguish this from everything nearby. The likelihood is exact — no bound, no sampling estimate, no variational gap — so model comparison actually means something. And because determinants multiply while their logarithms add, composing many layers costs nothing extra in this term, which is why flows are built deep and narrow.
- the flow: an invertible map from data space to latent space, composed of K simpler bijections
- the base distribution, chosen so it can be both sampled and evaluated in closed form
- the Jacobian of the k-th transformation; the whole design problem is making its determinant cheap
- the volume correction — positive where the map compresses, negative where it expands
- the number of composed layers; determinants add in log space, so depth is free in this term
How it is minimized
Adam on the exact negative log-likelihood, with the architecture constrained so the Jacobian determinant is triangular
- scale and translation networks — arbitrarily complex, because they are never inverted
- the frozen half; it conditions the transformation of the other half and passes through unchanged
- an elementwise affine map whose Jacobian is diagonal on the transformed block
- the log-determinant: a sum, not a determinant computation, because the Jacobian is triangular
The coupling layer is the idea worth extracting, and it is an unusually clean piece of design. Split the input in half; leave one half alone; transform the other half elementwise using parameters computed from the frozen half. The Jacobian of that map is block-triangular with a diagonal block, so its determinant is the product of the diagonal — a sum in log space, computed in linear time and without ever forming the matrix. The scale and translation networks can be arbitrarily deep and nonlinear because they are only ever evaluated forwards, never inverted, so all the expressive power hides in the part that does not need to be invertible. The cost is that half the dimensions pass through untouched, so the halves must be permuted between layers and depth becomes the only route to expressiveness. That is the trade the whole family makes: every architecture here buys a tractable determinant with a structural restriction, and deeper stacks are how the restriction is paid off.
coupling layers
Depth. The only route to expressiveness, since each layer transforms half the dimensions (8 to 32)
coupling network width
Capacity of the scale and translation networks; free to be large since they are never inverted (128 to 512)
permutation scheme
How the halves are shuffled between layers. Fixed reversal, random, or a learned 1x1 convolution (reverse / random / learned)
scale clamping
Bounds the exponential in the affine map. Without it the log-determinant diverges and training produces NaN (tanh-clamped to ±2 to ±5)
dequantization noise
Continuous flows on discrete data assign unbounded density to the exact values unless noise is added (uniform or variational)
base distribution
Almost always a standard Gaussian; must be both samplable and exactly evaluable (standard normal)
Training is stable and pleasant compared with adversarial alternatives — there is one loss, it is exact, and it goes down. The failures are architectural rather than optimizational. The most common is a diverging log-determinant: the affine scale is exponentiated, so an unclamped scale network can drive it to infinity and the loss becomes NaN in a single step, which is why clamping is mandatory rather than advisable. The second is the expressiveness ceiling — because the map must be invertible and dimension-preserving, the model cannot compress, so a distribution on a low-dimensional manifold embedded in a high-dimensional space is fundamentally awkward for it, and the symptom is a model that fits fine and samples poorly. The third is specific to discrete data: without dequantization the likelihood is unbounded, and the model will happily chase it to infinity rather than learning anything.
O(K · d · W) for a forward pass with K coupling layers and width W, both for density evaluation and for sampling, since both directions cost the same in a coupling flow. Memory is substantially higher than a comparable non-invertible model because the map is dimension-preserving at every layer — nothing is ever compressed, so the activation width never shrinks.
How, where, and why it is used
Time-Series Forecasting
AdaptedCondition the flow on an encoding of the history and let it model the distribution of the next values jointly. Because the likelihood is exact, the model is trained directly on the predictive density rather than on a point loss with an assumed error distribution — and because it is a joint density over the horizon, the correlation structure between horizons is modelled rather than assumed independent.
- Multivariate probabilistic forecasting where the correlation between series matters as much as each marginal
- Forecasting with genuinely non-Gaussian predictive distributions — multimodal, skewed, or heavy-tailed
- Risk applications needing an exact joint density over a horizon rather than per-step quantiles
- Settings where the predictive distribution itself is the deliverable rather than a summary of it
Where the predictive distribution is genuinely non-Gaussian, this models it without assuming a family — which is a real advantage over a quantile forecaster that emits marginals, since a flow gives the joint and the joint is what any decision aggregating over time or across series actually needs. Against it, the objections are substantial and mostly practical: it is far more machinery than the alternatives, the invertibility constraint forces the latent dimension to equal the output dimension so a long multivariate horizon becomes a wide flow, and on most business series the predictive distribution is close enough to Gaussian that a quantile head captures everything that matters for a fraction of the cost. The honest trigger is a genuinely complicated predictive distribution that you can demonstrate is not Gaussian.
- Normalize per series, since the base distribution is standard and an unscaled input makes the flow spend its capacity fixing the scale
- Condition every coupling layer on the history encoding rather than only the first, or the conditioning washes out with depth
- Model the horizon jointly rather than per step, because the joint structure is the only reason to prefer this over quantile regression
- Clamp the affine scale, since a diverging log-determinant produces NaN in a single step
Continuous ranked probability score and joint log-likelihood on held-out windows, with a quantile-regression baseline on the same data — that comparison is what establishes whether the extra machinery bought anything, and it frequently has not. Rolling-origin backtesting throughout, since a random split leaks future structure.
- Assuming a non-Gaussian predictive distribution without demonstrating one, which is the only justification for the cost
- A wide multivariate horizon making the flow enormous, since dimension is preserved at every layer
- Conditioning only the first layer, so the history influence fades with depth
- Comparing likelihoods across models with different dequantization, which makes the numbers incomparable
Anomaly Detection
ViableFit the flow to normal data and score each point by its exact log-likelihood — low likelihood means anomalous. It is the most direct possible expression of density-based detection, with no reconstruction proxy, no bound and no threshold heuristic between the model and the score.
- Tabular and low-dimensional anomaly detection where a genuine density is wanted rather than a proxy
- Scientific and instrument monitoring where the density is itself the object of interest
- Fraud and risk screening on engineered features, where the feature space is modest and well understood
- Cases requiring a calibrated probability rather than an uncalibrated score
On modest-dimensional tabular data this works well and the score means exactly what it says, which is more than any reconstruction-based detector can claim. But this fit is capped at viable for a specific and important reason: on high-dimensional data, particularly images, flows have been shown to assign HIGHER likelihood to out-of-distribution inputs than to their own training data — a model trained on one image dataset can rate a completely different one as more probable. That is not a tuning failure but a property of how likelihood behaves in high dimensions, where typicality and density come apart, and it means "low likelihood implies anomalous" is simply false there. Knowing that result is more valuable than knowing the architecture, and it is why typicality-based tests exist as an alternative to raw likelihood.
- Standardize features, since the base distribution is standard and scale mismatch wastes model capacity
- Prefer a typicality test over a raw likelihood threshold in higher dimensions, where the two genuinely diverge
- Validate on known out-of-distribution data before trusting the score, because the high-dimensional failure is not visible from training
- Dequantize discrete features, or the likelihood is unbounded and the model chases it rather than learning
Precision and recall at a fixed alert budget, and crucially an explicit check that held-out out-of-distribution data scores lower than in-distribution data — that check is the one that catches the high-dimensional likelihood pathology, and it is almost never run. Compare against an isolation forest, which on tabular data is competitive at a small fraction of the cost.
- The high-dimensional likelihood inversion, where out-of-distribution inputs score higher than training data
- Training data containing the anomalies, which the density then rates as likely
- Comparing likelihoods across models with different dequantization schemes, which is meaningless
- Enormous cost against an isolation forest that performs comparably on tabular features
Optimization
AdaptedThe whole family is a constrained design problem: maximize expressiveness subject to the map being invertible and its Jacobian determinant being computable in less than cubic time. Every architecture is a different point on that frontier — coupling layers make the Jacobian triangular by construction, autoregressive flows do the same by ordering, and continuous flows replace the determinant with a trace. Underneath sits a general lesson about building constraints into a hypothesis class rather than enforcing them with a penalty.
- Structural constraint as architecture: invertibility is guaranteed by construction rather than penalized in the loss
- Triangular Jacobians as the device that turns a cubic determinant into a linear sum
- The expressiveness-tractability frontier, where each architecture trades one for the other explicitly
- Change of variables as an exact reparameterization, against variational bounds as an approximate one
Worth studying because it is the clearest demonstration in this reference that a hard constraint can be designed into a model rather than imposed on it — nothing in the training ever checks invertibility, because the architecture makes violation impossible. That is a far stronger guarantee than a penalty, and the pattern transfers. The comparison with the VAE is the other lesson: both model the same kind of thing, and the flow pays for exactness in architectural restriction while the VAE pays for freedom in a variational gap. Neither is free, and seeing where each puts the cost is more instructive than either in isolation.
- Permute or mix dimensions between coupling layers, or the untransformed half never changes and depth buys nothing
- Clamp the scale network output, since the exponential makes an unbounded scale a single-step divergence
- Prefer a learned invertible mixing layer to a fixed permutation where cost permits, since it uses depth more efficiently
- Verify invertibility numerically once as a test — the architecture guarantees it, but an implementation bug does not
Compare architectures at matched parameter count on the same data, reporting both likelihood and sample quality, because they diverge more than expected. Measure the forward and inverse cost separately: some flows are cheap in one direction and expensive in the other, and which direction matters depends entirely on whether you are evaluating densities or sampling.
- Forgetting to permute between layers, which leaves half the dimensions untransformed throughout
- An unclamped scale producing a diverging log-determinant and NaN within a step
- Choosing an architecture whose cheap direction is not the direction the application uses
Computer Vision
AdaptedMulti-scale coupling flows over images, with squeeze operations trading spatial resolution for channels and invertible mixing between couplings. This was a serious line of image generation before diffusion, and the architecture is still used where an exact likelihood is genuinely needed.
- Exact-likelihood image density estimation, where a bound or a score will not do
- Invertible feature extraction, where the ability to recover the input exactly is the requirement
- Lossless compression, where a density model translates directly into a code length
- As an invertible component inside a larger model rather than as a generator in its own right
It remains the right tool when an exact likelihood over images is genuinely required — lossless compression is the clearest case, since code length is log-likelihood and a bound is not good enough. As a generative model for images it has been comprehensively overtaken by diffusion, which produces far better samples at a fraction of the parameter count, and the reason is structural: dimension preservation makes flows enormous, and the invertibility constraint limits what each layer can do. The out-of-distribution likelihood pathology is also at its sharpest here, which is worth knowing before using image likelihoods for anything consequential.
- Dequantize pixel values, or the density on discrete inputs is unbounded and training diverges to it
- Use multi-scale architecture with squeeze operations, since a flat flow at image resolution is unaffordable
- Prefer learned invertible mixing over fixed permutations between couplings, which uses depth far more efficiently
- Report bits per dimension rather than raw likelihood, and state the dequantization used, or the numbers are incomparable
Bits per dimension on held-out data with the dequantization scheme stated explicitly, plus sample quality assessed separately — the two diverge sharply and a flow can have excellent likelihood and poor samples. Compare against a diffusion model at matched compute, which is usually unflattering.
- Skipping dequantization and training toward an unbounded likelihood
- Enormous parameter count from dimension preservation at every layer
- Out-of-distribution inputs scoring higher than training data, which is at its worst in this domain
- Reporting likelihood without the dequantization scheme, making the number meaningless for comparison
Risk & Fraud
ViableModel the joint density of engineered risk features on normal activity, then use the exact likelihood as an unsupervised score alongside a supervised classifier. Because the density is joint, it captures feature combinations that are individually unremarkable and jointly implausible, which is precisely the fraud signature that per-feature rules miss.
- Unsupervised risk scoring as a complement to a supervised model trained on confirmed cases
- Novel-attack detection, where the pattern has not been seen and so cannot have been labelled
- Density-based feature validation, flagging inputs the downstream model was never trained on
- Calibrated risk probabilities where a raw score would not be actionable
The joint-density argument is genuinely strong here: fraud frequently looks normal on every individual feature and implausible in combination, and a density model catches that where thresholds cannot. Feature spaces in risk are also usually modest in dimension, which keeps the high-dimensional likelihood pathology at bay. Against it: adversaries adapt faster than the model is retrained, the density says nothing about whether an unusual pattern is malicious or merely rare, and a gradient-boosted classifier on labelled cases beats it wherever labels exist. It belongs as a complement rather than a replacement.
- Standardize and, where possible, decorrelate features first, since the flow otherwise spends depth learning the covariance
- Fit on a verified-normal period, because a density fitted on contaminated data rates the contamination as likely
- Combine the likelihood with supervised scores rather than using it alone, since it cannot distinguish rare from malicious
- Refit on a rolling window so the notion of normal tracks legitimate behavioural drift
Lift over the supervised model alone at a fixed review budget, measured on confirmed cases with a strictly time-based split. Report how many of the flagged cases were novel patterns rather than ones the supervised model already caught, since that is the only thing this contributes.
- Training on contaminated data, so the fraud is rated as likely
- Treating low likelihood as evidence of malice when it only indicates rarity
- Adversarial drift outpacing the retrain cadence
What it costs to run
Hours to days on a GPU, and considerably more than a comparable VAE at equal quality because the model is dimension-preserving at every layer and depth is the only route to expressiveness. The likelihood objective is exact and stable, so the cost is compute rather than tuning.
Density evaluation and sampling both cost one pass through the stack, which is an unusual property — most generative models are cheap in one direction and expensive in the other. For coupling flows both directions are the same cost, so a single deployment serves both use cases without a second model.
Driven by distribution drift rather than by the model. In anomaly detection the usual uncomfortable question applies: retraining on recent data absorbs any slow-developing anomaly into the new definition of normal.
- Track the log-likelihood distribution on known-normal data rather than the alert count, since a shifting floor means drift and not more anomalies
- Monitor the log-determinant magnitude, because a growing one is the precursor to the divergence that produces NaN
- Validate against held-out out-of-distribution data periodically, since the high-dimensional likelihood pathology is invisible from training alone
- Alert on inputs outside the range the flow was fitted on, where the density is extrapolated rather than estimated
- Likelihoods are only comparable under identical preprocessing and dequantization. Two models with different schemes produce numbers that cannot be compared at all, and this is routinely violated in published results
- The scale network must be clamped. An unbounded exponential makes the log-determinant diverge and produces NaN in one step, with no gradual warning
- On discrete data the continuous density is unbounded without dequantization, so the model trains toward infinity rather than toward the data
- Dimension preservation means memory does not shrink with depth, unlike every encoder in this reference — a flow is much larger than its parameter count suggests
- Higher likelihood does not mean more in-distribution in high dimensions. Do not build an out-of-distribution detector on raw image likelihoods without validating it
Pros and cons, in context
Exact likelihood, not a bound
The defining property and the reason to choose it over a VAE. Model comparison, lossless compression and calibrated density all require the real number, and a variational bound cannot supply any of them.
Invertibility is guaranteed by construction
Nothing in training ever checks it, because the architecture makes violation impossible. A far stronger guarantee than a penalty, and the pattern transfers well beyond this model.
Sampling and density evaluation cost the same
Unusual: most generative models are cheap in one direction and expensive in the other. One deployment serves both, with no second model and no approximation in either direction.
Stable, single-objective training
One exact loss that goes down, with none of the equilibrium problems of adversarial training or the term-balancing of an ELBO. A genuine practical advantage that is easy to undervalue until you have debugged the alternatives.
Dimension preservation makes it large
No layer can compress, so activation width never shrinks and the model is far bigger than its parameter count suggests. This is the structural reason diffusion overtook flows for image generation.
Invertibility caps expressiveness per layer
Each coupling layer transforms only half its input, so depth is the only route to flexibility — and a distribution on a low-dimensional manifold is fundamentally awkward for a map that cannot reduce dimension.
High-dimensional likelihoods can invert
Flows have been shown to assign higher likelihood to out-of-distribution data than to their own training data. Not a tuning failure but a property of density in high dimensions, and it invalidates the obvious anomaly-detection use there.
Diverges hard if the scale is unclamped
The exponential in the affine coupling makes the log-determinant unbounded, so one bad step gives NaN. Trivially fixed and trivially forgotten, and the failure is abrupt rather than gradual.
Likelihoods are not portable
Comparable only under identical preprocessing and dequantization, which makes cross-paper comparison unreliable and cross-model comparison inside one system easy to get wrong.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""A RealNVP coupling flow, transcribed the way the change of variables reads.
log p(x) = log p_Z(f(x)) + log |det df/dx|
The second term is the whole design problem. A general Jacobian determinant
costs O(d^3); the coupling layer makes it a SUM by construction:
y[:d] = x[:d] frozen half, passes through
y[d:] = x[d:] * exp(s(x[:d])) + t(x[:d]) affine, conditioned on the half
The Jacobian of that is block-triangular with a diagonal block, so its
determinant is the product of the diagonal - a sum in log space, computed in
linear time without ever forming the matrix.
Plain loops, no libraries.
"""
import math
import random
SEED = 13
LOG_2PI = math.log(2.0 * math.pi)
def tanh_network(inputs, params):
"""The scale-and-translate network.
Never inverted, so it can be arbitrarily deep and nonlinear. ALL the
expressive power of a flow hides in the part that does not need to be
invertible - that is the trick, stated plainly.
"""
hidden = []
for row, bias in zip(params['w1'], params['b1']):
total = bias + sum(w * v for w, v in zip(row, inputs))
hidden.append(math.tanh(total))
return [
bias + sum(w * h for w, h in zip(row, hidden))
for row, bias in zip(params['w2'], params['b2'])
]
def clamp_scale(raw, limit=3.0):
"""Bound the scale before exponentiating.
NOT optional. The affine map exponentiates s, so an unbounded scale
network makes the log-determinant diverge and the loss becomes NaN in a
single step - with no gradual warning at all.
"""
return limit * math.tanh(raw / limit)
def coupling_forward(x, split, params):
"""Data -> latent. Returns the output and the log-determinant.
The log-determinant is a SUM of the scales. No determinant is computed
anywhere, because the architecture guarantees the Jacobian is triangular.
"""
frozen = x[:split]
transformed = x[split:]
raw = tanh_network(frozen, params)
scales = [clamp_scale(raw[i]) for i in range(len(transformed))]
shifts = raw[len(transformed):]
output = list(frozen)
log_determinant = 0.0
for i, value in enumerate(transformed):
output.append(value * math.exp(scales[i]) + shifts[i])
log_determinant += scales[i]
return output, log_determinant
def coupling_inverse(y, split, params):
"""Latent -> data. The SAME network, evaluated forwards.
This is what makes the construction work: inverting the coupling never
requires inverting the scale network, because that network reads only the
frozen half - which is unchanged and therefore already known.
"""
frozen = y[:split]
transformed = y[split:]
raw = tanh_network(frozen, params)
scales = [clamp_scale(raw[i]) for i in range(len(transformed))]
shifts = raw[len(transformed):]
output = list(frozen)
for i, value in enumerate(transformed):
output.append((value - shifts[i]) * math.exp(-scales[i]))
return output
def permute(vector, permutation):
"""Shuffle dimensions between couplings.
Without this the same half is frozen in every layer and depth buys
nothing at all - the untransformed dimensions would pass straight from
input to output unchanged.
"""
return [vector[index] for index in permutation]
def inverse_permutation(permutation):
inverse = [0] * len(permutation)
for position, index in enumerate(permutation):
inverse[index] = position
return inverse
def flow_forward(x, layers):
"""Compose the couplings. Determinants MULTIPLY, so logs ADD.
That is why flows are built deep: composing K layers costs nothing extra
in the log-determinant term, it is just a longer sum.
"""
state = list(x)
total_log_determinant = 0.0
for layer in layers:
state = permute(state, layer['permutation'])
state, log_determinant = coupling_forward(state, layer['split'], layer['params'])
total_log_determinant += log_determinant
return state, total_log_determinant
def flow_inverse(z, layers):
"""Run the stack backwards, inverting each layer."""
state = list(z)
for layer in reversed(layers):
state = coupling_inverse(state, layer['split'], layer['params'])
state = permute(state, inverse_permutation(layer['permutation']))
return state
def standard_normal_log_density(z):
"""log p_Z(z) for a standard Gaussian base."""
return -0.5 * sum(value * value + LOG_2PI for value in z)
def log_likelihood(x, layers):
"""The exact log-density. No bound, no sampling estimate, no gap.
Contrast a VAE, which can only bound this quantity. That difference is
the entire reason to choose a flow.
"""
z, log_determinant = flow_forward(x, layers)
return standard_normal_log_density(z) + log_determinant
def negative_log_likelihood(batch, layers):
"""The training objective, and the only one there is."""
return -sum(log_likelihood(x, layers) for x in batch) / len(batch)
def sample(layers, dimension, rng):
"""Draw from the base and push it through the inverse.
Sampling and density evaluation cost the SAME here, which is unusual:
most generative models are cheap in one direction and expensive in the
other, and this needs no second model for either.
"""
z = [rng.gauss(0.0, 1.0) for _ in range(dimension)]
return flow_inverse(z, layers)
def check_invertibility(x, layers, tolerance=1e-8):
"""The architecture guarantees this - an implementation bug does not.
Worth running once as a test: nothing in training ever checks
invertibility, so a transposed index or a sign error produces a model
that trains happily and cannot sample.
"""
z, _ = flow_forward(x, layers)
recovered = flow_inverse(z, layers)
return max(abs(a - b) for a, b in zip(x, recovered)) < tolerance
O(K * d * W) per density evaluation with K coupling layers and width W, and the same cost for sampling. Illustrative, not a measured benchmark: the log-determinant is O(d) because the architecture makes the Jacobian triangular — a general determinant at the same dimension would be O(d^3) and would make the whole approach unusable.
How this language actually wants it written.
"""The same flow, with the bijection as a protocol and invertibility tested.
Two things change. A bijection becomes an explicit interface with forward,
inverse and log-determinant, because composing them is the whole method and a
uniform contract is what makes composition safe. And the numerical guards move
from comments into code: the scale clamp is enforced at construction rather
than hoped for, and invertibility becomes a testable property rather than an
architectural claim nobody verified.
"""
from __future__ import annotations
import math
import random
from dataclasses import dataclass
from typing import NamedTuple, Protocol, Sequence
class NotInvertible(RuntimeError):
"""Raised when forward then inverse does not recover the input.
Its own type because the architecture GUARANTEES invertibility and an
implementation bug does not: a transposed index or a sign error produces
a model that trains happily and cannot sample, with nothing in the loss
to indicate it.
"""
class DivergentScale(ValueError):
"""Raised when the affine scale is unbounded.
The coupling exponentiates the scale, so an unclamped network makes the
log-determinant diverge and the loss becomes NaN in a single step, with
no gradual warning. Refusing the configuration is cheaper than debugging
the NaN.
"""
class ShapeMismatch(ValueError):
"""Raised on a dimension violation instead of computing nonsense."""
class Transformed(NamedTuple):
"""Output AND log-determinant.
They come back together because they are produced together: the
log-determinant of a coupling is a by-product of the scales it already
computed, and returning only the output would force a second pass.
"""
values: list[float]
log_determinant: float
class Bijection(Protocol):
"""A single invertible transformation.
A protocol rather than a base class because composition is the whole
method: a flow is a list of these, and every one must honour the same
contract for the composition to be invertible at all.
"""
def forward(self, x: Sequence[float]) -> Transformed: ...
def inverse(self, y: Sequence[float]) -> list[float]: ...
@property
def dimension(self) -> int: ...
@dataclass(frozen=True)
class CouplingConfig:
"""Frozen, because the split and the clamp define the transformation."""
dimension: int
split: int
scale_limit: float = 3.0
def __post_init__(self) -> None:
if not 0 < self.split < self.dimension:
raise ShapeMismatch(
f'split {self.split} must leave both halves non-empty in '
f'{self.dimension} dimensions'
)
# Guard clause for the single-step divergence.
if not 0.0 < self.scale_limit <= 10.0:
raise DivergentScale(
f'scale limit {self.scale_limit} is unusable; the coupling '
'exponentiates the scale and an unbounded one gives NaN in one step'
)
@property
def transformed_width(self) -> int:
return self.dimension - self.split
class AffineCoupling:
"""The construction that makes the determinant tractable.
Split the input; leave one half alone; transform the other elementwise
with parameters computed from the frozen half. The Jacobian is
block-triangular with a diagonal block, so its determinant is the product
of the diagonal — a sum in log space, in linear time, without ever
forming the matrix.
"""
def __init__(self, config: CouplingConfig, weights: dict[str, list]) -> None:
self._config = config
self._weights = weights
@property
def dimension(self) -> int:
return self._config.dimension
def _scale_and_shift(self, frozen: Sequence[float]) -> tuple[list[float], list[float]]:
"""The network is NEVER inverted, so it can be arbitrarily complex.
All the expressive power of a flow hides in the part that does not
need to be invertible. That is the trick, and it is why the scale
network here is a free design choice rather than a constrained one.
"""
hidden = [
math.tanh(bias + math.fsum(w * v for w, v in zip(row, frozen)))
for row, bias in zip(self._weights['w1'], self._weights['b1'])
]
raw = [
bias + math.fsum(w * h for w, h in zip(row, hidden))
for row, bias in zip(self._weights['w2'], self._weights['b2'])
]
width = self._config.transformed_width
limit = self._config.scale_limit
# Clamped through a tanh rather than a hard clip: a hard clip has zero
# gradient outside the range, so a saturated unit could never recover.
scales = [limit * math.tanh(value / limit) for value in raw[:width]]
shifts = list(raw[width:width * 2])
return scales, shifts
def forward(self, x: Sequence[float]) -> Transformed:
if len(x) != self._config.dimension:
raise ShapeMismatch(f'expected {self._config.dimension} values, got {len(x)}')
split = self._config.split
scales, shifts = self._scale_and_shift(x[:split])
values = list(x[:split])
values.extend(
value * math.exp(scale) + shift
for value, scale, shift in zip(x[split:], scales, shifts)
)
# The log-determinant is a SUM of the scales. No determinant is
# computed anywhere — the architecture guarantees triangularity.
return Transformed(values=values, log_determinant=math.fsum(scales))
def inverse(self, y: Sequence[float]) -> list[float]:
"""The SAME network, evaluated forwards.
This is what makes the construction work: inverting never requires
inverting the scale network, because that network reads only the
frozen half — which is unchanged and therefore already known.
"""
split = self._config.split
scales, shifts = self._scale_and_shift(y[:split])
values = list(y[:split])
values.extend(
(value - shift) * math.exp(-scale)
for value, scale, shift in zip(y[split:], scales, shifts)
)
return values
class Permutation:
"""Shuffles dimensions between couplings.
Without this the same half is frozen in every layer and depth buys
nothing — the untransformed dimensions would pass straight from input to
output unchanged. The inverse is precomputed because it is needed on
every sample and recomputing it is pure waste.
"""
def __init__(self, order: Sequence[int]) -> None:
if sorted(order) != list(range(len(order))):
raise ShapeMismatch('a permutation must be a bijection on its indices')
self._order = list(order)
self._inverse = [0] * len(order)
for position, index in enumerate(order):
self._inverse[index] = position
@property
def dimension(self) -> int:
return len(self._order)
def forward(self, x: Sequence[float]) -> Transformed:
# A permutation has determinant plus or minus one, so its log
# magnitude is exactly zero — it contributes nothing to the density.
return Transformed(values=[x[i] for i in self._order], log_determinant=0.0)
def inverse(self, y: Sequence[float]) -> list[float]:
return [y[i] for i in self._inverse]
class Flow:
"""A composition of bijections, with the invariant checked on request."""
def __init__(self, layers: Sequence[Bijection]) -> None:
if not layers:
raise ShapeMismatch('a flow needs at least one bijection')
dimensions = {layer.dimension for layer in layers}
if len(dimensions) != 1:
raise ShapeMismatch(f'every layer must share a dimension, found {dimensions}')
self._layers = list(layers)
self._dimension = dimensions.pop()
def forward(self, x: Sequence[float]) -> Transformed:
"""Determinants MULTIPLY, so logs ADD.
That is why flows are built deep: composing K layers costs nothing
extra in the log-determinant term, it is just a longer sum.
"""
values = list(x)
total = 0.0
for layer in self._layers:
result = layer.forward(values)
values, total = result.values, total + result.log_determinant
return Transformed(values=values, log_determinant=total)
def inverse(self, z: Sequence[float]) -> list[float]:
values = list(z)
for layer in reversed(self._layers):
values = layer.inverse(values)
return values
def log_likelihood(self, x: Sequence[float]) -> float:
"""The EXACT log-density. No bound, no estimate, no variational gap.
Contrast a VAE, which can only bound this quantity. That difference
is the entire reason to choose a flow.
"""
result = self.forward(x)
base = -0.5 * math.fsum(
value * value + math.log(2.0 * math.pi) for value in result.values
)
return base + result.log_determinant
def sample(self, rng: random.Random) -> list[float]:
"""Sampling and density evaluation cost the SAME here.
Unusual: most generative models are cheap in one direction and
expensive in the other, and this needs no second model for either.
"""
return self.inverse([rng.gauss(0.0, 1.0) for _ in range(self._dimension)])
def assert_invertible(self, x: Sequence[float], tolerance: float = 1e-8) -> None:
"""The architecture guarantees this — an implementation does not.
Nothing in training ever checks invertibility, so a transposed index
or a sign error produces a model that trains happily and cannot
sample. This belongs in a test suite, permanently.
"""
recovered = self.inverse(self.forward(x).values)
error = max(abs(a - b) for a, b in zip(x, recovered))
if error > tolerance:
raise NotInvertible(
f'forward then inverse differs by {error:.2e}; the architecture '
'guarantees invertibility, so this is an implementation bug'
)
def dequantize(values: Sequence[int], levels: int, rng: random.Random) -> list[float]:
"""Add uniform noise to discrete data before fitting a continuous density.
Without this the likelihood is UNBOUNDED: a continuous density can put
arbitrarily much mass on a finite set of exact values, so the model
chases infinity rather than learning the distribution. The likelihood is
also only comparable between models using the same scheme, which is why
a published bits-per-dimension without its dequantization is meaningless.
"""
return [(value + rng.random()) / levels for value in values]
What changed: Two changes. A bijection becomes an explicit protocol with forward, inverse and log-determinant, because composition is the whole method and a uniform contract is what makes composing arbitrary layers safe — the flow then validates that every layer shares a dimension, which is the invariant the composition depends on. The second is that the numerical guards move from comments into code. The scale clamp is enforced at construction with its own exception type, because the coupling exponentiates the scale and an unclamped network produces NaN in a single step with no gradual warning; it is applied through a tanh rather than a hard clip, since a hard clip has zero gradient outside its range and a saturated unit could never recover. Invertibility becomes a testable assertion with its own error type, because the architecture guarantees it and an implementation does not — a transposed index produces a model that trains happily and cannot sample, and nothing in the loss would indicate that. The permutation precomputes its inverse since that is needed on every sample, and dequantization is provided with an explicit note that the likelihood is unbounded without it and incomparable across schemes.
Identical asymptotics; the change is in what fails loudly. Illustrative, not a measured benchmark: the precomputed inverse permutation turns sampling from O(d log d) per layer into O(d), which matters when sampling is the production path.
Faster, and the justification for each change.
"""Batched, masked, and fused. The split becomes a mask, not a slice.
The structural change worth understanding: the literal version splits the
vector, which means the two halves live in different arrays and every layer
concatenates them back. Masking instead keeps the vector whole and multiplies
by a binary mask - so the coupling becomes elementwise arithmetic over the
full vector, the permutation becomes a matrix product, and nothing is ever
split, concatenated or copied.
Three changes:
1. Masked coupling. One array throughout; the frozen half is selected by
multiplication rather than by slicing.
2. The scale network becomes GEMMs over the batch.
3. The log-determinant is one masked reduction, which is where the
architecture's linear-time determinant becomes visible as a single sum.
"""
from __future__ import annotations
import numpy as np
from numpy.typing import NDArray
FLOAT = np.float32
LOG_2PI = float(np.log(2.0 * np.pi))
class MaskedCoupling:
"""An affine coupling expressed with a binary mask rather than a split.
mask == 1 marks the frozen dimensions. The transformation is then
elementwise over the WHOLE vector, with the mask selecting what changes:
y = mask * x + (1 - mask) * (x * exp(s) + t)
Algebraically identical to the split form, and structurally much better:
no slicing, no concatenation, and the scale network reads a masked copy
rather than a separate array.
"""
def __init__(self, mask: NDArray[np.float32], w1: NDArray[np.float32],
b1: NDArray[np.float32], w2: NDArray[np.float32],
b2: NDArray[np.float32], scale_limit: float = 3.0) -> None:
self._mask = np.ascontiguousarray(mask, dtype=FLOAT)
self._w1 = np.ascontiguousarray(w1, dtype=FLOAT)
self._b1 = np.ascontiguousarray(b1, dtype=FLOAT)
self._w2 = np.ascontiguousarray(w2, dtype=FLOAT)
self._b2 = np.ascontiguousarray(b2, dtype=FLOAT)
self._limit = FLOAT(scale_limit)
def _scale_and_shift(
self, x: NDArray[np.float32]
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
"""Two GEMMs for the whole batch. The network is never inverted, so
it can be as large as you like — all the expressive power hides
here, in the part that does not have to be invertible."""
hidden = (x * self._mask) @ self._w1
hidden += self._b1
np.tanh(hidden, out=hidden)
projected = hidden @ self._w2
projected += self._b2
width = projected.shape[1] // 2
scale, shift = projected[:, :width], projected[:, width:]
# Clamp in place through a tanh: NOT optional, since the coupling
# exponentiates this and an unbounded scale is a one-step NaN. A tanh
# rather than a clip because a clip has zero gradient outside its
# range and a saturated unit could never recover.
scale /= self._limit
np.tanh(scale, out=scale)
scale *= self._limit
return scale, shift
def forward(
self, x: NDArray[np.float32]
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
"""Returns the output and the per-sample log-determinant."""
scale, shift = self._scale_and_shift(x)
free = 1.0 - self._mask
# Fused: exponentiate in place, apply the affine map, and select with
# the mask in one expression — no intermediate array for the
# transformed half exists at any point.
transformed = np.exp(scale * free)
transformed *= x
transformed += shift * free
output = self._mask * x + free * transformed
# The log-determinant is a masked SUM of the scales. No determinant is
# computed anywhere: the architecture makes the Jacobian triangular,
# which is the whole reason this is O(d) rather than O(d^3).
log_determinant = np.einsum('ij,j->i', scale, free, optimize=True)
return output, log_determinant
def inverse(self, y: NDArray[np.float32]) -> NDArray[np.float32]:
"""The SAME network, evaluated forwards.
Inverting never requires inverting the scale network, because that
network reads only the masked half — which is unchanged and therefore
already known. That is the entire construction.
"""
scale, shift = self._scale_and_shift(y)
free = 1.0 - self._mask
recovered = y - shift * free
recovered *= np.exp(-scale * free)
return self._mask * y + free * recovered
class LinearMixing:
"""A learned invertible mixing layer, replacing a fixed permutation.
A permutation is a fixed, sparse special case of this. A general
invertible matrix uses depth far more efficiently — but its determinant
is no longer free, so it is kept as an LU factorization where the
log-determinant is a sum over the diagonal and stays O(d).
"""
def __init__(self, lower: NDArray[np.float32], upper: NDArray[np.float32],
permutation: NDArray[np.int32]) -> None:
self._lower = lower
self._upper = upper
self._permutation = permutation
# Precomputed once: the log-determinant of an LU factorization is a
# sum over the upper diagonal, and it does not depend on the input.
self._log_determinant = float(np.sum(np.log(np.abs(np.diag(upper)))))
self._weight = (lower @ upper)[permutation]
def forward(
self, x: NDArray[np.float32]
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
return x @ self._weight.T, np.full(x.shape[0], self._log_determinant, dtype=FLOAT)
class BatchedFlow:
"""A composition of masked couplings and mixing layers."""
def __init__(self, layers: list) -> None:
self._layers = layers
def log_likelihood(self, batch: NDArray[np.float32]) -> NDArray[np.float32]:
"""Exact log-density for a whole batch. No bound, no estimate.
Determinants MULTIPLY, so logs ADD — composing K layers is just a
longer sum, which is why flows are built deep.
"""
state = batch
total = np.zeros(batch.shape[0], dtype=FLOAT)
for layer in self._layers:
state, log_determinant = layer.forward(state)
total += log_determinant
# Base log-density contracted in one pass rather than squaring into a
# temporary and summing it.
squared = np.einsum('ij,ij->i', state, state, optimize=True)
base = -0.5 * (squared + FLOAT(batch.shape[1] * LOG_2PI))
return base + total
def sample(self, count: int, dimension: int, rng: np.random.Generator) -> NDArray[np.float32]:
"""Sampling and density evaluation cost the SAME here.
Unusual: most generative models are cheap in one direction and
expensive in the other, and this needs no second model for either.
"""
state = rng.standard_normal((count, dimension)).astype(FLOAT)
for layer in reversed(self._layers):
state = layer.inverse(state)
return state
def dequantize(
values: NDArray[np.int32], levels: int, rng: np.random.Generator
) -> NDArray[np.float32]:
"""Uniform noise before fitting a continuous density to discrete data.
Without it the likelihood is UNBOUNDED: a continuous density can put
arbitrarily much mass on a finite set of exact values, so the model
chases infinity rather than learning. Likelihoods are also only
comparable between models using the same scheme, which is why a published
bits-per-dimension without its dequantization is meaningless.
"""
noise = rng.random(values.shape, dtype=FLOAT)
return (values.astype(FLOAT) + noise) / FLOAT(levels)
def bits_per_dimension(log_likelihood: NDArray[np.float32], dimension: int, levels: int) -> float:
"""The reporting convention, with the dequantization correction included.
Omitting the level correction is the standard way published numbers
become incomparable — it is a constant offset, so it never looks wrong.
"""
nats = -log_likelihood.mean() / dimension
return float(nats / np.log(2.0) + np.log2(levels))
What changed: The structural change worth understanding is that the split becomes a mask. The literal version slices the vector, which means the two halves live in different arrays and every layer concatenates them back; multiplying by a binary mask keeps the vector whole, so the coupling becomes elementwise arithmetic over the full width and nothing is ever split, concatenated or copied. The formulation is algebraically identical and structurally far better: the scale network reads a masked copy rather than a separate array, the affine map fuses into one expression with no intermediate for the transformed half, and the log-determinant becomes a single masked contraction — which is where the architecture’s linear-time determinant becomes visible as one sum rather than an argument in a comment. The scale clamp runs in place through a tanh rather than a clip, deliberately, because a clip has zero gradient outside its range and a saturated unit could never recover. The fixed permutation is also generalized to a learned invertible mixing layer kept in LU form, where the log-determinant is a precomputed sum over the diagonal and therefore still free.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
The scale network is two GEMMs for the whole batch, the mixing layer is one, and both the log-determinant and the base density contract in a single einsum pass rather than squaring into temporaries.
Trade-off: Masking means the scale network computes outputs for the frozen dimensions too and then discards them, so roughly half the network arithmetic is wasted — the split form avoids that waste and pays for it in copies instead.
Fuse adjacent operations so intermediates are never materialized
The exponential, the affine map and the mask selection happen in one chained expression, and the scale clamp runs in place over the projection buffer.
Trade-off: The raw scale is destroyed by the in-place clamp, so diagnosing a scale network that has saturated — which is why a flow stops improving — needs an unfused pass.
Pre-allocate output arrays and use in-place operations
The tanh, the clamp and the bias adds all write through existing buffers, so a forward pass allocates only the coupling output.
Trade-off: The layers mutate arrays derived from their input, so a caller that keeps a reference to an intermediate state finds it changed — a real hazard when caching activations for a gradient check.
Ensure contiguous memory layout and a single dtype
float32 contiguous weights and masks keep every product on the BLAS fast path, and the mask is stored as float rather than bool so it multiplies without a conversion.
Trade-off: float32 accumulation of the log-determinant across a deep stack drifts, and because the density is the sum of that with the base term the error lands directly in the reported likelihood — which is the one number this model exists to get exactly right.
A density evaluation becomes two GEMMs and a contraction per coupling layer for the whole batch. Illustrative, not a measured benchmark: the log-determinant stays O(d) per layer because the architecture guarantees triangularity — a general determinant at the same dimension would be O(d^3) and would make the approach unusable at any width.