Continuous Spaces and Morphisms¶

This page covers the ContinuousSpace hierarchy, the ContinuousMorphism interface, composition operations, the boundary between discrete and continuous models, and normalizing flows. The full registry of parameterized distribution families lives in continuous families.

ContinuousSpace hierarchy¶

Continuous morphisms act between continuous measurable spaces. quivers provides a hierarchy of standard spaces:

ContinuousSpace (abstract)
├── Euclidean         R^n with Lebesgue measure
├── UnitInterval      (0, 1) with Lebesgue measure (defined via Euclidean)
├── Simplex           {x in R^dim : x >= 0, sum x = 1}
├── PositiveReals     (0, inf) with Lebesgue measure
├── CholeskyFactor    K x K lower-triangular Cholesky factors with unit-norm rows
├── ProductSpace      Cartesian product of spaces (and discrete objects)
├── Sphere            unit sphere S^{N-1} in R^N
├── Ball              closed ball of radius r in R^N
├── Covariance        cone of D x D symmetric positive-definite matrices
├── Correlation       D x D correlation matrices (unit diagonal, PD)
├── Orthogonal        orthogonal group O(D)
├── Stiefel           Stiefel manifold V_K(R^N) of N x K orthonormal-column matrices
├── LowerTriangular   D x D lower-triangular matrices
└── Diagonal          D x D diagonal matrices (identified with R^D)

Euclidean space¶

from quivers.continuous.spaces import Euclidean

R3 = Euclidean(name="R3", dim=3)                                  # R^3
R2_bounded = Euclidean(name="R2", dim=2, low=0.0, high=1.0)        # [0,1]^2

Unit interval¶

from quivers.continuous.spaces import UnitInterval

U = UnitInterval(name="U")     # (0, 1); pass dim=k for [0, 1]^k

Simplex¶

from quivers.continuous.spaces import Simplex

S3 = Simplex(name="S3", dim=3)   # 2-simplex in R^3 (3 categories with sum = 1)

PositiveReals¶

from quivers.continuous.spaces import PositiveReals

P2 = PositiveReals(name="P2", dim=2)   # (0, inf)^2

ProductSpace¶

from quivers.continuous.spaces import ProductSpace, Euclidean, UnitInterval

R3 = Euclidean(name="R3", dim=3)
U = UnitInterval(name="U")
P = ProductSpace(components=(R3, U))   # R^3 x (0, 1); also R3 * U

ContinuousMorphism¶

A continuous morphism \(f : X \to Y\) between spaces \(X\) and \(Y\) defines a conditional distribution \(p(y \mid x)\) via two operations:

log_prob(x, y): evaluate \(\log p(y \mid x)\).
rsample(x): generate reparameterized samples from \(p(\cdot \mid x)\).

from quivers.continuous.morphisms import ContinuousMorphism
from quivers.continuous.spaces import Euclidean
import torch

class MyMorphism(ContinuousMorphism):
    def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """Log density log p(y | x)."""
        raise NotImplementedError()

    def rsample(
        self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
    ) -> torch.Tensor:
        """Reparameterized samples; shape (*sample_shape, batch, codomain.dim)."""
        raise NotImplementedError()

Morphisms are nn.Module subclasses, so they are trainable.

Composition: SampledComposition¶

When composing continuous morphisms, the intermediate is continuous. The result is computed via ancestral sampling:

import torch
from quivers.continuous.spaces import Euclidean
from quivers.continuous.morphisms import SampledComposition
from quivers.continuous.families import ConditionalNormal

# f: X -> Y (continuous), g: Y -> Z (continuous)
X = Euclidean(name="X", dim=3)
Y = Euclidean(name="Y", dim=4)
Z = Euclidean(name="Z", dim=5)
f = ConditionalNormal(X, Y)
g = ConditionalNormal(Y, Z)

# Composition: (g . f)(x) samples from g(f(x))
composed = f >> g
assert isinstance(composed, SampledComposition)

# rsample: sample from intermediate
x = torch.randn(2, 3)
z_samples = composed.rsample(x, sample_shape=torch.Size((50,)))  # (50, 2, 5)

When an intermediate is discrete (a FinSet), composition uses exact marginalization:

\[p(z \mid x) = \sum_y p(y \mid x) \, p(z \mid y)\]

ProductContinuousMorphism¶

Tensor product of continuous morphisms:

from quivers.continuous.spaces import Euclidean, UnitInterval
from quivers.continuous.morphisms import ProductContinuousMorphism
from quivers.continuous.families import ConditionalNormal, ConditionalBeta

