The approach of training sequence models using supervised learning and next-step prediction suffers from known failure modes. For example, it is notoriously difficult to ensure multi-step generated sequences have coherent global structure. We propose a novel sequence-learning approach in which we use a pre-trained Recurrent Neural Network (RNN) to supply part of the reward value in a Reinforcement Learning (RL) model. Thus, we can refine a sequence predictor by optimizing for some imposed reward functions, while maintaining good predictive properties learned from data. We propose efficient ways to solve this by augmenting deep Q-learning with a cross-entropy reward and deriving novel off-policy methods for RNNs from KL control. We explore the usefulness of our approach in the context of music generation. An LSTM is trained on a large corpus of songs to predict the next note in a musical sequence. This Note-RNN is then refined using our method and rules of music theory. We show that by combining maximum likelihood (ML) and RL in this way, we can not only produce more pleasing melodies, but significantly reduce unwanted behaviors and failure modes of the RNN, while maintaining information learned from data.
1 INTRODUCTION
Generative modeling of music with deep neural networks is typically accomplished by training a RNN such as a Long Short-Term Memory (LSTM) network to predict the next note in a musical sequence (e.g. Eck & Schmidhuber (2002)). Similar to a Character RNN (Mikolov et al., 2010), these Note RNNs can be used to generate novel melodies by initializing them with a short sequence of notes, then repeatedly sampling from the model’s output distribution generated to obtain the next note. While melodies and text generated in this way have recently garnered attention1, this type of model tends to suffer from common failure modes, such as excessively repeating tokens, or producing sequences that lack a consistent theme or structure. Such sequences can appear wandering and random (see Graves (2013) for a text example).
Music compositions adhere to relatively well-defined structural rules, making music an interesting sequence generation challenge. For example, music theory tells that groups of notes belong to keys, chords follow progressions, and songs have consistent structures made up of musical phrases. Our research question is therefore whether such music-theory-based constraints can be learned by an RNN, while still allowing it to maintain note probabilities learned from data.
To approach this problem we propose RL Tuner, a novel sequence learning approach in which RL is used to impose structure on an RNN trained on data. The reward function in our framework combines task-related rewards with the probability of a given action originally learned by the pre-trained RNN. Thus, our model directly preserves inforamtion about the original probability distributions learned from data, while allowing us to explicitly control the trade-off between the influence of data
1http://www.theverge.com/2016/6/1/11829678/google-magenta-melody-art-generative-artificialintelligence
and heuristic rewards. This is an important novel direction of research, because in many tasks the available reward functions are not a perfect metric that alone will lead to the best task performance in the real world (e.g. BLEU score). Unlike previous work (e.g. (Ranzato et al., 2015), (Bahdanau et al., 2016), (Norouzi et al., 2016), (Li et al., 2016)) we do not use ML training as a way to simply bootstrap the training of an RL model, but rather we rely mainly on information learned from data, and use RL only as a way to refine characteristics of the output by imposing structural rules.
This paper contributes to the sequence training and RL literature by a) proposing a novel method for combining ML and RL training; b) showing the connection between this approach and Stochastic Optimal Control (SOC)/KL-control with a pre-trained RNN as a prior policy; c) showing the explicit relationships among a generalized variant of Ψ-learning (Rawlik et al., 2012), G-learning (Fox et al.), and Q-learning with log prior augmentation; d) being the first work to explore generalized Ψ-learning and G-learning with deep neural networks, serving as a reference for exploring KLregularized RL objectives with deep Q-learning; e) empirically comparing generalized Ψ-learning, G-learning, and Q-learning with log prior augmentation for the first time; and f) applying this new technique to the problem of music generation, and showing through an empirical study that this method produces melodies which are more melodic, harmonious, interesting, and rated as significantly more subjectively pleasing, than those of the original Note RNN. We suggest that the RL Tuner method could have potential applications in a number of areas as a general way to refine existing recurrent models trained on data by imposing constraints on their behavior.
2 BACKGROUND
2.1 DEEP Q-LEARNING
In RL, an agent interacts with an environment. Given the state of the environment at time t, st, the agent takes an action at according to its policy π(at|st), receives a reward r(st, at), and the environment transitions to a new state, st+1.The agent’s goal is to maximize reward over a sequence of actions, with a discount factor of γ applied to future rewards. The optimal deterministic policy π∗ is known to satisfy the following Bellman optimality equation,
Q(st, at;π ∗) = r(st, at) + γEp(st+1|st,at)[maxat+1 Q(st+1, at+1;π ∗)] (1)
where Qπ(st, at) = Eπ[ ∑∞ t′=t γ
t′−tr(st′ , at′)] is the Q function of a policy π. Q-learning techniques (Watkins & Dayan, 1992; Sutton et al., 1999) learn this optimal Q function by iteratively minimizing the Bellman residual. The optimal policy is given by π∗(a|s) = arg maxaQ(s, a). Deep Q-learning(Mnih et al., 2013) uses a neural network called the deep Q-network (DQN) to approximate the Q function Q(s, a; θ). The network parameters θ are learned by applying stochastic gradient descent (SGD) updates with respect to the following loss function,
L(θ) = Eβ [(r(s, a) + γmax a′
Q(s′, a′; θ−)−Q(s, a; θ))2] (2)
where β is the exploration policy, and θ− is the parameters of the Target Q-network (Mnih et al., 2013) that is held fixed during the gradient computation. The moving average of θ is used as θ− as proposed in (Lillicrap et al., 2016). Exploration can be performed with either the -greedy method or Boltzmann sampling. Additional standard techniques such as replay memory (Mnih et al., 2013) and Deep Double Q-learning (Hasselt et al., 2015) are used to stablize and improve learning.
2.2 MUSIC GENERATION WITH LSTM
Previous work with music generation using deep learning (e.g. (Eck & Schmidhuber, 2002), (Sturm et al., 2016)) has involved training an RNN to learn to predict the next note in a monophonic melody; we call this type of model a Note RNN. Often, the Note RNN is implemented using a Long ShortTerm Memory (LSTM) network (Gers et al., 2000). LSTMs are networks in which each recurrent cell learns to control the storage of information through the use of an input gate, output gate, and forget gate. The first two gates control whether information is able to flow into and out of the cell, and the latter controls whether or not the contents of the cell should be reset. Due to these properties, LSTMs are better at learning long-term dependencies in the data, and can adapt more rapidly to new data (Graves, 2013). A softmax function can be applied to the final outputs of the network to obtain
the probability the network places on each note, and softmax cross-entropy loss can be used to train the model via back propagation through time (BPTT) (Graves & Schmidhuber, 2005). However, as previously described, the melodies generated by this model tend to wander, and lack musical structure; we will show that they are also perceived as less musically pleasing by listeners. In the next section, we will show how to improve this model with RL.
3 RL TUNER DESIGN
Given a trained Note RNN, the goal is to teach it concepts about music theory, while still maintaining the information about typical melodies originally learned from data. To accomplish this task, we propose RL Tuner, a novel sequence training method incorporating RL. We use an LSTM trained on data (the Note RNN) to supply the initial weights for three networks in RL Tuner: the Q-network and Target Q-network in the DQN algorithm as described in Section 2.1, and a Reward RNN. Therefore, the Q-network is a recurrent LSTM model, with architecture identical to that of the original Note RNN. The Reward RNN is used to supply part of the reward value used to train the model, and is held fixed during training.
In order to formulate music generation as an RL problem, we treat placing the next note in the melody as taking an action. The state of the environment s consists of the previous note, and the internal state of the LSTM cells of both the Q-network and the Reward RNN. Thus, Q(a, s) can be calculated by initializing the recurrent Q-network with the appropriate memory cell contents, running it for one time step using the previous note, and evaluating the output value for the action a. The next action can be selected with either a Boltzmann sampling or -greedy exploration strategy.
Given action a, the reward can be computed by combining probabilities learned from the training data with knowledge of music theory. We define a set of music-theory based rules (described in Section 3.2) to impose constraints on the melody that the model is composing through a reward signal rMT (a, s). For example, if a note is in the wrong key, then the model receives a negative reward. However, it is necessary that the model still be “creative,” rather than learning a simple melody that can easily exploit these rewards. Therefore, we use the Reward RNN — or equivalently the trained Note RNN — to compute log p(a|s), the log probability of a note a given a melody s, and incorporate this into the reward function. Figure 1 illustrates these ideas.
The total reward given at time t is therefore:
r(s, a) = log p(a|s) + rMT (a, s)/c (3)
where c is a constant controlling the emphasis placed on the music theory reward. Given the DQN loss function in Eq. 2 and modified reward function in Eq. 3, the new loss function and learned policy for RL Tuner are,
L(θ) = Eβ [(log p(a|s) + rMT (a, s)/c+ γmax a′
Q(s′, a′; θ−)−Q(s, a; θ))2] (4)
πθ(a|s) = δ(a = arg max a Q(s, a; θ)). (5)
Thus, the modified loss function forces the model to learn that the most valuable actions are those that conform to the music theory rules, but still have high probability in the original data.
3.1 RELATIONSHIP TO KL CONTROL
The technique described in Section 3 has a close connection to stochastic optimal control (SOC) (Stengel, 1986) and in particular, KL control (Todorov, 2006; Kappen et al., 2012; Rawlik et al., 2012). SOC casts the optimal planning in stochastic environments as inference in graphical models, and enables direct application of probabilistic inference techniques such as ExpectationMaximization (EM) and message passing for solving the control problem (Attias, 2003; Toussaint & Storkey, 2006; Toussaint, 2009). Rawlik et al. (2012); Kappen et al. (2012) then introduced KL control, a generic formulation of the SOC as Kullback-Leibler (KL) divergence minimization, and connected to prior work on RL with additional KL cost (Todorov, 2006). Since our primary focus is to connect with DQNs, we specifically focus on the work by Rawlik et al. (2012) as they derive a temporal-difference-based approach on which we build our methods.
KL control formulation defines a prior dynamics or policy, and derives a variant of the control or RL problem as performing approximate inference in a graphical model. Let τ be a trajectory of state and action sequences, p(τ) be a prior dynamics, and r(τ) be the reward of the trajectory. Then, an additional binary variable b is introduced and a graphical model is defined as p(τ, b) = p(τ)p(b|τ), where p(b = 1|τ) = er(τ)/c and c is the temperature variable. An approximation to p(τ |b = 1) can be derived using the variational free-energy method, and this leads to a cost with a similar form to the RL problem previously defined, but with an additional penalty based on the KL divergence from the prior trajectory,
log p(τ |b = 1) = log ∫ p(τ)p(b|τ)dτ (6)
≥ Eq(τ)[log p(τ)p(b|τ)− log q(τ)] (7) = Eq(τ)[r(τ)/c− KL[q(τ)||p(τ)]] = Lv(q) (8)
where q(τ) is the variational distribution. Rewriting the variational objective Lv(q) in Eq. 6 in terms of policy πθ, we get the following RL objective with KL-regularization, also known as KL control,
Lv(θ) = Eπ[ ∑ t r(st, at)/c−KL[πθ(·|st)||p(·|st)]]. (9)
In contrast, the objective in Section 3 is, Lv(θ) = Eπ[ ∑ t r(st, at)/c+ log p(at|st)]. (10)
The difference is that Eq. 9 includes an entropy regularizer, and thus a different off-policy method from Q-learning is required. A generalization of Ψ-learning (Rawlik et al., 2012), and G-learning (Fox et al.)2 are two off-policy methods for solving the KL-regularized RL problem, where additional generalized-Ψ and G functions are defined and learned instead of Q. We implement both of these algorithms as well, treating the prior policy as the conditional distribution p(a|s) defined by the trained Note RNN. To the best of our knowledge, this is the first application of KL-regularized off-policy methods with deep neural networks to sequence modeling tasks. The two methods are given below respectively,
L(θ) = Eβ [(log p(a|s) + rMT (s, a)/c+ γ log ∑ a′ eΨ(s ′,a′;θ−) −Ψ(s, a; θ))2] (11)
πθ(a|s) ∝ eΨ(s,a;θ) (12)
L(θ) = Eβ [(rMT /c(s, a) + γ log ∑ a′ elog p(a ′|s′)+G(s′,a′;θ−) −G(s, a; θ))2] (13)
πθ(a|s) ∝ p(a|s)eG(s,a;θ). (14) 2The methods in the original papers are derived for different motivations and presented in different forms as
described in Section 4, but we refer them using their names as the derivations follow closely from the papers.
Both methods can be seen as instances of KL-regularized deep Q-learning, and they also subsume entropy-regularized deep Q-learning by removing the log p(a|s) term. The main difference between the two methods is the definition of the action-value functions generalized-Ψ and G. In fact G-learning can be directly derived from generalized Ψ-learning by reparametrizing Ψ(s, a) = log p(a|s)+G(s, a). TheG-function does not give the policy directly but instead needs to be dynamically mixed with the prior policy probabilities. While this computation is straight-forward for discrete action domains as here, extensions to continuous action domains require additional considerations such as normalizability of advantage function parametrizations (Gu et al., 2016). The KL control-based derivation also has another benefit in that the stochastic policies can be directly used as an exploration strategy, instead of heuristics such as -greedy or additive noise (Mnih et al., 2013; Lillicrap et al., 2016). The derivations for both methods are included in the appendix for completeness.
3.2 MUSIC-THEORY BASED REWARD
A central question of this paper is whether RL can be used to constrain a sequence learner such that the sequences it generates adhere to a desired structure. To test this hypothesis, we developed several rules that we believe describe more pleasant-sounding melodies, taking inspiration from a text on melodic composition (Gauldin, 1995). We do not claim these characteristics are exhaustive, strictly necessary for good composition, or even particularly interesting. They simply serve the purpose of guiding the model towards traditional composition structure. It is therefore crucial to apply the RL Tuner framework to retain the knowledge learned from real songs in the training data.
Following the principles set out on page 42 of Gauldin’s book (Gauldin, 1995), we define the reward function rMT (a, s) to encourage melodies to have the following characteristics. All notes should belong to the same key, and the melody should begin and end with the tonic note of the key; e.g. if the key is C-major, this note would be middle C. This note should occur in the first beat and last 4 beats of the melody. Unless a rest is introduced or a note is held, a single tone should not be repeated more than four3 times in a row. To encourage variety, we penalize the model if the melody is highly correlated with itself at a lag of 1, 2, or 3 beats. The penalty is applied when the auto-correlation coefficient is greater than .15. The melody should avoid awkward intervals like augmented 7ths, or large jumps of more than an octave. Gauldin also indicates good compositions should move by a mixture of small steps and larger harmonic intervals, with emphasis on the former; the reward values for intervals reflect these requirements. When the melody moves with a large interval (a 5th or more) in one direction, it should eventually be resolved by a leap back or gradual movement in the opposite direction. Leaping twice in the same direction is negatively rewarded. The highest note of the melody should be unique, as should the lowest note. Finally, the model is rewarded for playing motifs, which are defined as a succession of notes representing a short musical “idea”; in our implementation, a bar of music with three or more unique notes. Since repetition has been shown to be key to emotional engagement with music (Livingstone et al., 2012), we also sought to train the model to repeat the same motif within a melody.
4 RELATED WORK
Generative modeling of music with RNNs has been explored in a variety of contexts, including generating Celtic folk music (Sturm et al., 2016), or performing Blues improvisation (Eck & Schmidhuber, 2002). Other approaches have examined RNNs with richer expressivity, latent-variables for notes, or raw audio synthesis (Boulanger-Lewandowski et al., 2012; Gu et al., 2015; Chung et al., 2015). Recently, impressive performance in generating music from raw audio has been attained with convolutional neural networks with receptive fields at various time scales (Dieleman et al., 2016).
Although the application of RL to RNNs is a relatively new area, recent work has attempted to combine the two approaches. MIXER (Mixed Incremental Cross-Entropy Reinforce) (Ranzato et al., 2015) uses BLEU score as a reward signal to gradually introduce a RL loss to a text translation model. After initially training the model using cross-entropy, the training process is repeated using cross-entropy loss for the T −∆ tokens in a sequence (where T is the length of the sequence), and
3While the number four can be considered a rough heuristic, avoiding excessively repeated notes and static melodic contours is Gauldin’s first rule of melodic composition (Gauldin, 1995).
using RL for the remainder of the sequence. Another approach (Bahdanau et al., 2016) applies an actor-critic method and uses BLEU score directly to train a critic network to output the value of each word, where the actor is again initialized with the policy of an RNN trained with next-step prediction. Reward-augmented maximum likelihood (Norouzi et al., 2016) augments the standard ML with a sequence-level reward function and connects it with the above RL training methods. These approaches assume that the complete task reward specification is available. They pre-train a good policy with supervised learning so that RL can be used to learn with the true task objective, since training with RL from scratch is difficult. RL Tuner instead only uses rewards to correct certain properties of the generated data, while learning most information from data. This is important since in many sequence modeling applications such as music or language generation, the true reward function is not available or imperfect and ultimately the model should rely on learning from data. The RL Tuner method provides an elegant and flexible framework for correcting undesirable behaviors of RNNs that can arise from limited training data or imperfect training algorithms.
SeqGAN (Yu et al., 2016) applies RL to an RNN by using a discriminator network — similar to those used in Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) — to classify the realism of a complete sequence, and this classifier-based reward is used as a reward signal to the RNN. The approach is applied to a number of generation problems, including music generation. Although the model obtained improved MSE and BLEU scores on the Nottingham music dataset, it is not clear how these scores map to the subjective quality of the samples (Huszár, 2015), and no samples are provided with the paper. In contrast, we provide both samples and quantitative results demonstrating that our approach improves the metrics defined by the reward function. Further, we show that RL Tuner can be used to explicitly correct undesirable behaviors of an RNN, which could be useful in a broad range of applications.
Also related to our work is that of Li and colleagues Li et al. (2016), in which the authors pre-train a model with MLE and then use RL to impose heuristic rules designed to improve the dialog generated by the model. However, after pre-training, only the heuristic rewards are used for further training, which alters the model to optimize only for the heuristic rewards, whereas our approach allows the model to retain information learned from data, while explicitly controlling the trade-off between the influence of data and heuristic reward with the c parameter. While Li and colleagues do use the outputs of the pre-trained model as part of one of the heuristic reward functions, it is only to teach the model to choose dialog turns that minimize the probability that the pre-trained model places on “dull” responses, such as “I don’t know”. However, our approach directly penalizes divergence from the probability distribution learned by the MLE model for every response, allowing the model to retain information about the full space of sequences originally learned from data.
Finally, as discussed in Section 3.1, our approach is related to stochastic optimal control (SOC) (Stengel, 1986) and KL control (Todorov, 2006; Kappen et al., 2012; Rawlik et al., 2012), in particular the two off-policy, model-free methods, Ψ-learning (Rawlik et al., 2012) and Glearning (Fox et al.). Both approaches solve a KL-regularized RL problem, in which a term is introduced to the reward objective to penalize KL divergence from some prior policy. While our methods rely on similar derivations presented in these papers, there are some key differences. First, these techniques have not been applied to DQNs or RNNs, or as a way to fine-tune a pre-trained RNN with additional desired charateristics. Secondly, our methods have different motivations and forms from the original papers: original Ψ-learning (Rawlik et al., 2012) restricts the prior policy to be the policy at the previous iteration and solves the original RL objective with conservative, KL-regularized policy updates, similar to conservative policy gradient methods (Kakade, 2001; Peters et al., 2010; Schulman et al., 2015). The original G-learning (Fox et al.) penalizes divergence from a simple uniform prior policy in order to cope with over-estimation of target Q values, and includes scheduling for the temperature parameter c. Lastly, our work includes the Q-learning objective with additional cross-entropy reward as a comparable alternative, and provides for the first time comparisons among the three methods for incorporating prior knowledge in RL.
5 EXPERIMENTS
To train the Note RNN, we extract monophonic melodies from a corpus of 30,000 MIDI songs. Melodies are quantized at the granularity of a sixteenth note, so each time step corresponds to one sixteenth of a bar of music. We encode a melody using two special events plus three octaves of notes.
The special events are used to introduce rests and notes with longer durations, and are encoded as 0 = note off, 1 = no event. Three octaves of pitches, starting from MIDI pitch 48, are then encoded as 2 = C3, 3 = C#3, 4 = D3, ..., 37 = B5. For example, the sequence {4, 1, 0, 1} encodes an eighth note with pitch D3, followed by an eighth note rest. As the melodies are monophonic, playing another note implicitly ends the last note that was played without requiring an explicit note off event. Thus the sequence {2, 4, 6, 7} encodes a melody of four sixteenth notes: C3, D3, E3, F3. A length-38 one-hot encoding of these values is used for both network input and network output.
The Note RNN consists of one LSTM layer of 100 cells, and was trained for 30,000 iterations with a batch size of 128. Optimization was performed with Adam (Kingma & Ba, 2014), and gradients were clipped to ensure the L2 norm was less than 5. The learning rate was initially set to .5, and a momentum of 0.85 was used to exponentially decay the learning rate every 1000 steps. To regularize the network, a penalty of β = 2.5 × 10−5 was applied to the L2 norm of the network weights. Finally, the losses for the first 8 notes of each sequence were not used to train the model, since it cannot reasonably be expected to accurately predict them with no context. The trained Note RNN eventually obtained a validation accuracy of 92% and a log perplexity score of .2536.
The learned weights of the Note RNN were used to initialize the three sub-networks in the RL Tuner model. Each RL Tuner model was trained for 1,000,000 iterations, using the Adam optimizer, a batch size of 32, and clipping gradients in the same way. The reward discount factor was γ=.5. The Target-Q-network’s weights θ− were gradually updated to be similar to those of the Q-network (θ) according to the formula (1 − η)θ− + ηθ, where η = .01 is the Target-Q-network update rate. We replicated our results for a number of settings for the weight placed on the music-theory rewards, c; we present results for c=.5 below because we believe them to be most musically pleasing. Similarly, we replicated the results using both -greedy and Boltzmann exploration, and present the results using -greedy exploration below.
We compare three methods for implementing RL Tuner: Q-learning, generalized Ψ-learning, and G-learning, where the policy defined by the trained Note RNN is used as the cross entropy reward in Q-learning and the prior policy in G- and generalized Ψ-learning. These approaches are compared to both the original performance of the Note RNN, and a model trained using only RL and no prior policy. Model evaluation is performed every 100,000 training epochs, by generating 100 melodies and assessing the average rMT and log p(a|s). All of the code for RL Tuner, including a checkpointed version of the trained Note RNN is available at https://github.com/natashamjaques/magenta/tree/rl-tuner.
6 RESULTS
Table 1 provides quantitative results in the form of performance on the music theory rules to which we trained the model to adhere; for example, we can assess the fraction of notes played by the model which belonged to the correct key, or the fraction of melodic leaps that were resolved. The statistics were computed by randomly generating 100,000 melodies from each model.
The results above demonstrate that the application of RL is able to correct almost all of the targeted “bad behaviors” of the Note RNN, while improving performance on the desired metrics. For example, the original LSTM model was extremely prone to repeating the same note; after applying RL, we see that the number of notes belonging to some excessively repeated segment has dropped from 63% to nearly 0% in all of the RL Tuner models. While the metrics for the G model did not improve as consistently, the Q and Ψ models successfully learned to play in key, resolve melodic leaps, and play motifs. The number of melodies that start with the tonic note has also increased, melody auto-correlation has decreased, and repeated motifs have increased slightly. The degree of improvement on these metrics is related to the magnitude of the reward given for the behavior. For example, a strong penalty of -100 was applied each time a note was excessively repeated, while a reward of only 3 was applied at the end of a melody for unique extrema notes (which most likely explains the lack of improvement on this metric). The reward values could be adjusted to improve the metrics further, however we found that these values produced the most pleasant melodies.
While the metrics indicate that the targeted behaviors of the RNN have improved, it is not clear whether the models have retained information about the training data. Figure 2a plots the average log p(a|s) as produced by the Reward RNN for melodies generated by the models every 100,000 training epochs; Figure 2b plots the average rMT . Included in the plots is an RL only model trained using only the music theory rewards, with no information about log p(a|s). Since each model is initialized with the weights of the trained Note RNN, we see that as the models quickly learn to adhere to the music theory constraints, log p(a|s) falls from its initial point. For the RL only model, log p(a|s) reaches an average of -3.65, which is equivalent to an average p(a|s) of approximately 0.026. Since there are 38 actions, this represents essentially a random policy with respect to the distribution defined by the Note RNN. Figure 2a shows that each of our models (Q, Ψ, and G) attain higher log p(a|s) values than this baseline, indicating they have maintained information about the data probabilities. The G-learning implementation scores highest on this metric, at the cost of slightly lower average rMT . This compromise between data probability and adherence to music theory could explain the difference in G model’s performance on the music theory metrics in Table 1. Finally, while c = 0.5 produced melodies that sounded better subjectively, we found that by increasing the c parameter it is possible to train all the models to have even higher average log p(a|s).
The question remains whether the RL-tuned models actually produce more pleasing melodies. To answer it, we conducted a user study via Amazon Mechanical Turk in which participants were asked to rate which of two randomly selected melodies they preferred on a Likert scale. A total of 192 ratings were collected; each model was involved in 92 of these comparisons. Figure 3 plots the number of comparisons in which a melody from each model was selected as the most musically pleasing. A Kruskal-Wallis H test of the ratings showed that there was a statistically significant difference between the models, χ2(3) = 109.480, p < 0.001. Mann-Whitney U post-hoc tests revealed that the melodies from all three RL Tuner models (Q, Ψ, and G) had significantly higher ratings than the melodies of the Note RNN, p < .001. The Q and Ψ melodies were also rated as significantly more pleasing than those of the G model, but did not differ significantly from each other. The sample melodies used for the study are available here: goo.gl/XIYt9m; we encourage readers to judge their quality for themselves.
Listening to the samples produced by the Note RNN reveals that they are sometimes dischordant and usually dull; the model tends to place rests frequently, repeat the same note, and produce melodies with little variation. In contrast, the melodies produced by the RL Tuner models are more varied and interesting. The G model tends to produce energetic and chaotic melodies, which include sequences of repeated notes. This repetition is likely because the G policy as defined in Eq. 14 directly mixes p(a|s) with the output of the G network, and the Note RNN strongly favours repeating notes. The most pleasant-sounding melodies are generated by the Q and Ψ models. These melodies stay firmly in key and frequently choose more harmonious in-
terval steps, leading to melodic and pleasant melodies. However, it is clear they have retained information about the training data; for example, the sample q2.wav in the sample directory ends with a seemingly familiar riff.
7 DISCUSSION AND FUTURE WORK
We have derived a novel sequence learning framework which uses RL rewards to correct properties of sequences generated by an RNN, while keeping much of the information learned from supervised training on data. We proposed and evaluated three alternative techniques for achieving this, and showed promising results on music generation tasks.
While we acknowledge that the simple monophonic melodies generated by these models — which are based on overly simplistic rules of melodic composition — do not approach the level of artistic merit of human composers, we believe this study provides a proof-of-concept that encoding domain knowledge using our method can help the outputs of an LSTM adhere to a more consistent structure. The musical complexity of the songs is limited not just by the heuristic rules, but also by the numerical encoding, which cannot represent the dynamics and expressivity of a musical performance. However, although these simple melodies cannot surpass those of human musicians, attempting to train a model to generate aesthetically pleasing outputs in the absence of a better metric of human taste than log-likelihood is a problem of broader interest to the artificial intelligence community.
In addition to the ability to train models to generate pleasant-sounding melodies, we believe our approach of using RL to refine RNN models could be promising for a number of applications. For example, it is well known that a common failure mode of RNNs is to repeatedly generate the same token. In text generation and automatic question answering, this can take the form of repeatedly generating the same response (e.g. “How are you?” → “How are you?” → “How are you?” ...). We have demonstrated that with our approach we can correct for this unwanted behavior, while still maintaining information that the model learned from data. Although manually writing a reward function may seem unappealing to those who believe in training models end-to-end based only on data, that approach it is limited by the quality of the data that can be collected. If the data contains hidden biases, this can lead to highly undesirable consequences. Recent research has shown that the word2vec embeddings in popular language models trained on standard corpora consistently contain the same harmful biases with respect to race and gender that are revealed by implicit association tests on humans (Caliskan-Islam et al., 2016). In contrast to relying solely on possibly biased data, our approach allows for encoding high-level domain knowledge into the RNN, providing a general, alternative tool for training sequence models.
ACKNOWLEDGMENTS
This work was supported by Google Brain, the MIT Media Lab Consortium, and Canada’s Natural Sciences and Engineering Research Council (NSERC). We thank Dzmitry Bahdanau, Greg Wayne, Sergey Levine, and Timothy Lillicrap for helpful discussions on RL and stochastic optimal control.
8 APPENDIX
8.1 OFF-POLICY METHODS DERIVATIONS FOR KL-REGULARIZED REINFORCEMENT LEARNING
Given the KL-regularized RL objective defined in Eq. 9, the value function is given by, V (st;π) = Eπ[ ∑ t′≥t r(st′ , at′)/c− KL[π(·|st′)||p(·|st′)]] (15)
8.1.1 GENERALIZED Ψ-LEARNING
The following derivation is based on modifications to (Rawlik et al., 2012) and resembles the derivation in Fox et al.. We define the generalized Ψ function as,
Ψ(st, at;π) = r(st, at)/c+ log p(at|st) (16) + Ep(st+1|st,at)Eπ[ ∑
t′≥t+1
r(st′ , at′)/c− KL[π(·|st′)||p(·|st′)]] (17)
= r(st, at)/c+ log p(at|st) + Ep(st+1|st,at)[V (st+1;π)] (18) The value function can be expressed as,
V (st;π) = Eπ[Ψ(st, at;π)] + H[π] (19) = Eπ[Ψ(st, at;π)− log π(at|st)] (20)
Fixing Ψ(st, at) = Ψ(st, at;π) and constraining π to be a probability distribution, the optimal greedy policy update π∗ can be derived by functional calculus, along with the corresponding optimal value function,
π∗(at|st) ∝ eΨ(st,at) (21)
V (st;π ∗) = log ∑ at eΨ(st,at) (22)
Given Eq. 18 and 22, the following Bellman optimality equation for generalized Ψ function is derived, and the Ψ-learning loss in Eq. 11 directly follows.
Ψ(st, at;π ∗) = r(st, at)/c+ log p(at|st) + Ep(st+1|st,at)[log ∑ at+1 eΨ(st+1,at+1;π ∗)] (23)
8.1.2 G-LEARNING
The following derivation is based on (Fox et al.) with small modifications. We define theG function as,
G(st, at;π) = r(st, at)/c+ Ep(st+1|st,at)Eπ[ ∑
t′≥t+1
r(st′ , at′)/c− KL[π(·|st′)||p(·|st′)]] (24)
= r(st, at)/c+ Ep(st+1|st,at)[V (st+1;π)] = Ψ(st, at;π)− log p(at|st) (25)
Similar derivation as above can be applied.
V (st;π) = Eπ[G(st, at;π)]− KL[π(·|st′)||p(·|st′)] (26)
= Eπ[G(st, at;π)− log π(at|st) log p(at|st) ] (27)
π∗(at|st) ∝ p(at|st)eG(st,at) (28)
V (st;π ∗) = log ∑ at p(at|st)eG(st,at) (29)
G(st, at;π ∗) = r(st, at)/c+ Ep(st+1|st,at)[log ∑ at+1 p(at+1|st+1)eG(st+1,at+1;π ∗)] (30)
Alternatively, the above expression for G-learning can be derived from Ψ-learning by simple reparametrization with Ψ(s, a) = G(s, a) + log p(a|s) in Eq. 23.

