O N D IFFICULTIES OF P ROBABILITY D ISTILLATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Probability distillation has recently been of interest to deep learning practitioners as it presents a practical solution for sampling from autoregressive models for deployment in real-time applications.
We identify a pathological optimization issue with the commonly adopted stochastic minimization of the (reverse) KL divergence, owing to skewed gradient distribution due to curse of dimensionality.
We also explore alternative principles for distillation, and show that one can achieve qualitatively better results than with KL minimization.
I NTRODUCTION
Deep autoregressive models are currently among the best choices for unsupervised density modeling tasks (Van den Oord et al., 2016b;a)
However, due to the factorization induced by such models on the data space by the chain rule of probability, sampling from such models requires O ( T ) sequential computation steps, where T is the dimension/sequence-length of the data.
This makes sampling slow and inefﬁcient for most practical purposes.
A recent solution to alleviating this bottleneck was proposed by Van den Oord et al.
(2018), who show how to distill an autoregressive WaveNet model (Van den Oord et al., 2016a)
into a student network , which is signiﬁcantly faster to sample from, and therefore much more suitable for deployment in real-time applications.
In this paper, we identify a fundamental issue with the approach taken by Van den Oord et al.
(2018), and provide alternative views and principles for distilling the sampling process of a teacher model.
The solution proposed by Van den Oord et al.
(2018) is to minimize the reverse Kullback-Leibler divergence (KL) between the student and the teacher distribution.
Generally, this relies on two essential components: (1) the gradient signal from the teacher network (the WaveNet model) and (2) invertibility of the student network (the Parallel WaveNet model).
This allows samples drawn from the student to be evaluated under the likelihood of both models, and the student can then be updated via stochastic backpropagation from the teacher.
The student is implemented with an inverse autoregressive transformation (Kingma et al., 2016) that, unlike an autoregressive model, admits fast sampling at the cost of slow likelihood estimation.
In theory, it seems that minimizing the KL should be sufﬁcient for distilling a teacher into a student.
However, in practice, it has been necessary to implement additional heuristics in order to achieve samples with similar levels of realism as that of the teacher (Van den Oord et al., 2018; Ping et al., 2018).
In particular, a power loss that biases the power across frequency bands of the student sampler to match that of human speech patterns has been crucial in recent work for generating synthesized speech without the student collapsing to unrealistic ‘whispering’.
The main contribution of this work is to identify that the reverse KL is ill-suited for the task of distillation, since the teacher distribution typically possesses a low-rank nature for high-dimensional structured data, which renders the effective gradient signal sparse.
We explore some possible alternatives for distillation, by recasting the problem of distilling a generative model as learning a transformation of probability density from a prior space to the data space, connecting decoderbased generative models with probability density distillation.
We present experimental results which demonstrate that it is possible to learn qualitatively better samplers than distillation with reverse KL.
A UTOREGRESSIVE M ODELS
Given a joint probability distribution p ( x ) , where x denotes a T -dimensional vector, one can factorize the distribution according to the chain rule of probability , for arbitrary ordering of dimensions:
[Equation]
For the tabular case, i.e. when x t can take V different possible values, the size of the table is
O ( V T ) .
When the event set of x t is uncountable, the joint density is not even tractable.
This motivates the use of a parametric model to compress the conditional probability p ( x t | x 1: t − 1 ) , where one has p θ ( x ) =
(cid:81)
Tt =1 p θ ( x t | x 1: t − 1 ) .
This is referred to as an autoregressive model .
Parameters are usually shared across the dimensions of x , since T might not stay constant, for example in the case of recurrent neural networks or convolutional neural networks, which have been empirically demonstrated to possess good inductive biases for tasks involving images (Van den Oord et al., 2016b) and speech data (Van den Oord et al., 2016a).
However, sampling is sequential, requiring T passes, which is why state-of-the-art autoregressive models have been slow to sample from, which makes them impractical for tasks requiring sampling, such as speech generation.
This has motivated the work of Van den Oord et al.
(2018), who propose probability density distillation to learn a student network with a structure that allows for parallel sampling, by distilling a state-of-the-art autoregressive teacher into it.
P ROBABILITY D ISTILLATION WITH N ORMALIZING F LOWS
Van den Oord et al.
(2018) propose to distill the probability distribution parameterized by a WaveNet model (denoted by T , which stands for the teacher network) by minimizing its reverse KL with a student network (denoted by S ):
[Equation]
The idea is to leverage recent advances in change of variable models (also known as normalizing ﬂows ) to parallelize the computation of the sampling process.
First, the student distribution is constructed by transforming an initial distribution p S ( z ) (e.g. normal or uniform distribution) in a way such that each dimension x t in the output x depends only on up to t preceding variables (according to a chosen ordering) in the input z :
[Equation]
where g t is an invertible map between z t and x t .
Unlike the sampling process of an autoregressive model, where one needs to accumulate all x 1: t − 1 to sample x t from p T ( x t | x 1: t − 1 ) , which scales O ( t ) , the transformations g t can be carried out independently of t , allowing for O (1) time sampling.
Second, the entropy term of p S can be estimated using the change of variable formula:
[Equation]
where g is the overall transformation g ( z ) t .
=
g t ( z t ; z <t ) .
Furthermore, owing to the partial dependency of g t , g has a triangular Jacobian matrix, reducing the computation of the log-determinant of
Jacobian in (2) to linear time:
[Equation]
Finally, when the teacher network has a tractable explicit density, one can evaluate the likelihood of samples drawn from p S under p T efﬁciently (in the case of autoregressive models such as WaveNets, one can use teacher forcing to compute log-likelihood in parallel).
[Figure]
3. C URSE OF D IMENSIONALITY
In this section, we explain the failure mode of distillation with reverse KL that we have identiﬁed.
Intuitively, for each sample x drawn from p S , the negative gradient of the integrand in Equation (1) with respect to x , is made up of two counteracting factors: one that pushes x away from the mode of the student density p S and one that pulls x towards the mode of the teacher density p T .
These two counteracting terms form the “error” component of the path derivative (Roeder et al., 2017) when the student is updated. (see Appendix C for more details).
Assume the student density p S lies within a certain family Q .
Ideally, if p T ∈ Q , the solution to the minimization problem in (1) would be p S = p T .
However, this is not trivial in practice when an oracle solver is not accessible.
We empirically ﬁnd out that training with reverse KL consists of two stages:
(i) p S starts to ﬁt to the mode of p T (Figure 2a)and (ii) p S gradually expands from the mode of p T to ﬁt the shape of the distribution
Stage (i) is fast due to the well-known zero forcing property of the reverse KL.
Turner & Sahani (2011) (see Fig 1.3) shows that p S tends to be more concentrated when independence assumption is made. We show that even when Q contains p T , stochastic optimization can result in slow convergence of stage (ii) and thus a more concentrated, suboptimal p S . This is because the gradient signal via the path derivative can be increasingly unlikely to be effective when the dimensionality in x grows and when p T is highly structured. This makes it harder for p S to expand its probability mass along the high density manifold under p T .
We illustrate this as follows. Consider the case when both the student and teacher are multivariate normal centered at the origin: assume without loss of generality 1 that p S = N (0 , I ) and p T = N (0 , Σ T ) (both centered atassuming stage (i) has been completed), where I is an n -by- n identity matrix and Σ T ∈ S n ++ is a positive deﬁnite matrix.
The following proposition establishes the connection between the eigenvalues of the covariance matrix Σ T and the probability of getting a positive signal that pushes x away from the origin.
Proposition 1.
Let p S =
N (0 , I ) and p T (0 , Σ T ) .
Draw x ∼ p S .
Let A U be the surface area of the unit sphere: U = { x : (cid:107) x (cid:107) 2 = 1 } and A U ∩ ρ be the surface area of { x ∈ U : (cid:80) i ρ i x 2 i > 0 } .
Then the probability of ∇ x [log p T − log q S ] pointing away from the origin along x is given by
[Equation]
[Figure]
where ρ i = 1 − 1 d 2 i and d 2 i is the i -th eigenvalue of the covariance matrix.
Proof.
Let g x .
= ∇ x (log p T ( x ) − log p S ( x )) be the gradient. Since both p S and p T are Gaussian distributions centered at the origin, the gradient when projected onto x , either points towards or away from the origin, depending on whether x (cid:62) g x < 0 or x (cid:62) g x > 0 .
By deﬁnition, we have g x = − Σ − 1 T x + x .
Let Σ T =
Λ D Λ − 1 be the eigen-decomposition of the covariance, where D ii =
d 2 i is the i -th eigenvalue and the columns of Λ are the eigenvectors.
Due to the rotational invariance and uniformity of the density on the level set { x : (cid:107) x (cid:107)
= r } for any r > 0 of the standard normal p S ,
[Equation]
What the proposition implies is that the chances of receiving a gradient signal that points outward depend on the eigenvalues of the covariance matrix of the teacher: the greater the number of eigenvalues that are smaller than 1 (more ill-conditioned), the lower the chances.
Consider Figure 1a for example, where the red contour plot and blue contour plot represent the density of p T and p S , respectively.
For a random sample drawn from p S , marked by the yellow star, the gradients of log p T and − log p S with respect to it are represented by the red arrow and blue arrow.
The net g x here can be decomposed into two parts: one that is perpendicular to x , g x, ⊥ , and one that is parallel with x , g x, (cid:107) .
In this example, x and g x, (cid:107) point towards the opposite direction, meaning the back-propagated signal would draw x towards the mode of p T .
On average, the chances of getting a stochastic gradient signal that push the points away from the mode is the percentage of the area of the unit sphere intersecting with the hypercone, represented by the shaded area in Figure 1b.
In practice, such a condition coefﬁcient can be very small, as it is well known that a high dimensional distribution over structured data is effectively low-rank.
In fact, an almost-sure convergence result of the smallest eigenvalue of a random Wishart matrix scaled by 1 /T was proven by Silverstein et al.
(1985) to be effectively zero with large enough T .
The Marchenko-Pastur Law describes a more general asymptotic distribution of the eigenvalues (Marchenko & Pastur, 1967).
3. E MPIRICAL DEMONSTRATION
We showed above that with increasing dimensionality, and for a structured teacher, there is a diminishing probability of gradient signals that can push the student to expand around the mode of the teacher. It seems intuitive that the distilled density of the student will therefore be collapsed around the mode of the teacher density.
We validate this hypothesis in the following experiment.
We take both p T and p S to be multivariate Gaussian distributions, with the sampling process deﬁned as x ← µ + R · z where µ ∈ R T , R ∈ R T × T and z ∼ N (0 , I ) .
We randomly initialize each element of R for T independently according to the standard Gaussian, set µ = [2 , ..., 2] (cid:62) to be a vector of T 2 ’s, and ﬁx them while training S to distill T .
For T ∈ [4 , 16 , 32 , 64] , training proceeds as follows: we sample x from the student, estimate log p S ( x ) using the change of variable formula, and evaluate x under log p T .
We use a minibatch size of 64 and learning rate of 0 .
005 with the Adam optimizer (Kingma & Ba, 2014), and make 5000 updates.
For evaluation, we draw 1000 samples from both p T and p S , and display the empirical distribution of log p T ( x ) in Figure 2a.
First, with increasing number of dimensions, we observe that p S does indeed concentrate more on the high density region of p T .
This suggests that the unbalanced gradient signal poses an optimization problem for distillation of higher dimensional structured distributions; getting a sample that gets pushed away along the thin manifold under the teacher density is very unlikely, which is consistent with Proposition 1.
Second, we also observe that the log-likelihood of the teacher samples deviate from 0 as dimensionality grows.
In fact, assuming x t is sampled i.i.d.
, the l 2 norm of ¯ x . = x √ T would almost surely converge.
This is a phenomenon known as the concentration of measure .
To see this,[-(cid:107) ¯ x (cid:107) 
 = ¯ x (cid:62) ¯ x = T (cid:88) t = ¯ x 2 t = T (cid:88)-]
