# Fast Jukebox: Accelerating Music Generation with Knowledge Distillation
## 4 Methodology  
At each level of encoding, an autoregressive model was trained to learn the prior over the compressed spaces. This process resulted in the following:  
\[p(z)=p\Big{(}z^{top},z^{middle},z^{bottom}\Big{)} \tag{1}\]  
\[p(z)=p\big{(}z^{top}\big{)}p(z^{middle}\Big{|}z^{top}),p(z^{bottom}\Big{|}z^{ top},z^{middle}) \tag{2}\]  
where \(p\big{(}z^{top}\big{)}\) is the top-level prior, while \(p\Big{(}z^{middle}\Big{|}z^{top}\Big{)}\) and \(p(z^{bottom}\Big{|}z^{top},\ z^{middle})\) are the up samplers. Each model level was trained separately. Sparse transformers were proposed to solve this autoregressive modeling problem. The sparsity patterns are shown in Figure 2. Due to the computational expensive nature, knowledge distillation was used to accelerate the sampling time.  
Figure 1: Hierarchical VQ-VAE of three levels of abstraction. Each level is a model with separate priors conditioned on upper levels. \(h_{t}\) is the latent vector obtained from the encoder, and \(e_{z_{t}}\) is its nearest vector from codebook \(\mathrm{C}=\langle e_{k}\rangle_{k=1}^{K}\).  
Knowledge DistillationKnowledge distillation [2] is a method proposed to reduce size and computation costs during the inference stage in different models, such as transformers. The idea was to produce a large computational model to transfer its information to a smaller one. Once the large teacher model had been trained, its information would be used to train the student network so it could mimic its behavior. The most intuitive method of carrying out knowledge distillation was proposed by [PERSON] et al. [32]; here, the loss function was defined with cross entropy between the teacher and the student at predicted logits of \(z^{T}\) and \(z^{s}\), respectively.  
\[L_{ce}=\sum_{i}z_{i}^{T}log(z_{i}^{s}) \tag{3}\]  
In addition, we applied attention-based distillation and hidden-state-based distillation, as in [13], to strength transfer the information between the teacher and student using learning intermediate structures:  
\[L_{hith}=MSE\Big{(}H^{T},H^{S}\Big{)} \tag{4}\]  
where \(H^{T}\in R^{l\times d}\) and \(H^{S}\in R^{l\times d}\) represent the values shown in Figure 3, which are matrices with a sequence length 1 and transformer dimension d. The MSE evaluation was carried out with attention matrices \(A^{S}\) and \(A^{T}\), also illustrated in Figure 3. The Jukebox used a sparse-attention transformer with a different attention pattern for each layer. To keep the same level of attention for the student, the correspondence between layers was carried out as shown in Figure 4. Each layer's attention pattern was grouped together and distilled the information to the group in the student transformer. The layer proportion depended on the size of both the student and teacher transformers.  
Figure 2: Attention patterns used in the sparse transformer. The first row shows the full attention of the vanilla transformer, since it is autoregressive all squares are marked up until the last iteration input (dark square). The second row shows the attention patterns of row attention, column attention, and previous row attention from left to right, respectively.  
Figure 3: Transformer layer distillation for the prior Jukebox model architecture. Which is the typical autoregressive Transformer structure. Teacher and student outputs are compared in the attention block residual stage and outputs of the layers.