Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Block removal for large language models through constrained binary optimization

David Jansen, Roman Rausch, David Montero, Roman Orus

TL;DR

This work tackles the combinatorial problem of transformer block pruning by casting block removal as a constrained binary optimization (CBO) mapped to an Ising model with fixed magnetization. It derives a second-order loss expansion, producing a Hessian-based objective min $\boldsymbol{x}^\top H^0 \boldsymbol{x}$ subject to $\sum_i x_i = M$, with $H^0$ approximated by $(1/m) A^\top A$ and computable once from a calibration set. The resulting low-energy configurations yield pruned models that outperform state-of-the-art block-removal methods on dense and mixture-of-experts architectures, including notable gains in MMLU at high compression and effective pruning without retraining in some MoE cases. The approach generalizes to heterogeneous architectures and offers a robust, architecture-agnostic framework for efficient LLM compression, with potential extensions such as diversity-based candidate selection and integration with other pruning strategies.

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.

Block removal for large language models through constrained binary optimization

TL;DR

This work tackles the combinatorial problem of transformer block pruning by casting block removal as a constrained binary optimization (CBO) mapped to an Ising model with fixed magnetization. It derives a second-order loss expansion, producing a Hessian-based objective min subject to , with approximated by and computable once from a calibration set. The resulting low-energy configurations yield pruned models that outperform state-of-the-art block-removal methods on dense and mixture-of-experts architectures, including notable gains in MMLU at high compression and effective pruning without retraining in some MoE cases. The approach generalizes to heterogeneous architectures and offers a robust, architecture-agnostic framework for efficient LLM compression, with potential extensions such as diversity-based candidate selection and integration with other pruning strategies.

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.
Paper Structure (9 sections, 6 equations, 5 figures, 2 tables)

This paper contains 9 sections, 6 equations, 5 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Method
  4. Experiments
  5. Low-lying excited states
  6. Comparison against SOTA approaches
  7. Generalization to Heterogeneous and Mixture-of-Experts Architectures
  8. Conclusion
  9. Appendix

Figures (5)

  • Figure 1: Sketch of our method. We formulate block pruning as a constrained binary optimization problem that can be mapped to an Ising model with fixed magnetization. Each feasible solution specifies which $M$ out of $N$ blocks are removed from the LLM. The optimization is constructed such that low-energy solutions correspond to pruned models with high performance across multiple benchmarks.
  • Figure 2: Sketch of the implementation of our algorithm. For each $i$ of the $N$ transformer blocks in the LLM, we introduce an additional parameter $\alpha_i$ that multiplicatively modulates the output of the attention and feedforward blocks.
  • Figure 3: Number of times each block is recommended for removal among the first 20 low-energy states for the Llama-3.1-8B-Instruct model.
  • Figure 4: Selected benchmark results for the Llama-3.1-8B-Instruct model with 16 blocks removed. The ground state removes the blocks [14,16--24,26--31] and 17th excited state the blocks [2,16,18--31].
  • Figure 5: Benchmark results for NVIDIA-Nemotron-3-Nano-30B-A3B-FP8. We use the following notation: Gs:2 means the ground state (lowest-energy eigenvector) of the CBO problem when removing two layers, while, e.g., 19th excited:3 means the 19th excited state when removing 3 layers. See Tab. \ref{['tabel2']} for the corresponding block index combinations.