AMC: pruning ratios as a reinforcement learning problem

Take a ResNet-50 and allow each of its roughly fifty prunable layers to keep one of ten possible fractions of its weights. That is 10^50 ways to spend one overall budget. The Sensitivity analysis of part 03 escapes the explosion by never looking at combinations at all: it prunes one layer at a time, runs 50 x 10 = 500 single-layer experiments, and reads each layer's Pruning ratio off its own curve where it crosses the accuracy threshold.

The saving comes from a hidden assumption: that the accuracy lost by pruning several layers at once is the sum of what each loses alone. The AMC authors tested exactly this and report that the sensitivity approach "assumes that errors of different pruned layers can be summed up linearly, which does not stand according to our experiments" (He et al., 2018). That is the Layer interaction that slide 24 says the analysis ignores. Pruning layer 3 changes what layer 4 receives, so the curve you measured for layer 4 on a dense network is the wrong curve once layer 3 is thin. The ratios are usable, but they are sub-optimal, and they belong to the network they were measured on. Change the architecture and every curve must be measured again.

The bottleneck is people

Slide 25 shows the practical consequence. Conventionally the allocation relies on human expertise and trial and error. Three customers with three engineers is a workable shop. Many customers overwhelm even a team of engineers (the red cross and the alarmed manager in the middle column), and every new model, new target device or new budget restarts the trial and error from scratch, because rules such as prune the first layer less and the fully connected layers more are, in the paper's words, "non-optimal, and doesn't transfer from one model to another" (He et al., 2018). The right column serves the same crowd with one engineer and an engine that takes the model and the constraint and returns the compressed network.

Slide 26 names the ambition: a push-the-button solution. Today an efficient network needs someone who is both a machine learning expert and a hardware expert. The goal is that a non-expert plus a hardware-centric AutoML tool reaches the same efficient network. This is the definition of Automated pruning in this lecture: an algorithm, not a person, chooses the per-layer ratios for an overall target, and it does so with the target hardware in the loop.

Follow one pass of AMC over MobileNet. Layer t-1 has just been pruned at 30%. The agent now looks at layer t: it receives a short vector describing that layer, its index, its channel counts, its kernel size, its FLOPs, how many FLOPs earlier layers already gave up and how many remain in later layers. It answers with one number, say 50%. The environment removes half of that layer's channels by magnitude and presents layer t+1. When the last layer is done, the whole pruned network is evaluated on a held-out set, and a single number comes back: minus the error. That number is the only feedback the agent ever gets, and it arrives once per pass.

That story contains every ingredient of a reinforcement learning problem, which is what slide 27 means by "pruning as a reinforcement learning problem". The paper's own description: the agent "processes the network in a layer-wise manner", "receives a layer embedding s_t", "outputs a precise compression ratio a_t", and after every layer is compressed the "validation accuracy of the pruned model with all layers compressed is evaluated without fine-tuning, which is an efficient delegate of the fine-tuned accuracy" (He et al., 2018). Only the best network found after the search is fine-tuned.

Why the action must be continuous, and why an actor and a critic

A pruning ratio is a real number, and the compressed model's accuracy is very sensitive to it. A discrete agent would need a grid fine enough to tell 48% from 52%, which explodes the number of actions, and a grid also throws away the order between neighbouring ratios (He et al., 2018), so the agent could not exploit the fact that neighbouring ratios behave alike. DDPG (Lillicrap et al., 2015) is the standard actor-critic method for continuous control: it concurrently learns a Q-function and a policy and can only be used with continuous action spaces (OpenAI Spinning Up). Sutton and Barto give the naming: when a value function is used "to assess, or criticize, the policy's action selections, then the value function is called a critic and the policy is called an actor". In AMC the critic learns, from replayed transitions, how much reward a state and ratio pair leads to; the actor is then nudged by gradient ascent on that estimate to propose ratios the critic scores higher. The critic judges, the actor proposes, and neither ever sees a gradient of the pruned network itself.

Why the reward is minus error under a constraint

Reinforcement learning maximizes reward, so maximizing -Error is minimizing error. The budget appears as a constraint rather than a bonus on purpose: if the agent were paid for shrinking, it would keep shrinking past the target and lose accuracy for nothing. With a hard constraint there is no reward for finishing under budget, so the agent "can precisely arrive at the target compression ratio" (He et al., 2018). The minus infinity branch is the slide's compact way of saying an infeasible network earns nothing useful.