X = Euclidean(name="X", dim=3)
f = ConditionalNormal(X, Euclidean(name="Y1", dim=2))
g = ConditionalBeta(X, UnitInterval(name="Y2"))

# Product: (f @ g)(x) ~ p(y_1, y_2 | x) = p(y_1 | x) . p(y_2 | x)
fg = f @ g

# Domain and codomain are products
assert fg.domain == f.domain * g.domain   # X x X
assert fg.codomain == f.codomain * g.codomain  # R^2 x (0,1)

Discrete / continuous boundary¶

Embed: embed a discrete domain into a continuous codomain¶

Embed a finite-set value into a continuous space:

import torch
from quivers.continuous.spaces import Euclidean
from quivers.continuous.boundaries import Embed
from quivers.core.morphisms import morphism
from quivers.core.objects import FinSet

X = FinSet(name="X", cardinality=4)
Y = FinSet(name="Y", cardinality=5)
discrete_f = morphism(X, Y)

# Give each element of X a learnable Gaussian embedding in R^5
Y_continuous = Euclidean(name="Y", dim=5)
embedded_f = Embed(domain=X, codomain=Y_continuous)

# Now can compose with continuous morphisms
y_cont = torch.randn(4, 5)
log_p = embedded_f.log_prob(torch.arange(4), y_cont)

Discretize: bin a continuous space¶

Discretize a continuous space into a finite set:

import torch
from quivers.continuous.spaces import Euclidean
from quivers.continuous.boundaries import Discretize

# Discretize a bounded interval into 20 bins
# (Discretize requires bounded Euclidean)
U = Euclidean(name="U", dim=1, low=0.0, high=1.0)
discretized = Discretize(U, n_bins=20)

# Discretize bins the continuous domain into the codomain FinSet.
assert discretized.domain == U
assert discretized.codomain.cardinality == 20

Integration with discrete morphisms¶

Continuous and discrete morphisms integrate transparently via the >> and @ operators:

# Discrete to Continuous
discrete_f = morphism(X, Y)  # X -> Y (finite)
cont_g = ConditionalNormal(
    domain=Y,                        # discrete domain, indexed lookup
    codomain=Euclidean(name="Z", dim=3),
)

composed = discrete_f >> cont_g  # exact marginalization over Y

# Continuous to Discrete: not directly composable without an
# explicit Discretize boundary; use Discretize to materialize a
# FinSet target first.

Normalizing flows¶

A normalizing flow transforms a simple base distribution via a chain of invertible transformations.

AffineCouplingLayer¶

An affine coupling layer partitions the input and applies an invertible affine transformation to each partition:

import torch
from quivers.continuous.flows import AffineCouplingLayer
from quivers.continuous.spaces import Euclidean

dim = 4
domain = Euclidean(name="X", dim=2)
layer = AffineCouplingLayer(domain, dim, hidden_dim=32)

# Forward (base to target), conditioned on x
x = torch.randn(10, 2)
z = torch.randn(10, dim)
z_out, log_det = layer.forward(x, z)

# Inverse (target to base)
z_back, log_det_inv = layer.inverse(x, z_out)

ConditionalFlow¶

A full normalizing flow conditioned on input:

import torch
from quivers.continuous.flows import ConditionalFlow
from quivers.continuous.spaces import Euclidean

domain = Euclidean(name="X", dim=5)
codomain = Euclidean(name="Y", dim=4)

flow = ConditionalFlow(
    domain=domain,
    codomain=codomain,
    n_layers=6,
    hidden_dim=64,
)

# Sample
x = torch.randn(8, 5)
y_samples = flow.rsample(x, sample_shape=torch.Size((50,)))

# Log probability
y = torch.randn(8, 4)
log_p = flow.log_prob(x, y)

The quivers.continuous.flows module exposes two public building blocks: AffineCouplingLayer, a single invertible affine coupling layer, and ConditionalFlow, a stack of coupling layers wrapped as a ContinuousMorphism.

References¶

Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. 2019. Neural spline flows. arXiv preprint arXiv:1906.04032.
Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. 2016. Improving variational inference with inverse autoregressive flow. arXiv preprint arXiv:1606.04934.
George Papamakarios, Theo Pavlakou, and Iain Murray. 2017. Masked autoregressive flow for density estimation. arXiv preprint arXiv:1705.07057.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. 2016. Density estimation using Real NVP. arXiv preprint arXiv:1605.08803.