DOTResize: Reducing LLM Width via Discrete Optimal Transport-based Neuron Merging

TL;DR

DOTResize introduces a novel method for compressing large language models by using optimal transport to reproject and merge neurons, preserving useful information and reducing computational costs.

Contribution

It proposes a new neuron width reduction technique based on discrete optimal transport, enhancing model compression beyond traditional pruning methods.

Findings

Achieves measurable reductions in computational cost.

Serves as an effective add-on to existing pruning techniques.

Maintains model performance while reducing size.

Abstract

Structured pruning methods designed for Large Language Models (LLMs) generally focus on identifying and removing the least important components to optimize model size. However, in this work, we question this prevalent approach by instead exploring how to recombine information from structures designated for pruning back into the reduced model. We specifically focus on neuron width reduction, and frame this problem as a Discrete Optimal Transport problem, and propose 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 compared to simple or state-of-the-art neuron width-pruning techniques, DOTResize serves as a useful add-on to pruning, while achieving measurable reductions in real-world computational cost.