KroneckerBERT: Learning Kronecker Decomposition for Pre-trained Language Models via Knowledge Distillation

1 Introduction

• Transformer-based models' size conflicts with edge-device memory, latency, and energy constraints, driving interest in compression methods.The cited approaches include low-bit quantization, pruning, matrix decomposition, and knowledge distillation.Source: §1 · p. 1
• Extreme compression above a factor of 10 remained challenging, with prior attempts beyond that level trading compression for significant performance drops.Large compression factors also reduce the Kronecker network's expressiveness, motivating intermediate-layer distillation from the uncompressed model.Source: §1 · pp. 1–2
• At 19× compression, KroneckerBERT significantly outperforms state-of-the-art compressed BERT models.The framework is applied to BERT_BASE, and the paper notes its evaluation is limited to BERT while the framework can be used for other Transformer-based NLP models.Source: §1 · p. 2
• The method applies Kronecker decomposition to embedding and Transformer-layer weights, with knowledge distillation proposed to improve the compressed model.Transformer matrices in multi-head attention and feed-forward networks are represented as products of smaller matrices.Source: §1 · p. 2

2 Related Work

• The review covers language-model compression and earlier Kronecker use for CNNs and RNNs.It frames the section around BERT compression work and Kronecker decomposition in other neural-network settings.Source: §2 · p. 2

2.1 Pre-trained Language Model Compression

• Prior PLM compression combines knowledge distillation with pruning or factorization, while other methods use layer-wise distillation and smaller vocabularies.BERT-PKD distills intermediate layers during fine-tuning; TinyBERT uses pre-training and fine-tuning distillation, and MobileBERT uses inverted bottlenecks.Source: §2.1 · pp. 2–3

2.2 Kronecker Factorization

• Earlier Kronecker compression targeted CNNs and small RNNs; this work presents the first Transformer-language-model application, using two-factor products for linear layers with distillation.Prior work also combined Kronecker factorization with an unfactorized upper matrix or a learned sparse matrix.Source: §2.2 · p. 3

3.1 Kronecker Product

• The proposed embedding compression represents the embedding matrix as a Kronecker product of a matrix and a row vector, storing the matrix in a lookup table to reduce computation overhead.The figure contrasts this representation with a conventional embedding lookup table.Source: Figure 2 · p. 3
• For A sized m1×n1 and B sized m2×n2, their Kronecker product has dimensions m1m2×n1n2.The section refers readers to Graham (2018) for further details.Source: §3.1 · p. 3
• The Kronecker product forms a block matrix whose (i,j) block is the corresponding A element multiplied by matrix B.Figure 1 illustrates the operation using two 2×2 matrices.Source: §3.1 · p. 3

3.2 Kronecker factorization

• A matrix W is represented as A ⊗ B, with factor dimensions chosen from the dimensions of W to set the compression factor.The dimensions of A can be any factors of W’s input and output dimensions; B’s dimensions follow from those choices.Source: §3.2 · p. 3

3.2.1 Memory and computation reduction

• Storage for a weight matrix falls from mn elements to m1n1 + m2n2 when represented by Kronecker factors.Source: §3.2.1 · p. 3
• The implementation computes (A ⊗ B)x through a matrix operation rather than reconstructing W, reducing the required FLOPs.The paper uses this method to implement its Kronecker layers.Source: §3.2.1 · pp. 3–4

• The embedding table X is factorized into AE and BE, with AE storing a distinct row for each word and BE defined as a row vector.Source: §3.3 · p. 4
• An individual word embedding can be obtained from its AE row and BE in O(d) computation.AE is stored as a lookup table.Source: §3.3 · p. 4

• The Transformer’s multi-head attention and feed-forward network are both compressed using Kronecker decomposition.Source: §3.4 · p. 4
• The per-head query, key, and value mappings are stacked, and Kronecker decomposition is applied after concatenation rather than separately to each head.The input is projected by WQ, WK, and WV to obtain Q, K, and V.Source: §3.4 · p. 4
• A shared Kronecker factor across attention heads provides additional compression; the output mapping and both FFN weight matrices are also compressed.Source: §3.4 · p. 4

3.5.1 Intermediate KD

• Intermediate distillation uses BERTBASE as teacher and its Kronecker-compressed counterpart as student, whose layer dimensions remain intact for direct output comparison.Source: §3.5.1 · pp. 4–5
• The intermediate loss compares embedding outputs, attention matrices, and FFN outputs between teacher and student networks.Source: §3.5.1 · p. 5
• A learnable projection of pooled final-layer attention and FFN features adds a projection loss, comparing student features with projected teacher features by MSE.The projection is used to obtain a richer representation of the knowledge, rather than to match feature dimensions.Source: §3.5.1 · pp. 5–6

