NAE: Normalizing AutoEncoder

2026-08-12

AI Summary

This work addresses a critical inconsistency between the existing loss functions used in flow-based autoencoders and their reconstruction objectives, which leads to suboptimal training dynamics. The paper provides the first theoretical analysis of this misalignment and introduces a novel conditional loss function designed to align the surrogate gradients of the encoder and decoder with the true reconstruction loss. By preserving the established architecture that combines normalizing flows with autoencoders, the proposed method significantly enhances generative performance. It achieves state-of-the-art results across diverse benchmarks—including molecular generation, tabular data modeling, and image synthesis—demonstrating both its effectiveness and broad applicability.

Abstract

We consider the setting of Normalizing flows with approximate inverses, an established paradigm spanning both full-dimensional ($d=D$) and bottleneck ($d<D$) settings, and group these models under the term flow autoencoders. We present a theoretical investigation into their training dynamics and prove that the proposed loss used by existing approaches is suboptimal; specifically, both encoder and decoder surrogates must be optimized in alignment with reconstruction loss. Guided by these insights, we propose Normalizing Autoencoder (NAE), which employs a novel conditional loss that aligns the surrogate loss gradient with that of reconstruction loss, directly improving upon the current standard. Extensive experiments across molecule generation, tabular data, and image benchmarks demonstrate that NAE achieves state of the art performance. Our work highlights the importance of loss alignment in flow autoencoders and establishes NAE as a powerful generative framework.

Problem

• Normalizing flows
• flow autoencoders
• loss alignment
• reconstruction loss
• generative modeling

Innovation

• Normalizing Autoencoder
• flow autoencoders
• loss alignment
• surrogate loss
• reconstruction loss