Phase-aware Training Schedule Simplifies Learning in Flow-Based Generative Models

TL;DR

This paper introduces a phase-aware training schedule for flow-based generative models, improving learning efficiency by explicitly modeling the different phases of mode and variance learning in high-dimensional Gaussian mixtures.

Contribution

It proposes a novel time dilation technique to address phase disappearance in training, and introduces a method to identify optimal training intervals for real data features.

Findings

Time dilation enables mode and variance learning phases.

Autoencoder simplifies by focusing on phase-relevant parameters.

Method improves training efficiency on real data.

Abstract

We analyze the training of a two-layer autoencoder used to parameterize a flow-based generative model for sampling from a high-dimensional Gaussian mixture. Previous work shows that the phase where the relative probability between the modes is learned disappears as the dimension goes to infinity without an appropriate time schedule. We introduce a time dilation that solves this problem. This enables us to characterize the learned velocity field, finding a first phase where the probability of each mode is learned and a second phase where the variance of each mode is learned. We find that the autoencoder representing the velocity field learns to simplify by estimating only the parameters relevant to each phase. Turning to real data, we propose a method that, for a given feature, finds intervals of time where training improves accuracy the most on that feature. Since practitioners take a uniform distribution over training times, our method enables more efficient training. We provide preliminary experiments validating this approach.