3.5.2 KD at pre-training

• Pre-training distillation transfers general-domain knowledge from BERTBASE using intermediate-layer losses on 5% of English Wikipedia for three epochs.The authors state that this stage improves initialization for task-specific distillation and makes training more efficient.Source: §3.5.2 · p. 6

3.6 Model Settings

• All MHA matrices use the same factor configuration, while the two FFN layers use swapped dimensions.The search sets one factor matrix’s shape; the other is then determined accordingly.Source: §3.6 · p. 6
• For the row-vector embedding matrix, only the n dimension must be selected.Source: §3.6 · p. 6
• Factor shapes are chosen for the desired compression factor, with the configuration minimizing FLOPS among eligible choices.The eligible choices are determined by factors of the original weight matrix dimensions.Source: §3.6 · p. 6

3.7 Implementation details

• The pretrained teacher initializes uncompressed student layers by copying their values, while L2-norm approximation supplies factors for decomposed layers.The factors approximate Kronecker components of the pretrained BERT_BASE weights.Source: §3.7 · p. 6
• Pretraining distillation uses 5% of English Wikipedia for 3 epochs, with batch size 64 and learning rate e-3.Source: §3.7 · p. 6
• After pretraining, the Kronecker model initializes the student’s Kronecker layers, and the fine-tuned teacher initializes the prediction layer.Fine-tuning hyperparameters are selected using dev-set performance for each task.Source: §3.7 · p. 6

4 Experiments

• The experiments compare compression baselines, benchmark performance, OOD robustness, distillation ablations, and deployment speed.The speed study includes two smartphones and a high-performance CPU.Source: §4 · p. 6, §4.6 · p. 9

4.1 Baselines

• Baselines are selected from compression factors below and above 10, with BERT_BASE as teacher; TinyBERT’s comparison involves more pretraining data and augmented fine-tuning data.The table caption states that TinyBERT uses all of Wikipedia for pretraining KD and augmented data for fine-tuning, unlike KroneckerBERT’s 5% Wikipedia pretraining and no fine-tuning augmentation.Source: §4.1 · pp. 6–7, Table 2 · p. 7

4.2 Results on GLUE benchmark

• At compression factors above 10, KroneckerBERT outperforms all baselines on available GLUE results; excluding CoLA, KroneckerBERT8 averages 81.7 versus TinyBERT’s 81.2.At lower compression, it beats other baselines on every task except TinyBERT, with similar TinyBERT results except on CoLA.Source: Table 2 · p. 7, §4.2 · p. 7

4.3 Out of distribution robustness

• KroneckerBERT8 beats TinyBERT on all three OOD experiments; its MRPC and RTE models slightly exceed BERT_BASE, while IMDb performance drops significantly after compression.The evaluations pair MRPC with PAWS, RTE with HANS, and SST-2 with IMDb.Source: §4.3 · pp. 7–8, Table 3 · p. 7

4.4 Results on SQuAD

• KroneckerBERT8 significantly outperforms other baselines on both SQuAD1.1 and SQuAD2.0 dev sets.Results for baselines with higher compression factors were unavailable.Source: Table 4 · p. 7, §4.4 · p. 8

4.5 Ablation study on the effect of KD

• Two-stage KD yields the highest performance on all three tested tasks; without KD, Kronecker decomposition causes a significant performance drop.Fine-tuning KD is more effective for MNLI-m, whereas pretraining KD is more effective for SST-2 and MRPC; the figure also shows more separable middle-layer features with KD on SST-2.Source: §4.5 · p. 8, Table 5 · p. 8, Figure 4 · p. 8

4.6 Inference speedup: Discussion and results

• At compression factor 19, KroneckerBERT reaches up to 2.4× speedup on edge devices, but does not speed up inference on the high-performance CPU.Serial matrix multiplication can limit parallel processing, so fewer FLOPS need not yield lower latency without underlying software or hardware modification.Source: §4.6 · pp. 8–9, Table 6 · p. 8

• At compression factors of at least 10×, the framework reduces model size and computations while outperforming state-of-the-art methods.Source: §5 · p. 9
• The framework applies Kronecker decomposition to the embedding layer and linear mappings within Transformer blocks.Source: §5 · p. 9
• The compressed BERT_BASE model is trained using an efficient two-stage knowledge-distillation method.Source: §5 · p. 9
• The methodology can be applied to other Transformer-based language models.Source: §5 · p. 9