Technical Summary: HiReLC – Hierarchical Reinforcement Learning for Neural Network Compression

1. Problem Statement

Deep neural networks, particularly Vision Transformers (ViTs), face significant deployment challenges on resource-constrained platforms due to high parameter counts, memory footprints, and inference latencies. While existing compression techniques such as static heuristics, one-shot mixed-precision assignment, and single-objective optimization exist, they often fail to capture cross-layer dependencies and lack adaptive refinement during fine-tuning. Furthermore, Reinforcement Learning (RL) approaches for compression have historically struggled with high variance in large action spaces, coarse granularity, and a lack of sensitivity-aware guidance.

The paper addresses the need for a general-purpose, architecture-agnostic framework capable of jointly optimizing structured pruning and mixed-precision quantization while navigating the complex trade-offs between model size reduction and accuracy retention.

2. Methodology: The HiReLC Framework

HiReLC (Hierarchical Reinforcement Learning for Neural Network Compression) is a two-tier hierarchical ensemble-RL framework designed to decompose the compression search space. The system is architecture-agnostic at the controller level, relying on a modular layer abstraction to decouple the RL environment from specific network topologies.

2.1 Hierarchical Agent Structure

The framework employs two distinct levels of agents:

• High-Level Agents (HLAs): These agents operate at the network level to coordinate global budget allocation. They assign a "LayerBudget" to each compressible block, determining target retained-size fractions ($R_i$), pruning limits, and bitwidth bounds. HLAs utilize Fisher Information-based sensitivity estimates to guide budget allocation, ensuring that blocks with high sensitivity receive lighter compression.
• Low-Level Agents (LLAs): Operating independently per block, LLAs refine the global budgets into specific per-kernel configurations. They select from a four-dimensional multi-discrete action space for each kernel:
  1. Bitwidth ($b_{i,k}$): Selecting precision levels.
  2. Keep-ratio ($\rho_{i,k}$): Determining the fraction of output channels to retain (structured pruning).
  3. Quantization Type ($\tau_{i,k}$): Choosing between Integer (INT) and Floating-point (FLOAT).
  4. Granularity ($\mu_{i,k}$): Selecting among uniform, log-domain, per-channel, or learned sign-magnitude schemes.

2.2 Ensemble Learning and Consensus

To mitigate the variance inherent in RL for large discrete spaces, HiReLC utilizes an ensemble approach.

• LLA Ensemble: Multiple parallel policies (combining PPO and A2C) are trained with heterogeneous reward weights (varying emphasis on accuracy vs. compression). A consensus action is derived via weighted voting, blending conservative, moderate, and aggressive strategies to produce a robust configuration.
• HLA Ensemble: Similarly, an ensemble of HLAs proposes budget allocations, with the final budget determined by majority voting or arithmetic mean aggregation.

2.3 Active Learning and Surrogate Modeling

To address the computational cost of evaluating every candidate configuration via full post-compression fine-tuning, HiReLC integrates an iterative active learning loop:

• Logit-MSE Proxy: During the "cold-start" phase (before sufficient data is collected), a lightweight proxy based on the Mean Squared Error (MSE) between the logits of the compressed model and the baseline model is used to estimate accuracy retention.
• MLP Surrogate: Once a minimum number of configurations are evaluated, a lightweight Multi-Layer Perceptron (MLP) surrogate is trained on the collected data to predict accuracy. This surrogate is used for reward shaping during RL optimization, while full fine-tuning is reserved for final cycle-level evaluations.
• Iterative Loop: The process alternates between RL policy optimization, configuration application, fine-tuning, and surrogate updates.

2.4 Reward Formulation

The reward functions are designed to balance accuracy, compression, and stability:

• LLA Rewards: Include terms for accuracy retention ($R_{acc}$), compression gain ($R_{comp}$), budget compliance ($R_{budget}$), and stability ($R_{stab}$). A sensitivity penalty ($P_{sens}$) is applied to discourage aggressive compression on highly sensitive blocks.
• HLA Rewards: Focus on balancing heterogeneity across blocks ($R_{bal}$), sensitivity-aware allocation ($R_{sens}$), global target adherence ($R_{global}$), and dynamic feedback based on previous cycle performance ($R_{dyn}$).

3. Key Contributions

The paper outlines three primary contributions:

1. HiReLC Framework: An architecture-agnostic hierarchical ensemble-RL framework for joint structured pruning and mixed-precision quantization. It achieves effective parameter-storage compression ratios of 5.99–6.72× across four architectures (DeiT-Small, CLIP ViT-B/32, ResNet18, MobileNetV2) with accuracy drops ranging from 0.55% to 5.62%, and even a 3.83% accuracy gain in one setting.
2. Sensitivity-Aware Guidance: The integration of Fisher Information sensitivity scores at both hierarchy levels. This ensures critical blocks receive lighter compression without requiring additional search-time evaluations, propagating gradient-based block importance through both reward shaping and action correction.
3. Surrogate-Augmented Active Learning: An iterative loop that interleaves RL optimization with post-compression fine-tuning. By using a lightweight MLP surrogate and a logit-MSE proxy, the framework amortizes expensive evaluations, enabling practical framework-level studies within a limited number of cycles.

4. Experimental Results

Experiments were conducted on CIFAR-10/100 and Tiny ImageNet across Vision Transformers and CNNs.

• Compression Performance: The framework consistently achieved compression ratios between 5.99× and 6.72× (equivalent to 83.3%–85.1% model-size reduction).
• Accuracy Impact:
  • CLIP ViT-B/32 on CIFAR-100: Achieved the smallest drop of 0.55% at 6.63× compression.
  • CLIP ViT-B/32 on CIFAR-10: Demonstrated a 3.83% accuracy gain (from 89.22% to 93.05%), attributed to compression acting as a regularizer for over-parameterized models.
  • ResNet18 and MobileNetV2 on Tiny ImageNet: Showed drops of 4.22% and 5.62% respectively, reflecting the challenges of compressing smaller, already-compact architectures on harder datasets.
• Ablation Studies:
  • Removing the HLA (using uniform budgets) led to over-compression and a significantly higher accuracy drop (6.80% vs. 1.72% for the full system).
  • Using a homogeneous A2C ensemble resulted in less stable trade-offs compared to the mixed PPO/A2C ensemble.
  • Fisher sensitivity guidance improved accuracy retention by 1.30 percentage points and produced more differentiated block-wise budget assignments.
• Mixed Precision: Experiments allowing per-kernel selection of INT/FLOAT types and log-scale granularity yielded lower accuracy drops compared to uniform INT settings, suggesting flexibility aids recovery.

5. Significance and Claims

The paper positions HiReLC as a practical design choice for joint neural network compression, emphasizing the following points:

• Unified Search Space: Unlike prior methods that target single compression axes (pruning or quantization) or rely on static heuristics, HiReLC jointly optimizes both, leveraging the synergies between structured pruning and mixed-precision quantization.
• Architecture Agnosticism: The framework successfully compresses diverse model families (Transformers, CLIP, ResNets, MobileNets) without requiring architecture-specific controller redesign.
• Sensitivity-Awareness: The use of Fisher Information allows the system to automatically identify and protect sensitive network components, improving the accuracy-compression trade-off.
• Practicality: By utilizing surrogate models and active learning, the framework reduces the computational burden of RL-based compression, making it feasible to explore large configuration spaces without exhaustive fine-tuning of every candidate.

The authors note that while the achieved compression ratios are significant, the results are context-dependent (e.g., evaluated on CIFAR/Tiny ImageNet rather than ImageNet-1K in the main tables) and assume packed structured storage. The framework is presented as a starting point for deployment scenarios where model size is the primary constraint, with future work needed to optimize for hardware-specific latency and energy metrics.