This paper uses a combination of likelihood and reward based learning to learn sequence models for music. The ability to combine likelihood and reward based learning has been long known, as a result of the unification of inference and learning first appearing in the ML literature with the EM formalism of Attias (2003) for fixed horizons, extended by Toussaint and Storkey (2006), to general horizon settings, Toussaint et al. (2011) to POMDPs and generalised further by Kappen et Al. (2012) and Rawlik et Al. (2012). These papers introduced the basic unification, and so any additional probabilistic or data driven objective can be combined with the reinforcement learning signal: it is all part of a unified reward/likelihood. Hence the optimal control target under unification is p(b=1|\tau)E_p(A,S) \prod_t \pi(a_t|s_t): i.e. the probability of getting reward, and probability of the policy actions under the known data-derived distribution, thereby introducing the log p(a_t|s_t) into (9) too.
The interpretation of the secondary objective as the prior is an alternative approach under a stochastic optimal control setting, but not the most natural one given the whole principle of SOC of matching control objectives to inference objectives. The SOC off policy objective still does still contain the KL term so the approach would still differ from the approach of this paper.
Though the discussion of optimal control is good, I think some further elaboration of the history and how reward augmentation can work in SOC would be valuable. This would allow SOC off-policy methods to be compared with the DQN directly, like for like.
The motivation of the objective (3) is sensible but could be made clearer via the unification argument above. Then the paper uses DCN to take a different approach from the variational SOC for achieving that objective.
Another interesting point of discussion is the choice of E_pi \log p(a_t|s_t) – this means the policy must “cover” the model. But one problem in generation is that a well-trained model is often underfit, resulting in actions that, over the course of a number of iterations, move the state into data-unsupported parts of the space. As a result the model is no longer confident and quickly tends to be fairly random. This approach (as opposed to a KL(p||pi) – which is not obvious how to implement) cannot mitigate against that, without a very strong signal (to overcome the tails of a distribution). In music, with a smaller discrete alphabet, this is likely to be less of a problem than for real valued policy densities, with exponentially decaying tails. Some further discussion of what you see in light of this issue would be valuable: the use of c to balance things seems critical, and it seems clear from Figure 2 that the reward signal needed to be very high to push the log p signal into the right range.
Altogether, in the music setting this paper provides a reasonable demonstration that augmentation of a sequence model with an additional reward constraint is valuable. It demonstrates that DQN is one way of learning that signal, but AFAICS it does not compare learning the same signal via other techniques. Instead for the comparator techniques it reverts to treating the p(a|s) as a “prior” term rather than a reward term, leaving a bit of a question as to whether DQN is particularly appropriate.
Another interesting question for the discussion is whether the music theory reward could be approximated by a differentiable model, mitigating the need for an RL approach at all.


