Unconstrained Neural Autoregressive Flow (UNAF)

​

Overview

​

Reference

​

Class Definition

zuko.flows.UNAF(
    features: int,
    context: int = 0,
    transforms: int = 3,
    randperm: bool = False,
    signal: int = 16,
    network: dict = {},
    **kwargs
)

​

Parameters

​

Usage Example

import torch
import zuko

# Create an unconditional UNAF
flow = zuko.flows.UNAF(
    features=5,
    transforms=5,
    signal=32,
    hidden_features=[128, 128]
)

# Sample from the flow
dist = flow()
samples = dist.sample((1000,))
print(samples.shape)  # torch.Size([1000, 5])

# Compute log probabilities
log_prob = dist.log_prob(samples)
print(log_prob.shape)  # torch.Size([1000])

​

Conditional Flow

# Create a conditional UNAF
flow = zuko.flows.UNAF(
    features=3,
    context=5,
    transforms=5,
    signal=24
)

context = torch.randn(5)
dist = flow(context)
samples = dist.sample((100,))

​

Training Example

import torch.optim as optim

flow = zuko.flows.UNAF(
    features=10,
    transforms=5,
    signal=32,
    hidden_features=[256, 256],
    network={'hidden_features': [64, 64, 64], 'activation': torch.nn.ELU}
)

optimizer = optim.Adam(flow.parameters(), lr=1e-3)

for epoch in range(100):
    for x in dataloader:
        optimizer.zero_grad()
        
        # Ensure data is in [-10, 10]
        x = torch.clamp(x, -10, 10)
        
        loss = -flow().log_prob(x).mean()
        loss.backward()
        optimizer.step()

​

Methods

​

forward(c=None)

• c (Tensor, optional): Context tensor of shape (*, context)

• NormalizingFlow: A distribution with:
  • sample(shape): Sample from the distribution
  • log_prob(x): Compute log probability of samples
  • rsample(shape): Reparameterized sampling

​

When to Use UNAF

• Maximum flexibility in monotonic transformations
• Complex, highly nonlinear distributions
• Research and experimentation
• When architectural constraints are limiting

• You need fast training (use NAF or NSF)
• You need fast sampling (use RealNVP)
• Your data is outside [-10, 10] and can’t be standardized
• You want simpler, more interpretable models (use MAF or NSF)

​

Tips

1. Standardize your data: UNAF requires features in [-10, 10]. Always normalize inputs.
2. Use ELU activation: UNAF uses ELU activation by default, which works well with integration.
3. Tune signal dimension: Start with signal=16. Increase for more complex distributions.
4. Be patient: UNAF can be slow to train due to integration but is very expressive.

​

Architecture Details

• Base distribution: Diagonal Gaussian N(0, I)
• Transformation: Integration of positive integrand functions
• Signal network: Masked MLP predicts signal vectors and constants
• Integrand network: Unconstrained MLP whose integral defines the transformation
• Softclip layers: Inserted between transformations

y_i = ∫[0 to x_i] g(t; signal_i) dt + constant_i

​

Unconstrained Monotonic Networks

• No weight constraints: Weights can be any value
• Integration-based: Monotonicity via integration, not positive weights
• Integrand function: Models dy/dx instead of y directly
• Numerical integration: Uses ODE solvers for inversion

# Simplified UMNN concept
class UMNN(nn.Module):
    def __init__(self, signal_dim):
        # Unconstrained network
        self.integrand = MLP(1 + signal_dim, 1)
    
    def forward(self, x, signal, constant):
        # Integrate from 0 to x
        def g(t):
            inp = torch.cat([t, signal], dim=-1)
            return exp(self.integrand(inp))  # Always positive
        
        integral = integrate(g, 0, x)
        return integral + constant

​

UNAF vs NAF

​

Integration Details

• Forward pass: Integrate from 0 to x
• Inverse pass: Root finding to solve integral equation
• Log determinant: Computed from integrand evaluations

g(t) = exp(network(t, signal) / (1 + abs(network(t, signal) / 7)))

​

Advanced Usage

​

Custom Integrand Network

flow = zuko.flows.UNAF(
    features=10,
    transforms=5,
    signal=32,
    network={
        'hidden_features': [128, 128, 128, 128],
        'activation': torch.nn.Tanh  # Custom activation
    }
)

​

High-Dimensional Data

# Use coupling for efficiency
flow = zuko.flows.UNAF(
    features=100,
    transforms=3,
    passes=2,  # Coupling transformations
    signal=32
)

​

Manual Construction

from zuko.flows.autoregressive import MaskedAutoregressiveTransform
from zuko.flows.neural import UMNN

transform = MaskedAutoregressiveTransform(
    features=10,
    context=5,
    univariate=UMNN(
        signal=32,
        hidden_features=[64, 64, 64],
        activation=torch.nn.ELU
    ),
    shapes=[(32,), ()],  # Signal and constant shapes
    hidden_features=[256, 256]
)

​

Computational Considerations

• Training: Slower than NAF due to integration
• Memory: Higher due to integration state
• Inversion: Requires numerical root finding
• Gradients: Computed via adjoint method

1. Use smaller networks for the integrand
2. Reduce signal dimensions
3. Use coupling for high-dimensional data
4. Adjust integration tolerances

​

Comparison with Other Flows

​

Research Applications

• Research: Exploring limits of flow expressivity
• Benchmarking: Comparing against other architectures
• Complex distributions: Multi-modal, long-tailed, irregular
• Ablation studies: Understanding monotonic transformations

​

Related

• NAF - Constrained monotonic neural networks
• NSF - Spline-based alternative
• UnconstrainedMonotonicTransform - The underlying transformation
• MAF - Simpler baseline