A wasted episode teaches the critic nothing, so the paper never lets one happen. Before each action it computes the duty: the share of the removal target alpha W_all that this layer must carry, given what earlier layers already removed (W_reduced) and the most the remaining layers could remove at the cap a_max (0.8 for convolutions and 0.98 for fully connected layers in fine-grained pruning; 0.8 for every layer in channel pruning). If the actor's proposal is too timid, it is raised to the duty. The budget is therefore met by construction, and the minus infinity case is a formal backstop, not a common outcome.

Slide 28 shows the other protocol. When the goal is the smallest model with no accuracy loss, the budget is unknown in advance, so the reward itself must push on size. The bubble chart on the slide (Canziani et al., 2016) shows that top-1 accuracy rises roughly with the logarithm of the operation count, so multiplying -Error by log(FLOPs) gives, in the paper's words, "a small incentive for reducing FLOPs" while staying "sensitive to Error". The ResNet-50 and MobileNet results in the next two concepts use the constrained reward, not this one.

The constraint need not be FLOPs. The paper notes that the resource can be "FLOPs or the actual inference time on mobile device", and slide 29 adds the mechanism: a pre-built Latency lookup table of measured per-layer times on the target phone, the same device-in-the-loop idea NetAdapt uses in part 05. Substituting measured time for FLOPs turns a FLOPs-constrained search into a latency-constrained one, and the MobileNet table at the end of this part shows why that matters.

Slide 30 puts the agent up against Song Han's own hand-pruned ResNet-50 from his PhD thesis. The human expert kept 29% of the weights, a 3.4x reduction. The agent kept 20%, a 5x reduction, at the same accuracy: 76.13% top-1 for the original and 76.11% for AMC (He et al., 2018). This experiment uses fine-grained weight pruning, not channel pruning, and it is run as Iterative pruning in four rounds at 50%, 35%, 25% and 20% overall density with 30 fine-tuning epochs after each round, which is the schedule part 06 explains.

What the agent noticed inside a residual block

The per-region bars hide the interesting pattern, which slide 31 shows layer by layer. ResNet-50 is built from bottleneck blocks: a 1x1 convolution that reduces the channel count, a 3x3 convolution, and a 1x1 convolution that restores it (He et al., 2016). Walk along the layer index and the kept density forms a sawtooth. The peaks, where the agent kept more, are the 1x1 layers. The crests, where it cut deep, are the 3x3 layers. The paper draws the conclusion directly: the agent "automatically learns that 3x3 convolution has more redundancy than 1x1 convolution and can be pruned more" (He et al., 2018). Nobody told it about kernel shapes; the kernel size k is just one of the eleven features in the state.

The finding makes sense once you think about where redundancy lives. A 1x1 kernel mixes channels and has no spatial extent, so every weight is a distinct channel-to-channel connection with nothing next to it to stand in for it. A 3x3 kernel has nine spatial taps per channel pair, many of which are near zero after training, and it also holds most of a block's weights: in a stage-3 bottleneck the 3x3 layer carries about 590K of the block's 1.1M parameters against 262K for each 1x1 layer (Han, 2017). Cutting the 3x3 layers hard is where the compression is, and it is also where accuracy is cheapest to spend, which is exactly the trade the reward rewards.

The final slide moves from weights to milliseconds. MobileNet is run with TF-Lite on a Samsung Galaxy S7 Edge with a Qualcomm Snapdragon SoC, single core, batch size 1, which is the latency-oriented setting of a phone serving one image at a time. Four models are measured: the full 1.0 MobileNet, two AMC results and the 0.75 MobileNet baseline, which is the same architecture with every layer's width scaled by 75%.

Two lessons in one table

First, FLOPs are a proxy. Halving the multiply-accumulates bought 1.85x, not 2x, because a phone does not spend time in proportion to arithmetic. MobileNet's depthwise layers have a low ratio of computation to memory traffic, so trimming them saves fewer milliseconds than the MAC count suggests; the paper reports about 2x speedup on the 1x1 convolutions but less on the depthwise ones (He et al., 2018). When the agent is instead constrained by measured time from the lookup table, it removes work where the phone is actually slow, and the search lands at 59.7 ms, a 1.99x speedup that all but hits the 2x target. That allocation is less accuracy-friendly, which is the 0.3 point cost. You choose the target you care about, and you pay in the other currency.

Second, non-uniform beats uniform. Uniform shrinking with the width multiplier is a single knob turned on every layer, and the thinner model must be trained from scratch (Howard et al., 2017). AMC prunes a pretrained network with a different ratio per layer, keeps 2.1 (50% FLOPs) or 1.8 (50% Time) more points of accuracy with fewer MACs than 0.75 MobileNet, and does it after a short fine-tune. In the accuracy versus latency plane the AMC points sit above and to the left, which is the definition of a better Pareto frontier.