MLP Fusion: Towards Efficient Fine-tuning of Dense and Mixture-of-Experts Language Models

TL;DR

This work tackles the costly fine-tuning of large pre-trained language models by introducing MLP Fusion, a one-shot compression strategy that preserves training dynamics via the Neural Tangent Kernel (NTK). By clustering sub-MLPs within dense and MoE FFN modules into centroid-based representations, the method builds a compact MLP whose forward and backward passes closely approximate the original NTK, enabling similar training behavior with fewer parameters. The approach is complemented by a layer-wise, task-specific tuning step and applies to both dense and MoE architectures, with extensive experiments showing improved NTK preservation and competitive or superior performance on NLU and NLG benchmarks compared to strong baselines. The results suggest that NTK-aware clustering-based fusion can unlock efficient fine-tuning on edge devices and provide a robust initialization for subsequent distillation-based improvements, with practical implications for scalable language model deployment.

Abstract

Fine-tuning a pre-trained language model (PLM) emerges as the predominant strategy in many natural language processing applications. However, this process is known to be expensive, especially on edge devices with low computing power. While general approaches (e.g. quantization and distillation) have been widely studied to reduce the compute/memory of PLM fine-tuning, one-shot compression techniques specifically designed for fine-tuning remain largely unexplored. In this paper, we investigate the neural tangent kernel (NTK)--which reveals the gradient descent dynamics of neural networks--of the multilayer perceptrons (MLP) modules in a PLM and propose to coin a lightweight PLM through NTK-approximating MLP fusion. By incorporating NTK into the compression process, MLP Fusion not only preserves the original model's output but also maintains its training dynamics. To achieve this, we reconsider the MLP as a bundle of sub-MLPs and cluster them into a given number of centroids, which can then be restored as a compressed MLP and surprisingly well approximate the NTK of the original PLM. Our approach is applicable to both standard MLP modules and Mixture-of-Experts (MoE) modules in PLMs, demonstrating its scalability and versatility. Additionally, we provide theoretical derivations to demonstrate how the proposed compression preserves the NTK. Extensive experiments of PLM fine-tuning on both natural language understanding and generation tasks are provided to verify the effectiveness of MLP fusion. Our code is available at https://github.com/weitianxin/MLP_Fusion.

Figures (5)

• Figure 1: NTK matrix approximation error of compression methods on the validation set of SST2.
• Figure 2: In this illustrative example of MoE layers, the Top- K Selector, along with the Gate Network--often referred to as the 'router'--selects Experts 1 and 3 based on their scores for the given input.
• Figure 3: An overview of MLP Fusion. The original MLP in Dense and MoE LLMs can be decomposed as an ensemble of $p_I$ sub-MLPs; through MLP fusion, we cluster the sub-MLPs and re-construct a smaller MLP, which is shown to approximate the NTK of the original MLP and is thus expected to enjoy a similar training dynamics to the full-size PLM.
• Figure 4: Accuracy of each baseline method with respect to various numbers of sketched layers on the SST2 data set.
• Figure 5: Training loss (dynamic) of our proposed MLP Fusion with respect to various intermediate dimensions on the SST2 data set.