Block removal for large language models through constrained binary optimization

TL;DR

This paper introduces a novel method for large language model compression by formulating block removal as a constrained binary optimization problem, leveraging an Ising model analogy to efficiently identify high-quality block removal configurations.

Contribution

It presents a new optimization-based approach for block removal in large language models, outperforming existing methods and applicable to various architectures.

Findings

Outperforms state-of-the-art block-removal methods on multiple benchmarks.

Achieves up to 6-point improvements on the MMLU benchmark.

Requires only a few forward/backward passes and an Ising solver.

Abstract

Compressing resource-intensive large language models by removing whole transformer blocks is a seemingly simple idea, but identifying which blocks to remove constitutes an exponentially difficult combinatorial problem. In this paper, we formulate block removal as a constrained binary optimization problem that can be mapped to a physical system (Ising model), whose energies are a strong proxy for downstream model performance. This formulation enables an efficient ranking of a large number of candidate block-removal configurations and yields many high-quality, non-trivial solutions beyond consecutive regions. We demonstrate that our approach outperforms state-of-the-art block-removal methods across several benchmarks, with performance gains persisting after short retraining, and reaching improvements of up to 6 points on the MMLU benchmark. Our method requires only forward and backward passes for a few active parameters, together with an (at least approximate) Ising solver, and can be readily applied to any architecture. We illustrate this generality on the recent NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 model, which exhibits a highly inhomogeneous and challenging block structure.