Deep models have achieved great success on a variety of challenging tasks. However, the models that achieve great performance often have an enormous number of parameters, leading to correspondingly great demands on both computational and memory resources, especially for fully-connected layers. In this work, we propose a new “density-diversity penalty” regularizer that can be applied to fullyconnected layers of neural networks during training. We show that using this regularizer results in significantly fewer parameters (i.e., high sparsity), and also significantly fewer distinct values (i.e., low diversity), so that the trained weight matrices can be highly compressed without any appreciable loss in performance. The resulting trained models can hence reside on computational platforms (e.g., portables, Internet-of-Things devices) where it otherwise would be prohibitive.
Deep models have achieved great success on a variety of challenging tasks. However, the models that achieve great performance often have an enormous number of parameters, leading to correspondingly great demands on both computational and memory resources, especially for fully-connected layers. In this work, we propose a new “density-diversity penalty” regularizer that can be applied to fullyconnected layers of neural networks during training. We show that using this regularizer results in significantly fewer parameters (i.e., high sparsity), and also significantly fewer distinct values (i.e., low diversity), so that the trained weight matrices can be highly compressed without any appreciable loss in performance. The resulting trained models can hence reside on computational platforms (e.g., portables, Internet-of-Things devices) where it otherwise would be prohibitive.
1 INTRODUCTION
Deep neural networks have achieved great success on a variety of challenging data science tasks (Krizhevsky et al., 2012; Hinton et al., 2012; Simonyan & Zisserman, 2014b; Bahdanau et al., 2014; Mnih et al., 2015; Silver et al., 2016; Sutskever et al., 2014). However, the models that have achieved this success have a very large number of parameters, a consequence of their wide and deep architectures. Although such models yield great performance benefits, the corresponding memory and computational costs are high, making such models inaccessible to lightweight architectures (e.g., portable devices, Internet-of-Things devices, etc.). In such settings, the deployment of neural networks offers tremendous potential to produce novel applications, yet the modern top-performing networks are often infeasible on these platforms.
Fully-connected layers and convolutional layers are the two most commonly used neural network structures. While networks that consist of convolutional layers, are particularly good for vision tasks, the fully-connected layers, even if they are in the minority, are responsible for the majority of the parameters. For example, the VGG-16 network (Simonyan & Zisserman, 2014a) has 13 convolutional layers and 3 fully-connected layers, but the parameters for 13 convolutional layers contain only ∼ 1/9 of the parameters of the three fully-connected layers. Moreover, for general tasks, convolutional layers may be not applicable (data might be one-dimensional, or there might be no local correlation among data dimensions). Therefore, compression of fully-connected layers is critical for reducing the memory and computational cost of neural networks in general.
We can use the characteristics of convolutional layers to reduce the memory and computational cost of fully-connected layers. Convolution is a special form of matrix multiplication, where the weights of the matrix are shared according to the convolution structure (low diversity), and most entries of the weight matrix are zeros (high sparsity). Both of these properties greatly reduce the information capacity of the weight matrix.
In this paper, we propose a density-diversity penalty regularization, which encourages low diversity and high sparsity on fully-connected layers. The method uses a pairwise L1 loss of weight matrices to the training objective. Moreover, we propose a novel “sorting trick” to efficiently optimize the density-diversity penalty, which would otherwise would be very difficult to optimize and would slow down training significantly. When initializing the weights to have a small portion of them to be zero, the density-diversity penalty effectively increases the sparsity of the trained weight matrices in addition to reducing the diversity, generating highly compressible weight matrices.
On two separate tasks, computer vision and speech recognition, we demonstrate that the proposed density-diversity penalty significantly reduces the diversity and increases the sparsity of the models, while keeping the performance almost unchanged.
2 PREVIOUS RELATED WORK
For this paper, we focus on reducing the number of actual parameters in fully-connected layers by using the density-diversity penalty to encourage sparsity and penalize diversity. Several previous studies have focused on the task of reducing the complexity of deep neural network weights.
Nowlan & Hinton (1992) introduced a regularization method to enforce weight sharing by modeling the distribution of weight values as a Gaussian mixture model. This approach is in the same spirit as our method, because both methods apply regularization during training to reduce the complexity of weights. However, the Nowlan et al. approach focuses more on the generalization of the network rather than the compression. Moreover, their approach involves explicitly clustering weights to group them, while in our approach, weights are grouped by the regularizer directly.
The “optimal brain damage” (LeCun et al., 1989) and “optimal brain surgeon” (Hassibi & Stork, 1993) methods are two approaches for pruning weight connections of neural networks based on information about the second order derivatives. Interestingly, Cireşan et al. (2011) showed that dropping weights randomly can lead to better performance. All three approaches focus more on pruning unnecessary weights of the network (increasing sparsity) rather than grouping parameters (decreasing diversity), while our approach addresses both tasks with a single regularizer.
Chen et al. (2015b) proposed a hashing trick to randomly group connection weights into hash buckets, so that the weight entries in the same hash bucket are tied to be a single parameter value. Such approaches force low diversity among weight matrix entries, and the grouping of weight entries are determined by the hash functions prior to training. In contrast, our method learns to tie the parameters during training.
Han et al. (2015b) first proposed a method to learn both weights and connections for neural networks by iteratively pruning out low-valued entries in the weight matrices. Based on that idea, they proposed a three stage pipeline (Han et al., 2015a): pruning low-valued weights, quantizing weights through k-means clustering, and Huffman coding. The resulting networks are highly compressed due to both sparsity (pruning) and diversity (quantization). The quantization of weights is applied to well-trained models with pruned weights, whereas our method prunes and quantizes weights simultaneously during training. Our approach is most relevant to that of Han et al., and we achieve comparable or better results. In some sense, it may be argued that our approach generalizes on that of Han et al. because the hyperparameter controlling the strength of the density-diversity penalty can be adaptively increased during training (although in the present paper, we keep it fixed during training).
3 DENSITY-DIVERSITY PENALTY
Suppose we are given a deep neural network of the following form:
ŷ = Wmφ(Wm−1φ(. . . φ(W1x))), (1)
where Wj represents the weight matrix of layer j, φ is a non-linear transfer function, x denotes the input data, and y and ŷ denote the true and estimated labels for x, respectively. Suppose weight matrix Wj has order (rj , cj).
Let the objective of the deep neural network be minL(ŷ, y), where L(·) is the loss function. We propose the following optimization to encourage low density and low diversity in weight matrices:
minL(ŷ, y) + m∑ j=1 λj  ∑ a=1:rj , b=1:cj ∑ a′=1:rj ,
b′=1:cj
|Wj(a, b)−Wj(a′, b′)|+ ‖Wj‖p  , (2) where Wj(a, b) denotes the entry of Wj in the ath row and bth column, λj > 0 is a hyperparameter, and ‖·‖p is a matrix p-norm (e.g., p = 2 gives the Frobenius norm, or p = 1 is a sparsity encouraging norm). We denote the density-diversity penalty for each Wj as DP(Wj).
In general, for weight matrix Wj , the proposed density-diversity penalty resembles the pairwise L1 difference over all entries of Wj . Intuitively, such regularization forces the entries of Wj to collapse into the same values, not just similar values. The regularizer thus reduces the diversity of Wj significantly. The diversity penalty, therefore, is a form of total variation penalty (Rudin et al., 1992) but where the pattern on which total variation is measured is not between neighboring elements (as in computer vision (Chambolle & Lions, 1997)) but rather globally among all pairs of all elements in the matrix.
Though the hyperparameter λj can be tuned for each layer, in practice, we only tune one λj for layer j, and for all layers j′ 6= j, we set λj′ = ( rj′cj′
rjcj )λj , because the magnitude of the density-diversity
penalty is directly correlated with the number of entries in the weight matrix.
3.1 EFFICIENTLY OPTIMIZING THE DENSITY-DIVERSITY PENALTY
While the gradient of the p-norm part of the density-diversity penalty is easy to calculate, computing the gradient of the diversity part of the density-diversity penalty is expensive: naive evaluation costs O(r2j c 2 j ) for a weight matrix Wj of order (rj , cj). If we suppose rj = 2000 and cj = 2000, which is common for modern deep neural networks, then the computational cost of the density-diversity penalty becomes roughly (2000)4 = 1.6× 1013, which is intractable even on modern GPUs. For simplicity, suppose we assign the subgradient of each L1 term at zero to be zero, which is typical in certain widely-used neural network toolkits such as Theano (Theano Development Team, 2016) and mxnet (Chen et al., 2015a). With a sorting trick, shown in Alg 1, we can greatly reduce the computational cost of calculating the gradients of density-diversity penalty from O(r2j c 2 j ) down to only O(rjcj(log rj + log cj)).
input : Wj , λj , rj , cj output: ∂DP(Wj)∂Wj W̄j = flatten(Wj) ; Ij = sort index(W̄j) ; I ′j = num greater(Ij); ∂DP(Wj)
∂Wj = reshape(λj ∗ (I ′j − Ij), (rj , cj));
return ∂DP(Wj)∂Wj Algorithm 1: Sorting Trick for Efficiently Calculating the Gradient of DensityDiversity Penalty on Weight Matrix Wj :
∂DP(Wj) ∂Wj
In the algorithm, flatten(Wj) transforms the matrix Wj , which is of order (rj , cj), into a vector of length rj ∗ cj , and sort index(·) outputs the sorted indices (in ascending order) of the input vector, e.g. sort index((3, 2, 1, 2, 3, 4)) = (3, 1, 0, 1, 3, 5). In other words, entry i in the sorted indices is the number of entries in the original vector smaller than the ith entry. Also note that the entries with the same value have the same sorted index. Correspondingly, num greater(·) outputs the number of elements greater than a certain entry based on the sorted index. For example, num greater((3, 1, 0, 1, 3, 5)) = (1, 3, 5, 3, 1, 0). Finally, reshape(·) transforms the input vector into a matrix of the given shape.
The computational cost of the sorting trick is dominated by the sorting step Ij = sort index(W̄j), which is of complexity O(rjcj(log rj + log cj)). We show that the sorting trick outputs the correct
gradient in the following equation:
∂DP(Wj) ∂Wj(a, b) = ∂ ∑ a′=1:rj ,b′=1:cj |Wj(a, b)−Wj(a′, b′)| ∂Wj(a, b)
= ∑
Wj(a,b)>Wj(a′,b′)
∂|Wj(a, b)−Wj(a′, b′)| ∂Wj(a, b)
+ ∑
Wj(a,b)<Wj(a′,b′)
∂|Wj(a, b)−Wj(a′, b′)| ∂Wj(a, b)
= ∑
Wj(a,b)>Wj(a′,b′)
1 + ∑
Wj(a,b)<Wj(a′,b′)
−1
= I ′j(arj + b)− Ij(arj + b), (3)
where Ij(arj + b) corresponds to the (arj + b)th entry of Ij , which is the ath row and bth column of the matrix formed by reshaping Ij into (rj , cj).
Intuitively, ∂DP(Wj)∂Wj(a,b) requires counting the number of entries in Wj with values greater than Wj(a, b), and the number of entries less than Wj(a, b). By sorting entries of Wj , we get Ij(arj + b) and I ′j(arj + b) for all pairs of (a, b) collectively; therefore, we can easily calculate the gradient for the density-diversity penalty.
Although the sorting trick is efficient for calculating the gradient for the density-diversity penalty, depending on the size of each weight matrix, the computational cost can still be high. In practice, to further reduce computational cost, for every mini-batch, we only apply the density-diversity penalty with a certain small probability (e.g. 1% to 5%). This approach still effectively forces the values of weight matrices to collapse, while the increase in the training time is not significant. For our implementation, to accelerate collapsing of weight entries, we truncate the weight matrix entries to have a limited number of decimal digits (e.g. 6), in which case entries with very small differences (e.g. 1e− 6) are considered to be the same value.
3.2 ENCOURAGING SPARSITY
The density-diversity penalty forces entries of a weight matrix to collapse into the same value, yet sparsity is not explicitly enforced. To encourage sparsity in the weight matrices, we randomly initialize every weight matrix with 10% sparsity (i.e., 90% of weight matrix entries are non-zero). Thereafter, every time we apply density-diversity penalty, we subsequently set the weight matrix value corresponding to the modal value to be zero. Because the value zero is almost always the most frequent value in the weight matrix (owing to the p-norm), weights are encouraged to stay at zero when using this method, because the density-diversity penalty encourages weights to collapse into same values. Our sparse initialization approach thus complements any sparsity-encouraging property of the p-norm part of the density-diversity penalty.
3.3 COMPRESSION WITH LOW DIVERSITY AND HIGH SPARSITY
Low diversity and high sparsity both can significantly reduce the number of bits required to encode the weight matrices of the trained network. Specifically, for low diversity, considering the weight matrix with d distinct entries, we only need dlog2de bits to encode each entry, in contrast to 32-bit floating point values for the original, uncompressed model. High sparsity facilitates encoding the weight matrices in a standard, sparse matrix representation using value and position pairs. Therefore, for a weight matrix in which s entries are not equal to the modal value of the matrix, 2s+min(rj , cj) entries are required for encoding, where the min(rj , cj) part is for indicating the row or column index, depending on whether compressed sparse row or column form is used. We note that further compression is possible: e.g., by encoding sparsity with values and increments in positions instead of absolute positions, so that the increments are often small values which require less bits to encode. Huffman Coding can also further be applied in the final step, as is done in (Han et al., 2015a), but we do not report this method in our results.
3.4 TYING THE WEIGHTS
Because the density-diversity penalty collapses weight matrix entries into same values, we tie the weights together so that for every distinct value v of weight matrix Wj , the entries that are equal to v are updated with the average of their gradients.
In practice, we design the training procedure to alternate between learning with the density-diversity penalty and learning with tied weights. During the phase of learning with the density-diversity penalty, we apply the density-diversity penalty with untied weights. This approach greatly reduces the diversity of the weight matrix entries. However, the performance of the resulting model could be inferior to that of the original model because this approach does not optimize the loss function directly. Therefore, for the phase of learning with tied weights, we train the network without the density-diversity penalty, but we tie the entries of the weight matrices according to the pattern learned from the previous phase. In this way, the network’s performance improves while the diversity patterns of the weight matrices are unchanged. We also note that during the phase of learning with tied weights, the sparsity pattern is also fixed, because it was learned in the previous density-diversity pattern learning phase. We show the full procedure of model training with the density-diversity penalty in Figure 1.
An alternative approach would be to train the model with tied weights and use the density-diversity penalty simultaneously. However, because the diversity pattern, which controls the tying of the weights, changes rapidly across mini-batches, the weights would need to be re-tied frequently based on the latest diversity pattern. This approach would therefore be computationally expensive. Instead, we choose to alternate between applying the density-diversity penalty and training tied weights. In practice, we train each phase for 5 to 10 epochs.
We note that the three key components of our algorithm, namely the sparse initialization, the regularization, and the weight tying, contribute to highly compressed network only when they are applied jointly. Independently applying any of the key component would result in inferior results (we stated developing our approach without the sparse initialization and weight tying and got worse compression results) as a good compression requires low diversity, which is achieved by applying the densitydiversity penalty, high sparsity, which is achieved by applying both the density-diversity penalty and the sparse initialization, and no loss of performance, which is achieved by training with weight tying.
4 RESULTS
We apply the density-diversity penalty (with p = 2 for now) to the fully-connected layers of the models on both the MNIST (computer vision) and TIMIT (speech recognition) datasets, and get significantly sparser and less diverse layer weights. This approach yields dramatic compression of models, whose original sizes are already quite conservative, while keeping the performance unchanged. For our implementation, we start with the mxnet (Chen et al., 2015a) package, which we modified by changing the weight updating code to include our density-diversity penalty.
To evaluate the effectiveness of the density-diversity penalty, we report the diversity and sparsity of the trained weight matrices. We define “diversity” to be number of distinct values divided by the total number of entries in the weight matrix , “sparsity” to be number of entries with the modal value divided by the total number of entries, and “density” to be 1−sparsity. Based on the observed diver-
sity and sparsity, we can estimate the compression rate of the model. For weight matrixWj , suppose kvaluej = dlog2(diversity(Wj) ∗ rj ∗ cj)e, and kindexj = dlog2(min(rj , cj))e, which represent the bits required to encode the value and position, respectively, in the sparse matrix representation. Thus, we have
CompressRate(Wj) = rjcjp
(1−sparsity(Wj))rjcj(kvaluej +kindexj )+diversity(Wj)rjcjp+min(rj ,cj) ,
(4)
where p is the number of bits required to encode the weight matrix entries used in the original model, and we choose p = 32 as used in most modern neural networks.
4.1 MNIST DATASET
The MNIST dataset consists of hand-written digits, containing 60000 training data points and 10000 test data points. We further sequester 10000 data points from the training data to be used as the validation set for parameter tuning. Each data point is of size 28× 28 = 784 dimensions, and there are 10 classes of labels.
We choose LeNet (LeCun et al., 1998) as the model to compress, because LeNet performs well on MNIST while having a restricted size, which makes compression hard. Specifically, we test on LeNet-300-100, which consists of two hidden layers, with 300 and 100 hidden units respectively, as well as LeNet-5, which contains two convolutional layers and two fully connected layers. Note that for LeNet-5, we only apply the density-diversity penalty on the fully connected layers. For optimization, we use SGD with momentum.
We report the diversity and sparsity of each layer of the LeNet-300-100 model trained with the diversity penalty in Table 1. The overall compression rate for the LeNet-300-100 model is 32.43X, using 10 bits to encode both value and index of the sparse matrix representation (the number of bits are based on the number of distinct values in the trained weight matrices). The error rate of the compressed model is 1.62%, while the error rate of the original model is 1.64%; thus, we obtain a highly compressed model without loss of performance. Compared to the state-of-the-art “deep compression” result (Han et al., 2015b), which achieves roughly 32 times compression rate (without applying Huffman Coding in the end), our method overall achieves a better compression rate. We also note that “deep compression” uses a more complex sparsity matrix representation, so that indices of values are encoded with many fewer bits.
In Table 2, we show the per-layer diversity and sparsity of the LeNet-5 convolutional model applied with the density-diversity penalty. For such a model, the overall compression rate is 15.78X. When considering only the fully-connected layers of the model, where most parameters reside and the
density-diversity penalty applies, the compression rate is 226.32X, using 9 bits for both value and index for sparse matrix representation. The error rate for the compressed model is 0.93%, which is comparable to the 0.88% error rate of the original model. Compared to the “deep compression” method (without Huffman Coding), which gives 33 times compression for the entire model, and 36 times compression for the fully-connected layers only, our approach achieves better results on the fully-connected layers.
4.2 TIMIT DATASET
The TIMIT dataset is for a speech recognition task. The dataset consists of a 462 speaker training set, a 50 speaker validation set, and a 24 speaker test set. Fifteen frames are grouped together as inputs, where each frame contains 40 log mel filterbank coefficients plus energy, along with their first and second temporal derivatives (Mohamed et al., 2012). Overall, there are 1.1M training data samples, 120k validation samples, and 50k test samples. We use a window of size 15 ± 7 of each frame, so that each data point has 1845 dimensions. Each dimension is normalized by subtracting the mean and dividing by the standard deviation. The label vector has 183 dimensions, consisting of three states for each of the 61 phonemes. For decoding, we use a bigram language model (Mohamed et al., 2012), and the 61 phonemes are mapped to 39 classes as done in (Lee & Hon, 1989) and as is quite standard.
We choose the model used in (Mohamed et al., 2012) and (Ba & Caruana, 2014) as the target for compression. In particular, the model contains three hidden fully-connected layers, each of which has 2048 hidden units. We choose ReLU as the activation function and AdaGrad (Duchi et al., 2011) for optimization, which performs the best on the original models without the density-diversity penalty.
Table 3 shows the per-layer diversity and sparsity for both the density-diversity penalty and “deep compression” applied to the 3-2048 fully-connected model on TIMIT dataset. We train the original, uncompressed model and observe a 23.30% phone error rate on the core test set. With our best effort of tuning parameters for the “deep compression” method, we get 23.35% phone error rate and a compression rate of 19.47X, using 64 cluster centers for the k-means quantization step. For our density-diversity penalty regularization, we get 21.45X compression and 23.25% phone error rate with 11 digits for value encoding and 11 digits for position encoding for the sparse matrix representation.
We visualize a part of the weight matrix either trained with or without the density-diversity penalty in Figure 2. We clearly observe that the weight matrix trained with the density-diversity penalty has significantly more commonality amongst entry values. In addition, the histogram of the entry values comparing the two weight matrices (Figure 3) shows that the weight matrix trained with density-diversity penalty has much less variance in the entry values. Both figures show that densitydiversity penalty effectively makes the weight matrix extremely compressed.
5 DISCUSSION
On both the MNIST and TIMIT datasets, compared to the “deep compression” method (Han et al., 2015b), the density-diversity penalty achieves comparable or even better compression rates on fullyconnected layers. Another advantage offered by the density-diversity penalty approach is that, rather than learning the sparsity pattern and diversity pattern separately as done in the “deep compression” method, the density-diversity penalty enforces high sparsity and low diversity simultaneously, which greatly reduces the effort involved in tuning parameters.
Specifically, comparing the diversity and sparsity of the trained matrices using the two compression methods, we find that the density-diversity penalty achieves higher sparsity but more diversity than the “deep compression” method. The “deep compression” method has two separate phases, where for the first phase, the sparsity pattern is learned by pruning away low value entries, and for the second phase, k-means clustering is applied to quantize the matrix entries into a chosen number of clusters (e.g., 64), thus generating weight matrices with very low diversity. In contrast, the densitydiversity penalty acts as a regularizer, enforcing low diversity and high sparsity simultaneously and during training, so that the diversity and sparsity of the trained matrices are more balanced.
6 CONCLUSION
In this work, we introduce density-diversity penalty as a regularization on the fully-connected layers of deep neural networks to encourage high sparsity and low diversity pattern in the trained weight matrices. To efficiently optimize the density-diversity penalty, we propose a “sorting trick” to make the density-diversity penalty computationally feasible. On the MNIST and TIMIT datasets, networks trained with the density-diversity penalty achieve 20X to 200X compression rate on fullyconnected layers, while keeping the performance comparable to that of the original model.
In future work, we plan to apply the density-diversity penalty to recurrent models, extend the densitydiversity penalty to convolutional layers, and test other values of p. Moreover, besides pairwise L1 loss for the diversity portion of the density-diversity penalty, we will investigate other forms of regularizations to reduce the diversity of the trained weight matrices (e.g., other forms of structured convex norms). Throughout this work, we have focused on the compression task, but the learned sparsity/diversity pattern of the trained weight matrices is also worth exploring further. For image and speech data, we know that we can use the convolutional structure to improve performance, whereas for other very different forms of data, where we have no prior knowledge about the structure (i.e., patterns of locality), the density-diversity penalty may be applied to discover the underlying hidden pattern of the data and to achieve improved results.

