Reading Notes on E³-PRUNER

This is a pruning method proposed by Huawei.

At a high level, it has two key steps:

1. Prune the model.
2. Fine-tune the model.

During the training loop, pruning and fine-tuning are optimized jointly in the early iterations. The remaining iterations perform fine-tuning only.

Searching for a Pruning Configuration

Naively, given a parameter $K$, finding the best way to remove $K$ layers from a model would require brute-force enumeration followed by evaluation. For every dataset, this means evaluating

$C(L,K)=\frac{L!}{K!(L-K)!}$

possible configurations.

Even setting aside the cost of millions or trillions of configurations multiplied by tens of thousands of samples, pruning $K$ layers all at once can easily produce a false optimum. Every candidate is likely to emit nonsense immediately after such aggressive pruning, and choosing the best result from a collection of broken outputs is not meaningful.

Removing layers one at a time can also lead to mistakes. Once a layer has been removed, the decision cannot be undone. A layer that appears unimportant on the current data may in fact be useful elsewhere, so a greedy procedure is unlikely to find the global optimum.

These problems arise from discrete pruning. Discrete decisions leave no room for error correction. $\text{E}^3\text{-PRUNER}$ instead uses $Gumbel\text{-}TopK$ to turn a discrete $Top\text{-}K$ choice into a differentiable approximation.

The Gumbel-TopK Algorithm

Requires: layer scores $S$, temperature $\tau$, and number of retained layers $k$.

$$\begin{aligned} &1: \quad S_1 = S + \text{Gumbel}(0,1) \\ &2: \quad \tilde{M}_1 = 0 \\ &3: \quad \textbf{for} \; i = 1 \; \textbf{to} \; k \; \textbf{do} \\ &4: \quad\quad S_{i+1} = S_i + \log(1 - \text{Softmax}(\frac{S_i}{\tau})) \\ &5: \quad\quad \tilde{M}_{i+1} = \tilde{M}_i + \text{Softmax}(\frac{S_i}{\tau}) \\ &6: \quad \textbf{end for} \\ &7: \quad M = \text{STE}(\text{TopK}(\tilde{M}_{k+1}, k)) \\ &8: \quad \textbf{return} \; \text{Pruning mask} \; M \end{aligned}$$

As shown above, every layer has a score $S$. A Gumbel noise term is added to each score at initialization, preventing tiny changes among nearly equal scores from producing disproportionately large changes in the selection.

Line 4 “penalizes” repeatedly selected high scores so that the method does not keep choosing the same few layers. The crucial component is the STE (Straight-Through Estimator), which makes backpropagation possible through the selection operation.

The temperature $\tau$ follows a linear annealing schedule, allowing broader exploration early in training.

The iterative update is:

$\tilde{M}_{i+1} = \tilde{M}_i + \text{Softmax}(\frac{S_i}{\tau}), \quad S_{i+1} = S_i + \log(1 - \text{Softmax}(\frac{S_i}{\tau}))$

Fine-Tuning

1. Joint Training

After a particular layer is removed, the model weights must be fine-tuned accordingly; otherwise accuracy drops sharply.

Algorithm 1:

Parameters:

• LLM $\Theta$
• Total number of iterations $T$
• Number of mask-search iterations $T_M$
• Layer scores $S$
• Layer mask $M$
• Total number of layers $L$

$$\begin{aligned} &1: \quad \textbf{for} \; t = 1 \; \textbf{to} \; T \; \textbf{do} \\ &2: \quad\quad \tau = 1 - \beta \cdot t/T, \; k' = L - \lceil (L - k) \cdot t/T \rceil \\ &3: \quad\quad \textbf{if} \; t \leq T_M \; \textbf{:} \\ &4: \quad\quad\quad M = \text{Gumbel-TopK}(S, \tau, k') \\ &5: \quad\quad\quad S = S - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial S \\ &6: \quad\quad \textbf{end if} \\ &7: \quad\quad \Theta = \Theta - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial \Theta \\ &8: \quad \textbf{end for} \\ &9: \quad \textbf{return} \; \text{Pruned model} \; \Theta^* \end{aligned}$$

The important detail is that the STE (Straight-Through Estimator) makes the backward pass possible.

The Gumbel-TopK procedure is written as Algorithm 2 in my notes:

Parameters:

• Layer scores $S$
• Temperature $\tau$
• Number of layers $k$

$$\begin{aligned} &1: \quad S_1 = S + \text{Gumbel}(0,1) \\ &2: \quad \tilde{M_1} = 0 \\ &3: \quad \text{for } i = 1 \text{ to } k \text{ do}\\ &4: \quad\quad\quad M = \text{Gumbel-TopK}(S, \tau, k') \\ &5: \quad\quad\quad S = S - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial S \\ &6: \quad\quad \textbf{end if} \\ &7: \quad\quad \Theta = \Theta - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial \Theta \\ &8: \quad \textbf{end for} \\ &9: \quad \textbf{return} \; \text{Pruned model} \; \Theta^* \end{aligned}$$

The loss function described in the paper is:

$\left\{\Theta^{*}, M^{*}\right\} = \arg\min_{\Theta, M} \mathcal{L}(x ; \Theta, M) \quad \text{s.t.} \; \|M\|_0 = k$

This leads to the adaptive distillation loss:

$\mathcal{L}_{\text{adaptive-KD}} = \frac{1}{N} \sum_{i=1}^{N} \mathcal{H}\left(\text{Softmax}\left(z_t^{(i)}\right)_c\right) \cdot \sum_{c \in S_i^K} \text{Softmax}\left(z_t^{(i)}\right)_c \cdot \log\left(\frac{\text{Softmax}\left(z_t^{(i)}\right)_c}{\text{Softmax}\left(z_s^{(i)}\right)_c}\right)$

It has two parts. First, the token loss:

$\mathcal{L}_{\text{token}} = \sum_{c \in S_i^K} \text{Softmax}\left(z_t^{(i)}\right)_c \cdot \log\left(\frac{\text{Softmax}\left(z_t^{(i)}\right)_c}{\text{Softmax}\left(z_s^{(i)}\right)_c}\right)$

Here, $z_s$ is the student’s output and $z_t$ is the teacher’s output. Because running the teacher online during training would be expensive, its top-$K$ logits are computed and stored in advance.

The weighting term is:

$\mathcal{W}_i = \mathcal{H}\left(\text{Softmax}\left(z_t^{(i)}\right)_c\right)$

2. Post-Training

As shown in Algorithm 1, after the joint-training phase the remaining iterations optimize only the model parameters. The goal is to recover accuracy while keeping the pruned architecture fixed.

Further Thoughts

Training against the output distribution often has a higher ceiling than directly aligning vector directions, because directional alignment does not necessarily make the final logits more accurate.

Once a large model produces a hidden representation $h$, the final logits are obtained from

$z = W_u h$

After pruning, the predicted hidden representation will deviate to some extent. Let $h'=h+\Delta h$. Then:

$z' = W_u h' = W_u (h + \Delta h) = z + W_u \Delta h$

The final logit distribution therefore does not depend only on the similarity of the vector $\Delta h$. It depends on the complete product $W_u\Delta h$.

Even when $\Delta h$ is large, the logits can still be approximated well if the distribution of its elements overlaps strongly with the null space of $W_u$.

Conversely, even a small $\Delta h$ can produce a large logit error $\Delta z$ if it happens to align with directions in which $W_u$ has large values. This is why directly aligning hidden-vector directions can have a lower performance ceiling.

The corresponding advantage is lower training cost. The learned vectors are also unlikely to differ drastically from those produced by output-distribution training. In theory, aligning the outputs during post-training should therefore recover better results.