🤖 AI

Loss-Aware Feature-Map Pruning in Convolutional Neural Networks Using Multi-Armed Bandits

This paper proposes a loss-aware feature-map pruning framework for convolutional neural networks that utilizes multi-armed bandit algorithms, specifically UCB1 and Thompson Sampling, to efficiently identify and remove redundant convolutional channels while preserving model accuracy across various datasets.

Original authors: Salem Ameen, Sunil Vadera

Published 2026-07-28

📖 4 min read☕ Coffee break read

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 are trying to teach a super-smart robot how to recognize pictures of cats, dogs, or flowers. To do this, the robot uses a "brain" made of math called a Convolutional Neural Network (or ConvNet for short). Think of this brain as a massive factory with thousands of tiny workers (called filters) and conveyor belts (called feature maps) that process an image layer by layer. The more workers and belts the factory has, the better it gets at spotting details, but the slower and heavier it becomes. If you try to run this factory on a small, battery-powered device like a smartphone or a drone, it might be too heavy to carry or too slow to react in time.

To fix this, scientists use a technique called "pruning." Imagine you are a gardener trimming a giant, overgrown bush. You want to cut away the dead or useless branches so the plant is smaller and easier to manage, but you must be careful not to cut off the parts that actually produce the fruit. In the world of AI, this means removing parts of the network that aren't helping much. The tricky part is figuring out which branches to cut. If you cut the wrong ones, the robot forgets how to recognize a cat. If you cut too many, it gets confused. The big question is: How do you find the useless parts without spending years testing every single branch one by one?

This paper introduces a clever, game-like strategy to solve that problem. The authors, Salem Ameen and Sunil Vadera, propose using a mathematical concept called "Multi-Armed Bandits." Imagine you are in a casino with a row of slot machines (the "arms"). You have a limited number of coins (the "budget") to play. Your goal is to figure out which machines pay out the most, but you can't afford to play every machine a thousand times. Instead, you play a few times, see which ones seem lucky, and then focus your remaining coins on the best ones.

In this study, the "slot machines" are the feature maps (the conveyor belts) in the AI's brain. The "coins" are the computer's time and energy. The researchers set up a system where the computer doesn't just guess which maps to cut based on how big they look (a common but often inaccurate method). Instead, it plays a game: it temporarily hides one map, sees if the robot still recognizes the picture correctly, and then brings the map back. If hiding the map didn't hurt the robot's performance, that map gets a "safe to remove" score. The Multi-Armed Bandit algorithm acts like a smart manager, deciding which maps to test next based on what it has learned so far, rather than testing them all randomly or in a fixed order.

The paper finds that this "smart manager" approach is incredibly effective. When they tested it on various image datasets (like MNIST for handwritten numbers, CIFAR for everyday objects, and even specific bird and flower datasets), the method successfully removed a significant chunk of the network—up to about 29% of the feature maps in some cases—without making the robot any worse at its job. In fact, in many cases, the pruned robot was just as accurate as the original, untrimmed version.

The researchers compared their method against two other common ways of pruning: "greedy" pruning (which just cuts whatever looks easiest to cut right now) and "magnitude" pruning (which cuts the smallest-looking parts). Their "bandit" method consistently beat these older methods. It was so good that, statistically, it performed just as well as the original, unpruned model, but with a much smaller, faster brain. The paper suggests that by using this adaptive, loss-aware strategy, we can shrink powerful AI models to fit on smaller devices without sacrificing their smarts, all while saving time and energy compared to the old, brute-force ways of testing every single part.

Technical Summary: Loss-Aware Feature-Map Pruning in Convolutional Neural Networks Using Multi-Armed Bandits

Problem Statement

Convolutional Neural Networks (ConvNets) often contain redundant feature maps that increase storage and inference costs without proportional gains in predictive performance. While pruning individual weights (unstructured pruning) reduces parameter count, it often results in irregular sparsity that is difficult to accelerate on standard hardware. Structured pruning, which removes complete feature maps and their associated filters, offers a more practical solution for reducing dense tensor dimensions and Floating Point Operations (FLOPs).

The primary challenge in structured feature-map pruning is determining which maps can be safely removed. "Oracle-style" or direct evaluation methods, which temporarily remove each candidate map and measure the resulting loss change, provide accurate importance estimates but are computationally prohibitive for large datasets and wide layers. Conversely, heuristic methods based on activation statistics or filter magnitude are computationally cheap but may fail to capture the true functional impact of a map on the network's loss. This paper addresses the trade-off between accurate importance estimation and computational cost.

Methodology

The authors propose a framework that formulates feature-map pruning as a fixed-budget, single-arm Multi-Armed Bandit (MAB) problem. The process is divided into two distinct phases: a search phase and a final pruning phase.

