Source-linked AI summary
EigenDamage: Structured Pruning in the Kronecker-Factored Eigenbasis
Chaoqi Wang, Roger Grosse, Sanja Fidler, Guodong Zhang
TL;DR
The paper addresses the cost of deploying over-parameterized neural networks while preserving accuracy. It reparameterizes layers in the KFE and applies Hessian-based structured pruning there. EigenDamage outperforms strong parameter-coordinate baselines, especially on challenging datasets and networks.
Problem
Over-parameterized networks impose high test-time computation and memory costs, while existing pruning methods can rely on restrictive independence assumptions or unstructured sparsity.
Method
EigenDamage reparameterizes fully connected and convolutional layers in the KFE, where the Hessian is closer to diagonal, then applies Hessian-based structured pruning.
Results
EigenDamage empirically outperforms strong baselines that prune in original parameter coordinates, especially on more challenging datasets and networks.
Takeaways & Limitations
KFE pruning yields low-rank bottleneck structures with fast computation and can automatically calibrate layer ranks from a single compression ratio.
Takeaways & Limitations
The Fisher approximation is justified under stated model-distribution conditions, including exponential-family predictive distributions such as categorical or Gaussian distributions.
Abstract
from arXiv · showhide
Reducing the test time resource requirements of a neural network while preserving test accuracy is crucial for running inference on resource-constrained devices. To achieve this goal, we introduce a novel network reparameterization based on the Kronecker-factored eigenbasis (KFE), and then apply Hessian-based structured pruning methods in this basis. As opposed to existing Hessian-based pruning algorithms which do pruning in parameter coordinates, our method works in the KFE where different weights are approximately independent, enabling accurate pruning and fast computation. We demonstrate empirically the effectiveness of the proposed method through extensive experiments. In particular, we highlight that the improvements are especially significant for more challenging datasets and networks. With negligible loss of accuracy, an iterative-pruning version gives a 10$\times$ reduction in model size and a 8$\times$ reduction in FLOPs on wide ResNet32.
1. Introduction
The paper targets the test-time cost of over-parameterized networks and develops structured pruning in a coordinate system designed to better capture weight independence. EigenDamage produces loss-aware bottleneck structures and performs especially well on challenging datasets and networks.
- Over-parameterization improves generalization but increases test-time computation and memory overhead on resource-constrained devices.
- 10× or more compression has been achieved with limited accuracy loss, using train, prune, and fine-tune stages.
- Magnitude pruning can remove important weights or retain unimportant ones, motivating Hessian-based criteria such as OBD and OBS.
- Unstructured weight pruning does not provide computational speedups without specialized hardware and software, whereas structured pruning preserves deployable convolutional structure.
- Existing factorization assumptions may hurt pruning because posterior weights can be highly coupled.
- EigenDamage reparameterizes layers in the KFE, where the Hessian is closer to diagonal, enabling lower-cost OBD pruning.
- EigenDamage creates loss-aware low-rank bottlenecks, applies across fully connected and convolutional layers, and automatically calibrates ranks from one compression ratio.
2. Background
The background develops curvature-based pruning through Laplace, Fisher, and K-FAC approximations, then introduces the KFE as an eigenbasis for a Kronecker-factored curvature matrix. Structured pruning removes filters or channels while preserving convolutional structure, but existing methods can require target models or specialized architectures.
- Laplace Approximation: The Laplace approximation models weights near a trained local MAP solution as a Gaussian using the Hessian, with the Fisher often used as a positive-semidefinite approximation.
- K-FAC: K-FAC approximates layerwise Fisher matrices with Kronecker factors under independence assumptions, extending to convolutions with additional spatial assumptions.
- K-FAC: K-FAC factorization avoids exact Fisher storage costs and enables efficient Fisher-vector products, inversion, and eigendecomposition.
- K-FAC: The KFE consists of the Kronecker product of the factor eigenvector matrices and serves as the eigenbasis of the approximated curvature matrix.
- Structured Pruning: Structured pruning reduces network size while retaining convolutional structure, commonly by removing channels or filters and their corresponding feature maps.
- Structured Pruning: Many channel-pruning methods require predefined target models or specialized architectures, limiting their ease of use.
3. Revisiting OBD and OBS
OBD and OBS estimate pruning damage from a second-order loss approximation, but they differ in how they treat correlations among weights. The paper motivates pruning in the KFE because diagonal approximations are more reasonable there, especially when pruning many weights simultaneously.
- OBD and OBS train to a local minimum, then prune weights predicted to cause the smallest increase in training error.
- The pruning cost is derived from a second-order Taylor expansion around the local mode and can be interpreted through the Laplace approximation.
- OBD: OBD approximates the Hessian with a diagonal matrix and therefore treats weights as uncorrelated.
- OBS: OBS accounts for correlations among weights and updates remaining weights to compensate after pruning.
- OBD and OBS correspond to different factorized-Gaussian KL objectives, with reverse KL underestimating variance and forward KL overestimating it.
- OBS can be advantageous when pruning one weight, whereas OBD can be better when pruning multiple weights simultaneously.
- Correlations are smaller in the KFE, motivating OBD there as a more reasonable diagonal approximation.
4. Methods
The method extends Hessian-based pruning to structured units, then reparameterizes layers in the Kronecker-factored eigenbasis so pruning operates where the Fisher is closer to diagonal. EigenDamage combines this basis with bottleneck-form pruning and optional factorization to reduce computation and model size.
- Structured pruning: OBD and OBS are extended from weight-level pruning to channel or filter-level pruning, with C-OBD summing weight importances and C-OBS using K-FAC for tractable inversion.The structured algorithms compute Kronecker factors and prune filters according to their estimated loss changes before fine-tuning.
- Structured pruning: K-FAC approximates the Fisher as S ⊗ A, sharing the input factor A across filters to reduce storage and make structured importance calculations tractable.This replaces an intractable per-filter Fisher matrix with shared Kronecker factors.
- EigenDamage: EigenDamage rotates weights into the KFE, where the Fisher is diagonal under K-FAC assumptions, and applies structured pruning to both input and output channels.The transformed weights are pruned using C-OBD or equivalently C-OBS because the Fisher is close to diagonal in this basis.
- EigenDamage: Eigenpruning removes rows or columns in the transformed weights and corresponding eigenbasis components, producing a low-rank bottleneck structure.The method is named EigenDamage and removes selected rows or columns according to their estimated loss changes.
- Compression: EigenDamage preserves input and output shapes, applies to convolutional or fully connected architectures without modification, and supports one global compression ratio.The eigenbasis and transformed weights can also be compressed using doubly factored Kronecker and depthwise separable decompositions.
- Iterative pruning: Iterative pruning addresses the breakdown of the Taylor approximation when many filters are removed by pruning only a few filters per iteration and repeating the procedure.Repeated iterations continue until the desired compression ratio is reached, while the resulting bottleneck structures can reduce model size and FLOPs.
- Eigenbasis compression: The eigenbasis size is reduced to 1/k^4 of the original one when its spatial structure is exploited, with Q_A efficiently implemented by a 1 × 1 convolution.The approximation ignores correlations among spatial locations within the same input channel.
5. Experiments
Experiments compare EigenDamage with Hessian-based, channel-level, and low-rank pruning methods across VGGNet, ResNet, and PreResNet on CIFAR10, CIFAR100, and Tiny-ImageNet. EigenDamage generally offers the strongest accuracy–resource tradeoffs, particularly at high pruning ratios, on challenging datasets, and on more sophisticated networks.
- One-pass pruning: At a 90% pruning ratio, channel-level variants outperform NN Slimming by approximately 2% test accuracy on CIFAR100 with VGGNet.On CIFAR100 with ResNet, EigenDamage also achieves approximately 2% better performance than other methods at the same ratio.
- One-pass pruning: EigenDamage performs best across almost all settings, with larger improvements at high pruning ratios and on more complicated networks such as ResNet and PreResNet.Its bottleneck structure preserves input and output dimensions, allowing application to fully connected or convolutional layers without modification.
- One-pass pruning: On Tiny-ImageNet with VGGNet, EigenDamage significantly outperforms all baselines, while producing a more balanced pruning ratio across convolutional layers.The comparison discusses layerwise pruning behavior and links NN Slimming’s poorer performance to retaining filters in top layers.
- Loss analysis: Without finetuning, EigenDamage produces lower training loss after pruning, indicating that KFE pruning more accurately reflects loss sensitivity than pruning in original weight coordinates.For VGG19 on CIFAR10, pruning 80% of weights causes negligible loss increase for EigenDamage.
- Iterative pruning: In iterative pruning, EigenDamage is slightly better than baselines with VGGNet and substantially better with ResNet, while being especially effective at reducing FLOPs.Performance gaps grow on CIFAR100 and on the more sophisticated ResNet architecture.
6. Conclusion
The paper introduces KFE reparameterization to approximately decouple weights for Hessian-based pruning, yielding effective pruning and fast computation through low-rank bottleneck structures.
- KFE reparameterization approximately satisfies entrywise independence, enabling effective Hessian-based weight pruning.
- The resulting pruned networks have low-rank bottleneck structures that support fast computation.
- EigenDamage outperforms strong baselines that prune in original parameter coordinates, especially on more challenging datasets and networks.
A. Derivation of Kron-OBD and Kron-OBS
Kron-OBD and Kron-OBS extend Hessian-based structured pruning using Kronecker-factored Fisher approximations, differing in whether filters are treated as independent or correlated.
- For a convolutional layer, the Fisher information is approximated as F = S ⊗ A using vectorized weights and two Kronecker factors.
- Kron-OBD assumes filters are independent, making S diagonal and allowing filter importance to be computed individually.
- Kron-OBS allows correlations between filters, so S is not diagonal in the structured pruning formulation.
- The constrained optimization for removing a filter is solved with a Lagrange multiplier to obtain the optimal weight change.
- Substituting the optimal weight change into the objective yields the importance of each filter.
B. Algorithm for Solving eqn. (18)
The algorithm solves the tensor-factor optimization problem with alternating least squares, using closed-form factor updates and tensor unfolding expressed through Khatri-Rao products.
- Tensor notation: The Khatri-Rao product is the column-wise Kronecker product of two matrices.
- Tensor notation: A low-rank CP tensor is represented by factor matrices in Kruskal notation.
- Tensor notation: The mode-k unfolding depends on Khatri-Rao products of the other factor matrices through Zk.
- Alternating Least Squares (ALS): Alternating least squares solves problems like eqn. (18) by updating each factor matrix while holding the others fixed.
- Alternating Least Squares (ALS): The factor updates are alternated until convergence or the maximum iteration count is reached, including a Mahalanobis-norm variant.
- Alternating Least Squares (ALS): The update formulation uses vec and unvec as inverse vectorization operators, with Z† defined as (Z⊤Z)−1Z⊤.
C. Additional Results on One-pass Pruning
Additional one-pass results present accuracy–compression trade-off curves and tables across datasets, architectures, pruning methods, and BatchNorm sparsity settings.
- The additional results compare accuracy against reductions in weights and FLOPs for one-pass pruning.
- The tables report one-pass pruning results for VGG19 and ResNet on CIFAR-10, CIFAR-100, and Tiny-ImageNet.
- VGGNet results cover CIFAR-10 and CIFAR-100 under settings with and without L1 BatchNorm sparsity.
- (Pre)ResNet results likewise plot weight and FLOP reductions against test accuracy under both BatchNorm sparsity settings.
- NN Slimming uses a base network pretrained with L1 BatchNorm sparsity, whereas the other methods use normally pretrained networks.