DotResize: Reducing LLM Width via Discrete Optimal Transport-based Neuron Merging
Abstract
Structured pruning methods designed for Large Language Models (LLMs) generally focus on identifying and removing the least important components to reduce model size. In this work, we question the need for total removal, instead exploring how to integrate information from less important structures back into the compressed model. We specifically focus on neuron width pruning by framing it as a discrete optimal transport problem, leading to our proposed method DotResize, a novel Transformer compression method that uses optimal transport theory to transform and compress model width. To ensure applicability within the Transformer architecture, we motivate and incorporate necessary entropic regularization and matrix factorization techniques into the transportation maps produced by our method. Unlike pruning-based approaches which discard neurons based on importance measures, DotResize re-projects the entire neuron width, allowing the retention and redistribution of useful signal across the reduced layer. Empirical results show that DotResize serves as a useful add-on to pruning techniques, while achieving measurable reductions in real-world computational cost.