Distilling the Knowledge from Normalizing Flows
Abstract
Normalizing ﬂows are a powerful class of generative models demonstrating strong performance in several speech and vision problems.
In contrast to other generative models, normalizing ﬂows are latent variable models with tractable likelihoods and allow for stable training.
However, they have to be carefully designed to represent invertible functions with efﬁcient Jacobian determinant calculation.
In practice, these requirements lead to overparameterized and sophisticated architectures that are inferior to alternative feed-forward models in terms of inference time and memory consumption.
In this work, we investigate whether one can distill ﬂow-based models into more efﬁcient alternatives.
We provide a positive answer to this question by proposing a simple distillation approach and demonstrating its effectiveness on state-of-the-art conditional ﬂow-based models for image super-resolution and speech synthesis.
1. Introduction
Normalizing ﬂows (NF) (Dinh et al., 2015; Rezende & Mohamed, 2015) are a class of generative models that construct complex probability distributions by applying a series of invertible transformations to a simple base density (typically multivariate Gaussian).
Such a design allows for exact likelihood computation via change-of-variables formula.
Therefore, NF can be straightforwardly trained via likelihood maximization.
This NF property is appealing for practitioners since alternative GAN-based models require adversarial optimization, which suffers from vanishing gradients, mode collapse, oscillating or cyclic behavior (Goodfellow, 2016).
To enable exact likelihood computation NF architectures must be composed of invertible modules that also support the efﬁcient calculation of their Jacobian determinant.
A
Copyright large number of such modules have been recently developed, including autoregressive, bipartite, linear and residual transformations (Huang et al., 2018; Kingma et al., 2016;
Papamakarios et al., 2017; van den Berg et al., 2018; Durkan et al., 2019;
Dinh et al., 2017; 2015; Rezende & Mohamed, 2015; Kingma & Dhariwal, 2018; Chen et al., 2019).
While some modules are signiﬁcantly more efﬁcient than others, normalizing ﬂows are generally inferior to feed-forward counterparts (e.g., GANs) in terms of sampling time.
In particular, autoregressive ﬂows (Huang et al., 2018; Papamakarios et al., 2017) use a slow sequential generation procedure, while bipartite ﬂows (Dinh et al., 2017; Kingma & Dhariwal, 2018) can require a lot of submodules with low expressive power.
Moreover, invertibility limits the size of the inner representation leading to impractically deep models.
The more comprehensive discussion of the background and related work is deferred to Appendix A.
However, in many applications, one does not need explicit density estimation but requires efﬁcient inference at deployment.
This raises a question whether one can relinquish the invertibility of normalizing ﬂows to improve their runtime and memory consumption after training.
In this work, we investigate whether this can be achieved through the knowledge distillation approach.
In particular, we describe how to distill knowledge from pretrained ﬂow-based models into efﬁcient architectures, which do not suffer from NF design limitations.
Our code and models are available online.
We summarize our contributions as follows:
• We propose a plain training strategy and student design that allow for knowledge distillation from conditional normalizing ﬂows to feed-forward models with streamlined inference and lower memory consumption.
To the best of our knowledge, this is the ﬁrst work that proposes the NF distillation to feed-forward architectures.
• In our experiments, we empirically conﬁrm the effectiveness of our method on the state-of-the-art ﬂow-based models for super-resolution (SRFlow (Lugmayr et al., 2020a)) and speech synthesis (WaveGlow (Prenger et al., 2018)).
We can achieve up to × 10 speedups with no perceptible loss in quality.
[Figure]
2. Method
This section formally describes the details of our distillation approach.
We also provide a general guide how one can design a student architecture for a given ﬂow-based teacher.
As illustrative examples, we use two state-of-the-art models from speech and vision domains.
2.1. Distillation procedure
A pretrained conditional ﬂow-based model deﬁnes a deterministic bijective mapping from a noise sample z ∼ N (0 , σI ) and contextual information c to an output sample x .
Our distillation approach approximates this mapping by a more efﬁcient model in a supervised manner.
The general scheme of our approach is presented in Figure 1.
2.2. Objective
The proposed distillation objective is a combination of reconstruction and feature losses:
[Equation]
where α is a hyperparameter to balance the loss terms.
For both applications, the reconstruction loss is an average
L 1 -norm between the student and teacher samples:
[Equation]
where T θ and S ψ correspond to the teacher and student models, respectively.
The feature loss for SR model distillation is a perceptual distance between generated images computed by LPIPS (Zhang et al., 2018b).
A pretrained VGG16 model (Simonyan &
Zisserman, 2015) is used for LPIPS calculation.
The feature loss for speech vocoder distillation is a multiresolution STFT 2 loss (van den Oord et al., 2018;
Yamamoto et al., 2020; 2019), which is a sum of STFT losses with different parameters (i.e., FFT size, window size and frame shift).
In more detail, a single STFT loss is the sum of the spectral convergence and log STFT magnitude terms (Yamamoto et al., 2020).
2.3. Student design
The proposed student design can be described by a general recipe: “just clone a teacher architecture and take advantage of not being a ﬂow anymore”.
In other words, our student models inherit the reduced teacher architecture and get free of the following constrains:
Invertibility .
In ﬂow-based models, the size of the inner representations cannot exceed the dimensions of input data vectors.
This restriction might lead to much deeper and hence inefﬁcient models.
Since the student no longer has to be reversible, one can vary representation dimensions arbitrarily.
As a result, student architectures with much fewer blocks can achieve similar performance.
Tractable Jacobian determinants .
Flow-based models usually have to exploit speciﬁc operations for easy-tocompute Jacobian determinants.
Therefore, the student models might beneﬁt from replacing these operations with more efﬁcient and expressive ones.
Below, we describe the particular design options on the example of ﬂow-based models with afﬁne coupling layers (Dinh et al., 2017; Kingma & Dhariwal, 2018).
The supportive illustration is provided in Figure 2.
a) The initial ﬂow-based model with afﬁne coupling layers, e.g., SRFlow or WaveGlow with fewer parameters.
b) One can consider adjusting the inner representation dimensions to increase the overall expressive power.
c)
Afﬁne coupling layers split the input z into two halves z a and z b .
Then, z b is used to predict the afﬁne transform parameters to process z a .
This split operation is needed to make the Jacobian determinant easily computable.
In the student model, one can remove the split operation and transform the entire input vector.
[Figure]
d) Finally, it might be helpful to substitute the entire afﬁne block with a commonly used feed-forward module for the corresponding task.
In our SRFlow and WaveGlow students, we adhere to the options b) and d).
Speciﬁcally, we tune the inner representation size and replace the ﬂow steps with stacked RRDBs (Wang et al., 2018) and WaveNet blocks (Oord et al., 2016), respectively.
The detailed descriptions of the student architectures are provided in
Appendix B.
The student design ablation on the example of WaveGlow is presented in Appendix F.
3. Experiments
In this section, we compare the learned student against the corresponding teachers and other relevant baselines from the literature.
3.1. Super-resolution
For super-resolution, we compare the student to the following set of models:
• ESRGAN (Wang et al., 2018), RankSRGAN (Mittal et al., 2013) — the recent GAN-based super-resolution models.
ESRGAN is trained by a combination of reconstruction, perceptual and adversarial losses.
RankSRGAN uses the additional ranker network to optimize non-differentiable perceptual metrics.
• RRDB
(Wang et al., 2018) — the same as ESRGAN but trained only with L 1 objective.
• SRFlow (Lugmayr et al., 2020a) — a ﬂow-based teacher model.
The model is optimized for efﬁcient inference by precomputing inversions for 1 × 1 invertible convolutions and fusing actnorm layers.
• SRFlow Student — a proposed student model designed according to Section 2.3.
Datasets .
The evaluation is performed on the DIV2K dataset (Agustsson & Timofte, 2017) — one of the established benchmarks for single image super-resolution.
In addition to the DIV2K train set, the Flickr2K dataset (Lim et al., 2017) is also used for training.
We consider × 4 and × 8 scaling factors between LR and HR images.
Evaluation metrics .
As the primary measure of perceptual performance, we report the established LPIPS
(Zhang et al., 2018b
), which is proven to correlate with human estimation (
Lugmayr et al., 2020b)
We also report the standard ﬁdelity-oriented metrics, Peak Signal to Noise Ratio (PSNR) and structural similarity index (SSIM)
(Wang
et al., 2004).
In addition, we evaluate consistency with the LR image by reporting the LR-PSNR, computed as PSNR between the downsampled
SR
image and the original LR image.
To evaluate diversity, we follow (Mao et al., 2019) and measure the average pairwise LPIPS distances between the generated samples for a particular LR image.
Larger pairwise distances indicate higher diversity.
Evaluation results.
We report metrics as well as corresponding inference time and a number of parameters in
Table 1.
For most baseline models, we reproduce the results on publicly released checkpoints.
Runtimes are measured on a single Tesla V100 on a single validation sample in half precision.
The SRFlow student achieves the same metric values as the teacher being by × 2 . 1 and × 4 . 9 faster for ×and × 8 scaling factors, respectively.
Also, we provide qualitative results for × 4 and × 8 scaling factors in the appendix, see Figure 4 and Figure 5, respectively.
The inference for the teacher and student models is performed for the same input noise vectors z .
We observe that the student produces samples with quality similar to the teacher.
Note that the student preserves the teacher diversity, which indicates that our distillation procedure is potentially able to produce student models with exploitable latent space (Lugmayr et al., 2020a).
[Table]
3.2. Speech Synthesis
Here, we consider the following set of models for speech synthesis evaluation:
• WaveGlow (Prenger et al., 2018) — a ﬂow-based teacher model.
For efﬁcient inference, we remove weight norms (Salimans & Kingma, 2016) and precompute inversions for convolutional layers.
• NanoFlow (Lee et al., 2020) — a recent ﬂow-based vocoder, which produces samples with comparable quality to WaveGlow, but has about × 30 fewer parameters.
• MelGAN (Kumar et al., 2019), Parallel WaveGAN (Yamamoto et al., 2020), HiFi-GAN (Kong et al., 2020) — recent GAN-based vocoders.
While MelGAN and ParallelWaveGAN are inferior to WaveGlow, HiFiGAN represents the current state-of-the-art.
• WG Student — a student model designed according to Section 2.3.
For evaluation, we consider three conﬁgurations with 4 / 4 / 2 WaveNet blocks of 128 / 96 / 96 hidden channels and denote them as V1/V2/V3, respectively.
Datasets .
All experiments are performed on the LJ speech dataset (Ito & Johnson, 2017), which is one of the most common benchmarks in speech synthesis.
We use a sampling rate of 22 , 050 kHz and produce mel-spectrograms of the original audio according to (Prenger et al., 2018).
While
Prenger et al.
(2018) originally provide train/validation/test splits for the LJ speech dataset, other methods are not always consistent with them.
Therefore, for a fair comparison of the pretrained models, we collect a novel evaluation set and provide its details in Appendix E.
Evaluation
results
We compare vocoders in the setting where models are conditioned on ground-truth melspectrograms as in previous works (Prenger et al., 2018;
Kim et al., 2019).
For all baseline models, we use ofﬁcially released pretrained models.
As a primary metric for generated audio evaluation, we report Mean Opinion Score
[Table]
(MOS) with corresponding 95 conﬁdence intervals (CI) in
Table 2.
In Appendix D, we describe the detailed protocol used for MOS evaluation.
The inference speed is reported in MHz, which stands for 10 6 audio samples per second.
The runtimes are measured on a single Tesla V100 in half precision with sufﬁciently large batch sizes and sequence lengths to suppress non-model related overheads.
V1 and V2 students demonstrate speech quality similar to the teacher and provide × 7 . 4 and × 9 .
7 faster speech generation, respectively.
Moreover, these students have × 9 . 2 and × 13 .
8 fewer parameters compared to the teacher.
4. Conclusion
In this work, we address the problem of high computational cost of normalizing ﬂows and explain how one can increase the efﬁciency by giving up invertibility and tractable Jacobian, which are often not necessary for deployed models.
In particular, we describe an effective knowledge distillation method from a ﬂow-based teacher to more lightweight student architectures.
We empirically demonstrate that the models distilled from normalizing ﬂows are an appealing combination of simplicity, stability and efﬁciency, needed for typical production pipelines.
A. Background & Related work
This section brieﬂy describes the basics of normalizing ﬂows and their applications in the image super-resolution and speech synthesis tasks.
Finally, we provide a short overview of knowledge distillation techniques.
A.1. Normalizing ﬂows
Normalizing ﬂows rely on change-of-variable formula to compute exact log likelihood as a sum of Jacobian logdeterminants.
In general, computing the Jacobian determinant of a transformation with N -dimensional inputs and outputs has O ( N 3 ) computational complexity, which is intractable for large N in typical deep learning applications.
Therefore, one of the challenges is to design invertible architectures with efﬁcient determinant calculation.
Currently, there are two main families of ﬂow-based models with easy-to-compute Jacobian determinants: based on autoregressive and bipartite transformations.
The autoregressive family includes autoregressive ﬂow (AF) (Papamakarios et al., 2017; Huang et al., 2018) and inverse autoregressive ﬂow (IAF) (Kingma et al., 2016).
AFs resemble autoregressive models (Uria et al., 2016), which allow for parallel density estimation but perform sequential synthesis.
In contrast, IAF poses parallel synthesis but sequential density estimation, making likelihood-based training very slow.
The second family includes bipartite transformations (Dinh et al., 2017; Kingma & Dhariwal, 2018), which provide both efﬁcient likelihood-based training and parallel synthesis.
However, bipartite ﬂows are usually less expressive than the autoregressive ones and hence the models require a larger number of layers and parameters.
Recent works develop more powerful and efﬁcient normalizing ﬂows.
Ping et al.
; Lee et al.
(2020) combine the ideas of autoregressive and bipartite transformations.
Others propose more expressive invertible layers and architectures (Ho et al., 2019;
Durkan et al., 2019; Karami et al., 2019; Chen et al., 2019; Behrmann et al., 2019) or continuous transformations (Grathwohl et al., 2019;
Chen et al., 2018).
At the moment, conditional NFs are gaining popularity in various practical speech and vision applications.
In particular, several recent works (Prenger et al., 2018;
Kim et al., 2019; Lugmayr et al., 2020a; Yang et al., 2019) exploit the ideas of Glow (Kingma & Dhariwal, 2018), RealNVP (Dinh et al., 2017) and continuous ﬂows (Chen et al., 2018) for waveform synthesis, image super-resolution, and point cloud generation.
In this work, we focus on the state-of-theart ﬂow-based models for image super-resolution (Lugmayr et al., 2020a) and speech synthesis (Prenger et al., 2018).
A.2. Super-resolution
Super-resolution (SR) is one of the fundamental image processing problems which aims to improve the quality of lowresolution (LR) images by upscaling them to high-resolution (HR) ones with natural high-frequency details.
Single image super-resolution approaches tend to either discover the most efﬁcient neural network architectures (Zhang et al., 2018d;c;a) or improve high-frequency details generation by introducing more sophisticated objectives (Ledig et al., 2017;
Wang et al., 2018; Mittal et al., 2013).
In contrast to these approaches, the SRFlow model (Lugmayr et al., 2020a) designs a ﬂow-based architecture that is trained by likelihood maximization.
As a result, SRFlow estimates a conditional distribution of natural HR images corresponding to a given LR image.
SRFlow is built upon the Glow (Kingma & Dhariwal, 2018) architecture that gradually transforms a sample from initial distribution z ∼ N (0 , I ) into the HR image.
The model is conditioned on representations of LR images delivered by a deterministic encoder (Wang et al., 2018).
A.3. Speech Synthesis
The state-of-the-art performance for speech synthesis is also achieved by deep generative models.
In typical speech synthesis pipelines, neural vocoders 3 play one of the most important roles.
Usually, a neural vocoder synthesizes timedomain waveform and is conditioned on mel-spectrograms from a text-to-spectrogram model (Arık et al., 2017; Ping et al., 2018;
Ren et al., 2019; 2020; Shen et al., 2018).
Long-standing state-of-the-art neural vocoders are autoregressive models (Oord et al., 2016; Kalchbrenner et al., 2018; Mehri et al., 2016) that provide impressive speech quality but suffer from slow sequential generation.
Therefore, a wide line of works aims to speedup their inference.
Flow-based models have been successfully applied for parallel waveform synthesis with ﬁdelity comparable to autoregressive models (van den Oord et al., 2018;
Prenger et al., 2018; Ping et al., 2020; Lee et al., 2020;
Kim et al., 2019; Ping et al., 2019; Serr`a et al., 2019).
Among ﬂowbased vocoders, WaveGlow (Prenger et al., 2018), WaveFlow (Ping et al., 2020), NanoFlow (Lee et al., 2020) and
FloWaveNet (Kim et al., 2019) have a particular emphasis on efﬁciency.
In this work, we speciﬁcally consider the WaveGlow model (Prenger et al., 2018) — one of the current state-of-the-art ﬂow-based vocoders.
Being efﬁcient enough, GAN-based vocoders (Kumar et al., 2019; Yamamoto et al., 2020) used to provide inferior speech ﬁdelity compared to the state-of-the-art ﬂow-based
[Figure]
Output
[Figure]
models.
However, recently proposed HiFi-GAN (Kong et al., 2020) efﬁciently delivers high-quality speech due to the training procedure with multi-scale (Kumar et al., 2019) and multi-period discriminators (Kong et al., 2020).
An alternative line of works proposes vocoders based on the denoising diffusion framework (Kong et al., 2021;
Chen et al., 2021).
These models demonstrate promising speech quality but are not efﬁcient enough for many usage cases.
A.4. Knowledge Distillation
Knowledge distillation is one of the most popular compression and acceleration techniques for large models and ensembles of neural networks (Hinton et al., 2015; Sau &
Balasubramanian, 2016; Romero et al., 2015).
Speciﬁcally, the idea of knowledge distillation is to train a single efﬁcient student on predictions produced by a computationally expensive teacher.
The closest works to ours distill knowledge from an expensive autoregressive neural vocoder to normalizing ﬂows with parallel inference.
In particular, Parallel WaveNet (van den
Oord et al., 2018) and ClariNet (Ping et al., 2019) transfer knowledge from a pretrained WaveNet (Oord et al., 2016) to
IAF (Kingma et al., 2016).
As opposed to these works, we address the computational inefﬁciency of ﬂow-based models and propose to transfer knowledge from normalizing ﬂows to feed-forward networks.
B. Student architectures
Below, we describe the architectures for WaveGlow and
SRFlow students in more detail and depict them in Figure(Left) and (Right), respectively.
B.1. SRFlow Student
LR encoder.
Both teacher and student models condition on the LR representations produced by the LR encoder.
The
LR encoder is a feed-forward SR architecture based on
Residual-in-Residual Dense Blocks (RRDB)
(Wang et al., 2018).
Originally, RRDBs are composed of 3 Residual
Dense Blocks (RDB) and each RDB consists of 5 convolutional layers.
In our student models, we vary the numbers of RDB blocks as well as convolutional layers within them.
The LR encoder is unchanged for × 4 scaling factor.
For ×scaling factor, the RRDB module is substituted by a single
RDB with 3 convolutional layers.
Base network.
The base architecture consists of L levels.
At each level l , the model combines the corresponding LR representation C l , the activations from the previous level h l and the noise vector z l by the Concat2d layer.
Here, one can vary h channel dimensions, according to Section 2.3 b).
Then, instead of a sequence of ﬂow steps in the teacher model, stacked RRDB modules follow.
For × 4 scaling factor, we use 6 RRDB modules of 2 RDBs with 3 convolutional layers.
For × 8 , a single original RRDB is used.
Finally, similarly to the SRFlow teacher, the transition step and unsqueeze operation are applied.
The transition step is Conv 1 × 1 followed by the Actnorm2d layer (Kingma & Dhariwal, 2018).
The unsqueeze operation doubles the spatial resolution of feature maps and reduces their channel dimension by a factor of 4 .
B.2. WaveGlow Student
WaveGlow student represents a sequence of conditional non-causal WaveNet blocks (Oord et al., 2016).
We start and ﬁnish with Conv 1 × 1 layers to adjust the input channel dimension.
Each WaveNet block is organized into 8 residual layers, each consisting of a dilated convolution followed by a gated activation unit (van den Oord et al., 2016) and 1 ×convolution.
The upsampled mel-spectrograms are added to the intermediate activations before the gated activation unit.
In contrast to the teacher, a student does not inject noise between intermediate WaveNet blocks but obtains the entire z at the very beginning.
C. Training details
SRFlow students.
The models are trained on patches of size 128 × 128 by the Adam optimizer (Kingma & Ba, 2014) with learning rate of 2 e − 4 for 100 epochs.
The learning rate is dropped by a factor of 2 at each 25 -th epoch.
The batch sizes of 80 and 128 are used for 4 and 8 scaling factors, respectively.
The loss coefﬁcient α is 10 .
WaveGlow students.
The models are optimized on short segments of length 16000 samples by the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 1 e − 4 for 200000 iterations.
For faster convergence, we use One Cycle learning rate schedule (Smith & Topin, 2018) withwarm up steps.
The batch size is 256 .
The loss coefﬁcient α is 1 e − 2 for most student settings.
All hyperparameters are tuned on the hold-out train samples, which do not participate in the ﬁnal evaluation.
The training is performed on 8 × Tesla V100.
D. Mean Opinion Score
We crowd-sourced MOS tests on 400 English-speaking individuals per model.
Each rater had to pass training onvalidation audio samples from various models with golden results and corresponding clariﬁcations.
Then, they listened to 10 audio samples from a single model and rated their naturalness on a ﬁve-point scale, see Table 3.
Each rater was asked to wear headphones and work in a quiet environment.
To aggregate MOS results for each model, we follow the protocol from (Kumar et al., 2019).
E. Evaluation set for Speech Synthesis
The collected set consists of 50 utterances trimmed from publicly available recordings by Linda Johnson from LibriVox.
We carefully choose recordings to be similar to the original dataset and normalize audio samples accordingly.
This test set will be publicly available online for reproducibility purposes.
[Table] [Table]
F. Student design ablation
To identify the core change to the architecture that improves the student performance, we ablate the student design principles described in Section 2.3 on the speech synthesis task.
In this experiment, we consider the student models with
WaveNet blocks of 96 channels and evaluate the following conﬁgurations in correspondence with Figure 2:
a) Flow Student — a student corresponds to the smaller WaveGlow model with all NF restrictions mentioned in Section 2.3.
b) Wide Flow Student — a ﬂow student where the inner representation channel dimension is increased from 8 to 96 .
Note that this model is not invertible anymore.
c) Afﬁne Student — the same as a ﬂow student but the partition operations are removed, see Figure 2c).
This student is still invertible but has O ( N 3 ) complexity of Jacobian determinant calculation.
d) WG Student — a student model where afﬁne coupling layers are replaced with WaveNet blocks.
According to Table 4, once any of the NF constraints are lifted, the corresponding student can achieve signiﬁcantly better speech quality compared to the ﬂow student of the same capacity.
Distilling
Knowledge
[Figure]
SRFlow Student
SRFlow Teacher
ESRGAN
HR
[Figure]
SRFlow Student
SRFlow Teacher
ESRGAN
HR
[Figure]
LR
Bicubic
SRFlow Student
SRFlow Teacher
ESRGAN
HR
[Figure]
Bicubic
SRFlow Student
SRFlow Teacher
ESRGAN
HR
[Figure]
LR
Bicubic
SRFlow Student
SRFlow Teacher
ESRGAN
HR
[Figure]
SRFlow
[Figure]
SRFlow
[Figure]
ESRGAN
HR
[Figure]
ESRGAN
HR
[Figure]
ESRGAN
HR
[Figure]
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher
HR
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher HR
IMG
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher
HR
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher
HR
IMG
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher
HR
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher
HR
IMG
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher
HR
LR
Bicubic
RRDB
SRFlow Student
SRFlow Teacher
HR
IMG