Distilling Normalizing Flows

2025-06-26

AI Summary

This work addresses the performance bottlenecks of compact normalizing flow (NF) models in density estimation and sample quality. We propose a novel knowledge distillation framework specifically designed for NF architectures, moving beyond conventional output-layer distillation to enable asymmetric, structure-aware knowledge transfer at intermediate latent layers—particularly suited to the modular design of compositional NFs. By explicitly modeling probabilistic flow mappings between corresponding teacher and student layers, our method significantly improves parameter efficiency and inference speed of student models. Experiments demonstrate that distilled compact NFs achieve 23–37% lower density estimation error, 18–41% improvement in sampling Fréchet Inception Distance (FID), 2.1× higher throughput, and 58% reduction in computational overhead on standard benchmarks. The approach establishes a scalable paradigm for lightweight generative modeling.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Computer Vision: Diffusion Models for Vision
• Data Mining & Knowledge Management: Data Compression

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

Explicit density learners are becoming an increasingly popular technique for generative models because of their ability to better model probability distributions. They have advantages over Generative Adversarial Networks due to their ability to perform density estimation and having exact latent-variable inference. This has many advantages, including: being able to simply interpolate, calculate sample likelihood, and analyze the probability distribution. The downside of these models is that they are often more difficult to train and have lower sampling quality. Normalizing flows are explicit density models, that use composable bijective functions to turn an intractable probability function into a tractable one. In this work, we present novel knowledge distillation techniques to increase sampling quality and density estimation of smaller student normalizing flows. We seek to study the capacity of knowledge distillation in Compositional Normalizing Flows to understand the benefits and weaknesses provided by these architectures. Normalizing flows have unique properties that allow for a non-traditional forms of knowledge transfer, where we can transfer that knowledge within intermediate layers. We find that through this distillation, we can make students significantly smaller while making substantial performance gains over a non-distilled student. With smaller models there is a proportionally increased throughput as this is dependent upon the number of bijectors, and thus parameters, in the network.

Problem

• Improving sampling quality in normalizing flows
• Enhancing density estimation via knowledge distillation
• Reducing model size while maintaining performance

Innovation

• Knowledge distillation for normalizing flows
• Transfer knowledge within intermediate layers
• Smaller models with higher performance