🤖 machine learning

Structured Neuron Pruning in Deep Neural Networks Using Multi-Armed Bandits

This paper proposes a structured neuron pruning framework that utilizes multi-armed bandit algorithms to identify and remove redundant neurons, demonstrating through extensive benchmarks that policies like UCB1 and Thompson Sampling effectively reduce model size while maintaining or improving performance compared to unpruned networks and other pruning methods.

Original authors: Salem Ameen, Sunil Vadera

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Salem Ameen, Sunil Vadera

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, over-stocked warehouse (a Deep Neural Network) filled with thousands of workers (neurons). These workers are great at solving problems, but the warehouse is so big that it costs a fortune to run, takes up too much space, and is slow to navigate. You want to shrink the warehouse to make it faster and cheaper, but you're terrified of firing the wrong people. If you fire the wrong worker, the whole operation might collapse.

This paper presents a clever, low-risk strategy to figure out exactly which workers can be let go without hurting the business. It calls this method Structured Neuron Pruning using Multi-Armed Bandits.

Here is how it works, broken down into simple concepts:

1. The Problem: Why not just guess?

Usually, when people try to shrink these networks, they use one of two bad approaches:

The "Guessing Game" (Magnitude Pruning): They look at who is working the hardest (or the least) based on a simple score, like how much weight they carry. They fire the "lightest" workers. The problem? A worker might carry a light load but be the only one who knows a secret code. Firing them breaks the system.
The "Brute Force" Test: They fire one worker, test the whole warehouse, fire another, test again, and so on. This is accurate, but it takes forever. If you have 1,000 workers, you'd have to run the warehouse 1,000 times just to decide who to fire. That's too expensive.

2. The Solution: The "Taste-Test" Strategy

The authors propose a method that acts like a smart manager running a series of quick "taste tests."

The Concept: Imagine you have a buffet with 50 different dishes (the neurons). You want to remove the ones that nobody really needs, but you don't want to ruin the meal.
The "Bandit" Game: In the world of gambling, a "Multi-Armed Bandit" is a row of slot machines. You don't know which machine pays out the most, so you have to pull levers (arms) to find out. You have a limited number of coins (a budget) to spend. You need to balance exploring (trying new machines to see if they are good) and exploiting (playing the machine that seems to be paying out the most).
Applying it to the Network:
1. The algorithm picks a "worker" (neuron) to test.
2. It temporarily tells that worker to take a break (masks the neuron).
3. It runs a quick test (a small batch of data) to see if the warehouse still runs smoothly.
4. The Reward: If the warehouse runs just as well (or even better!) without that worker, the worker gets a "safe to fire" point. If the warehouse slows down, the worker gets a "keep" point.
5. The worker goes back to work immediately. Nothing is permanently fired yet.

3. The Smart Manager (The Algorithms)

The paper tests different "managers" (algorithms) to decide which worker to test next. They are all trying to solve the same puzzle with a limited number of coins:

Epsilon-Greedy: Mostly picks the workers that look safe to fire, but occasionally picks a random worker just in case they missed something.
UCB1 (Upper Confidence Bound): This is the "optimist." It picks workers that have done well and workers that haven't been tested enough yet. It says, "I'm not sure about this guy, so I should test him to be sure." This turned out to be the most effective manager in the study.
Thompson Sampling: This is the "statistician." It keeps a mental file of probabilities. It thinks, "There's a 90% chance this worker is useless, but a 10% chance they are a genius. Let's test them to update our file."
Hedge & EXP3: These are more complex strategies designed for tricky situations, but in this specific test, the simpler "optimist" (UCB1) and "statistician" (Thompson Sampling) worked best.

4. The Results: What Happened?

The researchers tested this method on three types of "warehouses":

Tabular Data: Standard spreadsheets (like predicting house prices or diagnosing diseases).
Regression: Predicting continuous numbers.
Deep Learning: Complex image and text recognition systems (like recognizing faces or understanding sentences).

The Findings:

It Works: The method successfully removed huge chunks of the network (sometimes up to 60-70% of the workers in a specific layer) without breaking the system.
It Often Improves Performance: Surprisingly, in many cases, firing the "redundant" workers actually made the network faster and more accurate. It's like cleaning out a cluttered desk; sometimes, having fewer distractions helps you work better.
The Winners: The UCB1 and Thompson Sampling managers consistently outperformed the old "guessing" methods and even beat the original, unpruned networks in many tests.

5. Why This Matters

This approach is special because it doesn't just look at how "heavy" a worker is (static stats); it actually tests if the worker is needed by seeing what happens when they step aside. It does this efficiently, using a smart sampling strategy so it doesn't waste time testing everyone exhaustively.

In short: The paper shows that by using a smart, adaptive "taste-test" system, we can shrink massive, expensive AI models into smaller, faster, and sometimes even smarter versions, making them ready to run on everyday devices like phones and tablets.

Technical Summary: Structured Neuron Pruning in Deep Neural Networks Using Multi-Armed Bandits