This work introduces a number of techniques to compress fully-connected neural networks while maintaining similar performance, including a density-diversity penalty and associated training algorithm. The core technique of this paper is to explicitly penalize both the overall magnitude of the weights as well as diversity between weights. This approach results in sparse weight matrices comprised of relatively few unique values. Despite introducing a more efficient means of computing the gradient with respect to the diversity penalty, the authors still find it necessary to apply the penalty with some low probability (1-5%) per mini-batch.
The approach achieves impressive compression of fully connected layers with relatively little loss of accuracy. I wonder if the cost of having to sort weights (even for only 1 or 2 out of 100 mini-batches) might make this method intractable for larger networks. Perhaps the sparsity could help remove some of this cost?
I think the biggest fault this paper has is the number of different things going on in the approach that are not well explored independently. Sparse initialization, weight tying, probabilistic application of density-diversity penalty and setting the mode to 0, and alternating schedule between weight tied standard training and diversity penalty training. The authors don't provide enough discussion of the relative importance of these parts. Furthermore, the only quantitative metric shown is the compression rate which is a function of both sparsity and diversity such that they cannot be compared on their own. I would really like to see how each component of the algorithm affects diversity, sparsity, and overall compression.
A quick verification: Section 3.1 claims the density-diversity penalty is applied with a fixed probability per batch while 3.4 implies structured phases alternating between application of density-diversity and weight tied standard cross entropy. Is this scheme in 3.4 only applying the density-diversity penalty probabilistically when it is in the density-diversity phase?
Preliminary rating:
I think this is an interesting paper but lacks sufficient empirical evaluation of its many components. As a result, the algorithm has the appearance of a collection of tricks that in the end result in good performance without fully explaining why it is effective.
Minor notes:
Please resize equation 4 to fit within the margins (\resizebox{\columnwidth}{!}{ blah } works well in latex for this)