[Equation]
The concentration is due to the compromise between density and volume of space (which vanishes exponentially as dimensionality grows).
The consequence is that when one samples from a high dimensional Gaussian, the norm of the sample can be well described by its expected value, which means one is effectively sampling from the shell of the Gaussian ball.
This suggests that the use of
KL would result in a mismatch of certain important statistics (such as norm of the samples, which is a perceivable feature in images and audio frames) even when p S is fairly close to p T .
Finally, in the above study, we only identify this optimization difﬁculty in the convex setup.
However, it is also well known that the reverse KL tends to be mode-seeking (see Figure 3b,3c for example), and is not well-suited for learning multimodal densities (Turner & Sahani, 2011; Huang et al., 2018b).
P ROBABILITY D ISTILLATION WITH I NVERSE M ATCHING
In this section, we discuss possible alternatives for distilling a teacher.
We assume there exists an invertible mapping from a prior space Z to the data space X , such that one can trivially sample from a prior distribution z ∼ p T ( z ) and pass the sample through this invertible map such that the sample is distributed according to p T ( x ) .
For Gaussian conditional autoregressive models, for example, one would sequentially pass scalar standard Gaussian noise z t through the following recursive function x t = µ t ( x 1: t − 1 ) + σ t ( x 1: t − 1 ) · z t .
For notational convenience, we denote the “inverse” of this transformation by T : X → Z , as this is the inverse autoregressive transformation that can be parallelized.
The goal is to learn the sampling transformation, i.e. T − 1 .
Similarly, we deﬁne the forward pass of the student as the mapping S : Z → X .
Ideally, since the transformation is deterministic, it’s most natural to simply minimize the prediction loss according to some distance metric d ( T − 1 ( z ) , S ( z )) , where z ∼ p T ,z .
When this loss function equals zero almost everywhere, passing the prior sample through S would induce an identical distribution as p T .
We refer to this setup as distillation with oracle prediction .
However, preparing such a dataset of T − 1 ( z ) samples would typically be time-consuming.
We present the following two alternatives.
Distillation with z -reconstruction.
We consider minimizing d ( z, T ◦ S ( z )) , which is a reconstruction loss and the student network and teacher network are viewed as the encoder and decoder, respectively.
In this case, since T is invertible and ﬁxed, the only functional form of S that gives zero reconstruction would be T − 1 , which means the random variable S ( Z ) should also be distributed according
to p T .
In fact, minimizing the z -reconstruction loss corresponds to a parametric distance induced by the teacher network.
Deﬁne d T ( a, b ) .
=
d ( T ( a ) , T ( b )) , where d is a distance metric.
Then
[Equation]
Interestingly, d T is also a metric:
Proposition 2.
d T is a metric if and only if T is injective.
Proof.
Trivially, positive-deﬁniteness and symmetry are inherited from d if and only if T is an injection.
To see that subadditivity is also preserved, for some a , b and c , let T a = T ( a ) , T b = T ( b )
[Figure] and T c = T ( c ) .
Since d ( T a , T b ) ≤ d ( T a , T c ) + d ( T b , T c ) , due to the subadditivity of d , for any T a ,
T b and T c , we have d T ( a, b ) ≤ d T ( a, c )
+ d T ( b, c ) for any a , b and c .
This means z-reconstruction loss behaves like a distance between T − 1 ( z ) and S ( z ) .
So when zreconstruction is minimized, it implies S gets closer to T − 1 in the sense of the induced metric d T .
Distillation with x -reconstruction.
Finally, we consider minimizing the reconstruction loss
d ( x, S ◦ T ( x )) , where x ∼ p D , the (empirical) data distribution, taking the teacher network as the encoder, and the student network as the decoder.
When the teacher density coincides with the underlying data distribution, this would be equivalent to training with oracle prediction, as T ( X ) would be distributed according to p T ( z ) .
This is a reasonable assumption when p T approximates p D well, and this is in fact true as p T is usually trained with Maximum likelihood under p D .
Now we revisit the two essential components required for distillation with reverse KL:
(1) Invertibility :
None of the three training
criteria we explored involves estimating the entropy of p S , so in principle, we do not require invertibility of the student.
In fact the entropy of p S is implicitly maximized since T is bijective.
To prevent degenerate p S , one simply needs to avoid using hidden units of dimensionality smaller than the input size without skip connectivity, which compresses the noise.
(2) Differentiability :
For distillation with oracle prediction and x -reconstruction, we only require the translation between X and Z , via T and T − 1 , which is readily accessible for many standard distributions, e.g. linear map between Gaussians (both T and T − 1 ), logistic-linear map from mixture of logistics to uniform ( T only, but sampling is achievable by sampling the mixture component ﬁrst), and neural transformation (Huang et al., 2018a)
( T only).
We also note that it is possible to recover uniform density from discrete data, by injecting noise proportional to the probability per class to break ties when passing the data through the cumulative sum of the probability (CDF). [Figure]
4. E XPERIMENTS 4.1. L INEAR M ODEL WITH I NCREASING D IMENSIONALITY
We replicate the experiment in Section 3.2 with z -reconstruction and x -reconstruction loss (equivalent to oracle prediction in this case).
Mapping T from X to S is simply inverse of the sampling transformation.
We observe that both models outperform distillation with reverse KL, since likelihood of samples (under the teacher) drawn from them follow more closely the shape of the empirical distribution of likelihood of samples drawn from the teacher.
It is worth noting that z -reconstruction also starts to fail with an increasing number of dimensions, suggesting that it might be subject to poor gradient signal if the transformation T is more complex; we elaborate more on this in the next section.
On the other hand x -reconstruction is quite robust to dimensionality (and potentially complexity) of the underlying distribution.
4.1. D ISTILLATION WITH D IFFERENT L OSSES
In this section, we distill PixelCNN++ (Salimans et al., 2017) teacher networks trained on the
MNIST handwritten digits dataset (LeCun et al., 1998) and the Fashion-MNIST dataset (Xiao et al., 2017).
We trained the teacher model for 100 epochs and distilled it into the student with anotherepochs of updates, using minibatch size of 64, learning rate of 0.0005 for the Adam optimizer with a decay rate of 0.95 per epoch, 3 ResNet blocks per downsampling and upsampling convolution,hidden channels, and a single Gaussian conditional.
The data is preprocessed with uniform noise between pixel values and rescaled using the logit function.
We use l 2 loss for the reconstruction and prediction methods.
First, we observe that when trained with the reverse-KL loss, the students collapse on undesirable modes.
As shown in (Figure 3b), the inductive bias of the causal convolution leads to higher density of the samples with striped textures.
When trained with the z -reconstruction loss, the MNIST student samples all collapse to the same digit.
Interestingly, when we visualize the corresponding T − 1 ( z ) as we slightly perturb the norm of z (see Appendix A), we observe that the digits abruptly change identity.
This suggests that when moving onto a different sublevel set of norm in z -space, the corresponding x jumps from one digit manifold to another, and the direction of z does not preserve digit identity.
This might explain why the student collapses to a digit: this is due to the bad local minimum that corresponds to relatively low reconstruction cost in the z -space.
Next, we observe that the student trained with x -reconstruction loss (Figure 3d) does not have good quality samples while the reconstructions are visually perfect.
We hypothesize this is due to the well-known problem of mismatch between the empirical distribution of the encodings T ( x ) and the prior distribution p S ( z ) of training decoder based generative models (Kingma et al., 2016).
We contrast this with a student trained on oracle predictions (Figure 3e) and observe that the latter’s samples match the teacher’s samples better.
Finally, we see that the samples from the teacher trained on Fashion-MNIST with the xreconstruction loss (Figure 3i) have a smoother texture than the one trained with oracle samples (Figure 3j), which again, are perceptually closer to samples from the teacher.
We elaborate more on this discrepancy in the next section.
4.1. L EARNING TO D ISTILL AND L EARNING TO G ENERATE
Since we are not constrained in our modeling choice for the student, we experiment with a ResNet student which is trained to directly map an encoded datapoint from the teacher back to the datapoint.
The ResNet is deep enough so that the receptive ﬁeld at the output is sufﬁcient to span all of the encoded z , so that far-away inﬂuence is still exploitable.
Using the l 1 reconstruction cost leads to sharper samples from the student (contrast 4a with 4b).
It appears to us that the l 2 loss tends to maintain global details, while the l 1 loss can sometimes sacriﬁce global coherence for local structure, potentially due to the sparsity induced by it.
In 4e, we use an average of both losses in an attempt to maintain both these characteristics, but we notice evidence of the failings manifesting to some extent as well.
A signiﬁcant improvement in sample quality occurs upon adding Gaussian noise to the encodings before training the student (contrast 4a,4b with 4c,4d).
Our intuition for this is as follows: providing a decoder network with pairs of points in a data space and an encoded (or latent) space would typically result in almost all of z -space not being “trained”, especially since image data (and therefore the corresponding encodings in z -space) usually lies in a low-dimensional manifold.
Adding noise enhances the support of the distribution, effectively spreading the distribution of “responsibility” of an encoding to cover more volume, which smoothens the mapping learned by a decoder.
When the goal is to sample from a prior, training methods that encourage such z -space-ﬁlling strategies and smoother mappings improve sample quality, in the vein of decoder-based sampling models such as variational auto-encoders (Kingma & Welling, 2014).
This leads us to an important point about these experiments: since the student is trained on noised (encoded) points from the data distribution, this is no longer purely density distillation.
The student no longer aims to reproduce the sampling behavior of the teacher (as in 3e), but rather uses the teacher to provide structural information through its encodings.
This information, when “spread out” through noise-injection and used by a student to learn decodings into real data (through a reconstruction penalty or more sophisticated losses) results in a network that can now be considered as a stand-alone generator, with the teacher acting as an inference machine that preserves information in the latent space.
This can potentially allow a student to outperform its teacher in terms of sample quality, by enabling the learning of a smoother mapping from z -space to data space.
4.1. N EURAL V OCODER
We compare distillation with x -reconstruction and the reverse-KL approach on the neural vocoder Sotelo et al.
(2017) task for speech synthesis.
The neural vocoder is an essential component of many text-to-speech models proposed recently (Wang et al., 2017;
Shen et al., 2018; Arik et al., 2017; Ping et al., 2018).
We train our teacher to map vocoders (Morise et al., 2016) to raw audio using the SampleRNN (Mehri et al., 2016) model.
We model the conditional distribution of the teacher with a unimodal gaussian distribution, making it easier to compute the corresponding z .
We speciﬁcally compare against closed form regularized KL with Gaussian conditionals as proposed in (Ping et al., 2018).
The student network has a WaveNet architecture with six ﬂows, and performs sampling as in Parallel WaveNet (Van den Oord et al., 2018).
Each ﬂow is a dilated residual block of 10 layers with a convolution kernel width of 2 and 64 output channels.
We use our x -reconstruction method with l 1 as the reconstruction loss.
We empirically ﬁnd that our method results in a student with samples without the characteristic whispering of the reverse-KL trained student.
We have uploaded samples for comparison here 2 .
C ONCLUSION
In this paper, we investigate problems with distilling an autoregressive generative model under a reverse KL cost between the student and the teacher, where the student can perform efﬁcient parallel generation.
Speciﬁcally we show that distillation with the reverse KL can suffer from imbalanced gradient signals due to the curse of dimensionality.
Further, we explore different alternatives which work qualitatively better when compared with distillation with reverse KL.
A V ISUALIZATION OF T − 1 ( z ) WITH PERTURBED NORM OF z
B A DDITIONAL RESULT ON DISTILLING P IXEL CNN TRAINED ON CIFARWITH X - RECONSTRUCTION ( NOISE INJECTION )
C P ATH DERIVATIVE
When x is real-valued and when reparameterization of the sample drawn from p S ( x ) allows for separation of a random variable z ∼ p ( z ) independent of the parameters of p S and a deterministic transformation x = S φ ( z )
[Equation]
we can decompose the gradient of the KL divergence (1) with the parameters φ (of S ) as (Roeder et al., 2017):
[Equation]
] where the last equality is the total derivative because the term p S φ ( x ) depends on φ through both the sample being evaluated, x = S φ ( z )
[Equation]
and the evaluating log-likelihood function, log p S φ .
The score function is in expectation zero, since
[Equation]
so it can be thought of as a control variate: an unbiased term used to reduce variance in gradient estimate.
The path derivative term measures the dependency on the parameters φ through the reparameterized sample x = S φ ( z ) .
As a result, the gradient direction wrt the sample x directly affects how the parameters of the distribution will be updated to change the shape of p S φ .