Problem Statement
Deep neural networks (DNNs) often contain redundant hidden units, leading to models that require excessive memory, computation, and energy, particularly hindering deployment on embedded and mobile systems. While weight-level pruning can reduce parameter counts, it often results in unstructured sparsity (irregular sparse matrices) that is difficult to exploit for speedups on standard hardware without specialized libraries. Structured pruning, which removes entire neurons (and their associated incoming and outgoing connections), offers a more practical solution by reducing dense matrix dimensions. However, the central challenge remains identifying which neurons can be removed without degrading predictive performance. Direct evaluation of every neuron via exhaustive testing is computationally prohibitive, while simple local statistics (e.g., weight magnitude or activation variance) are cheap but often fail to accurately predict the impact on loss.

Methodology
The authors propose a framework that formulates structured neuron pruning as a fixed-budget Multi-Armed Bandit (MAB) problem.

Formulation: Each candidate neuron in a selected layer is treated as an "arm." Pulling an arm corresponds to temporarily masking (removing) that specific neuron, measuring the change in loss on a sampled mini-batch, and then restoring the neuron.
Reward Signal: The core metric is the change in mini-batch loss ( $\Delta$ ) caused by temporary removal. A positive $\Delta$ indicates the removal improved performance (or reduced loss), while a negative $\Delta$ indicates harm. This loss change is converted into a bounded reward (for most policies) or a binary success/failure indicator (for Thompson Sampling).
Algorithms: The framework supports a unified algorithm that adapts various MAB policies to select which neurons to evaluate next. The evaluated policies include:
- Stochastic/Exploration-based: Epsilon-Greedy, Softmax (Boltzmann) exploration.
- Optimistic/Uncertainty-based: UCB1 (Upper Confidence Bound).
- Bayesian: Thompson Sampling.
- Multiplicative-Weight/Adversarial: Hedge-style multiplicative weights and EXP3.
Process: The algorithm iterates through a fixed "play budget" ( $T$ ), adaptively allocating trials to explore uncertain neurons and exploit those that repeatedly appear safe to remove. Once the budget is exhausted, the neurons with the highest estimated safe-removal rewards are permanently pruned.

Key Contributions

Formal Formulation: The paper provides a mathematical formulation of neuron pruning as a fixed-budget MAB ranking problem, extending previous MAB-based weight pruning work from scalar weights to complete computational units.
Unified Framework: A single algorithmic framework is presented that supports both stochastic policies (Epsilon-Greedy, Softmax, UCB1, Thompson Sampling) and multiplicative-weight policies (Hedge, EXP3).
Comprehensive Evaluation: The method is empirically evaluated across three distinct domains:
- Tabular Classification: 16 datasets.
- Tabular Regression: 6 datasets (using $R^2$ as the metric).
- Deep Learning Benchmarks: Various architectures including LeNet, AlexNet, ConvNets, LSTMs, Siamese networks, and memory networks on image, text, and reasoning tasks.
Statistical Rigor: The study employs Friedman tests followed by Nemenyi post-hoc tests to establish statistically significant differences between the proposed methods and baselines.

Results

Tabular Classification: UCB1 achieved the highest mean rank among pruning policies, outperforming the unpruned neural network and ranking competitively against standard classifiers like SVMs and Random Forests.
Tabular Regression: UCB1 again obtained the highest mean rank for $R^2$ , showing statistical superiority or competitiveness against several standard regression models (e.g., OLS, LASSO, SVM regression).
Deep Learning: UCB1 and Thompson Sampling obtained the strongest ranks. In several cases (e.g., LeNet on MNIST, AlexNet on ImageNet), MAB-based pruning not only preserved but improved accuracy while removing substantial proportions of neurons (e.g., 62% of a layer in LeNet).
Comparison to Baselines: The MAB-based policies significantly outperformed the unpruned model, magnitude-based neuron pruning, and greedy activation-variation pruning. The non-bandit baselines consistently received the lowest mean ranks.
Policy Performance: While UCB1 and Thompson Sampling were generally the strongest, simple stochastic policies like Epsilon-Greedy also performed well, particularly in deep learning tasks. Hedge-style and EXP3 policies were less consistent in this setting, likely because the pruning task behaves more like a stochastic best-arm identification problem than a fully adversarial sequence.

Significance and Claims
The paper claims that adaptive, loss-aware bandit sampling provides a more effective criterion for structured neuron pruning than static magnitude-based or greedy activation-variation baselines. The significance lies in demonstrating that MAB algorithms can efficiently navigate the trade-off between exploration (testing uncertain neurons) and exploitation (removing safe neurons) under a limited computational budget.

The authors assert that this approach is computationally practical, avoiding the high costs of exhaustive evaluation or Hessian computation required by second-order sensitivity methods. The results suggest that MAB-based pruning can serve as a form of structural regularization, removing redundant units to improve or maintain performance while reducing model size. The paper concludes that UCB1 and Thompson Sampling offer the best balance of exploration, exploitation, and reliability for this specific task, though the framework is flexible enough to accommodate other policies. Future work is suggested to include iterative pruning with fine-tuning, joint pruning across multiple layers, and hardware-aware reward functions.