Reproduce “Transformer Neural Autoregressive Flows” on OpenResearch

About this paper

Massimiliano Patacchiola, Aliaksandra Shysheya, Katja Hofmann, Richard E. Turner

T-NAF models the density of multivariate real-valued data using a transformer that treats dimensions as tokens and uses causal attention to condition invertible scalar transformations on earlier dimensions. Its shared conditioner predicts transformation parameters, while monotone CDF transformations preserve invertibility. On 21-dimensional HEPMASS, one-flow T-NAF reports −13.20 ± 0.26 test log-likelihood, but scores lower than B-NAF on MINIBOONE and POWER.

Key insights

• Transformer architectures can efficiently enforce autoregressive constraints through attention masking, offering a more parameter-efficient alternative to traditional masked feed-forward networks.
• Decoupling strict monotonicity constraints to a compact, single input-output transformation (instead of the entire conditioner) leads to more stable training dynamics.
• Parameter sharing inherent in transformer designs allows T-NAFs to achieve comparable or superior density estimation performance with an order of magnitude fewer parameters than prior autoregressive flows, especially for high-dimensional data.

Summary by alphaXiv.