1. Search Phase (MAB Exploration)

In this phase, no structural changes are made to the network. Instead, the algorithm iteratively evaluates candidate feature maps to learn their "safe-removal" scores.

Arm Definition: Each candidate feature map in the target convolutional layers constitutes an "arm."
Play Mechanism: At each play step $t$ , the bandit policy selects one arm (a specific feature map).
Temporary Masking: The selected feature map is temporarily masked (set to zero) for a sampled mini-batch. The loss is computed, and the map is immediately restored.
Reward Calculation: The change in loss ( $\Delta$ $Δ$ ) caused by the temporary removal is converted into a reward.
- If the loss decreases or remains stable, the removal is considered "safe," yielding a positive reward.
- If the loss increases significantly, the removal is considered harmful.
Policies: The study evaluates two adaptive policies for allocating the evaluation budget:
- UCB1 (Upper Confidence Bound 1): Combines the empirical mean reward with an exploration bonus to balance exploitation of known safe maps and exploration of uncertain ones.
- Thompson Sampling: Uses a Bayesian approach with a Beta posterior distribution to sample success probabilities, naturally balancing exploration and exploitation.

2. Final Pruning Phase (Structural Deletion)

Once the fixed play budget ( $T$ ) is exhausted, the search concludes.

Ranking: All candidate feature maps are ranked based on their learned safe-removal scores (empirical mean for UCB1, posterior mean for Thompson Sampling).
Top-k Selection: The top- $k$ feature maps with the highest scores are selected for permanent removal.
Structural Deletion: The selected maps are removed along with their producing filters, biases, and the corresponding input-channel kernels in the subsequent layer (whether convolutional or fully connected). This ensures the resulting model is structurally smaller and computationally efficient.

Key Contributions

The paper outlines six specific contributions:

Formal Formulation: A formalization of convolutional feature-map pruning as a fixed-budget top-k selection problem using MAB, where each play evaluates one candidate map.
Loss-Aware Masking: A procedure for estimating removability via temporary masking without permanently altering the network during the search phase.
Algorithmic Implementation: The development of UCB1 and Thompson Sampling algorithms specifically adapted for feature-map pruning, utilizing loss-change rewards to guide selection.
Structural Deletion Rule: A precise rule for deleting producing filters, biases, and next-layer input kernels, covering both convolutional and fully connected transitions.
FLOPs Analysis: A comprehensive analysis of computational reduction, accounting for reductions in both the current layer and the subsequent layer's input channels.
Empirical Validation: Extensive evaluation across six datasets (MNIST, CIFAR-10/100, SVHN, CUB-200-2011, Oxford Flowers 102) with statistical comparisons using Friedman and Nemenyi tests.

Experimental Results

The methods were evaluated on LeNet (MNIST), ConvNet architectures (CIFAR, SVHN), and AlexNet (Birds-200, Flowers 102).

Agreement with Oracle: On LeNet/MNIST, the learned rewards from UCB1 and Thompson Sampling showed strong correlation (Pearson coefficients of 0.83 and 0.80, respectively) with direct oracle-style evaluation, validating that the MAB approach approximates exhaustive search efficiently.
Accuracy Preservation: Across all datasets, UCB1 and Thompson Sampling generally maintained or slightly improved the accuracy of the original unpruned models.
- For example, on CIFAR-10, the original accuracy was 0.81; UCB1 and TS achieved 0.82 after removing ~27% of feature maps.
- In contrast, magnitude-based pruning consistently yielded lower accuracy (e.g., 0.74 on CIFAR-10), and greedy pruning was less stable.
Statistical Significance: Friedman and Nemenyi tests confirmed statistically significant differences between methods.
- UCB1 achieved the highest mean rank (4.28), followed by Thompson Sampling (4.10).
- Both proposed methods significantly outperformed greedy and magnitude-based baselines.
- Crucially, the performance of the pruned models was statistically comparable to the original unpruned models, indicating that accuracy was preserved despite significant structural reduction.

Significance and Claims

The authors claim that this work provides a practical compromise between the accuracy of loss-aware pruning and the efficiency of heuristic methods. By utilizing MAB policies, the method avoids the computational expense of evaluating every candidate feature map on the full dataset while still identifying redundant components based on their actual impact on loss.

The study emphasizes that the separation of the search phase (temporary masking) and the pruning phase (structural deletion) is a key methodological distinction that avoids confusion with multi-play MAB approaches. The results suggest that UCB1 and Thompson Sampling are suitable policies for structured compression, enabling the reduction of FLOPs and model size while preserving predictive performance across diverse image classification tasks. The authors note that future work should focus on applying these methods to modern architectures (e.g., ResNets, Transformers) and measuring actual hardware latency and energy consumption.