RL-Guided Pruning of CNNs Using Graph Embeddings

RL-Guided Pruning of CNNs Using Graph Embeddings Karima Amrouche11, Yacine Ait Ali Yahia21, and Ilhem Kherroubi31 1Laboratoire de la Communication dans les Syst`emes Informatiques (LCSI), ´ Ecole Nationale Sup´erieure d’Informatique, BP 68M, 16309, Alger, Alg´erie Abstract This paper presents a novel method for compressing Convolutional Neural Networks (CNNs) to enable efficient deployment on low-capacity devices. The proposed approach combines neural network pruning with reinforcement learning (RL) and graph embedding. Each network is represented as a computational graph, and Graph Convolutional Networks (GCNs) are utilized to learn graph-level embeddings that inform pruning decisions. By applying Proximal Policy Optimization (PPO), we automate the selection of layer-wise pruning ratios, eliminating the need for manual tuning. Experiments on ResNet-34 and VGG-19, trained on the CIFAR-10 dataset, demonstrate that our method achieves up to 80% compression while maintaining or improving model accuracy through post-pruning rewinding. We evaluated both structured and unstructured pruning strategies, analyzing the trade-offs in accuracy, FLOPs, parameter count, and inference time. Keywords: Model Compression, Deep Neural Networks, Graph Embedding, Reinforcement Learning, Neural Networks, Pruning, Convolutional Neural Networks, CNN Acceleration. 1 Introduction The deployment of convolutional neural networks (CNN) in low-capacity devices, such as smartphones and IoT devices [1], presents significant challenges due to their high computational demands and large memory requirements. These constraints limit the use of CNNs in real-time applications, such as facial recognition or object detection, especially in environments where cloud resources are unavailable or introduce unacceptable latency. Model compression techniques, particularly neural network pruning, have effectively overcome these challenges. By reducing the complexity of CNNs, pruning reduces computational cost and memory usage, enabling deployment on resource-constrained devices without significant performance loss. However, determining the optimal pruning strategy for each layer remains challenging, as it often requires manual tuning and is sensitive to the network structure. We introduce a method that combines neural network pruning with reinforcement learning (RL) and graph embeddings to automate the compression process. This approach reduces the need for manual tuning while preserving accuracy and efficiency, key requirements for deploying CNNs on low-capacity devices. 2 Related Work Model compression techniques can be categorized into several approaches, including knowledge distillation, quantization, factorization, and pruning [2]. Among these, pruning offers a direct and effective way to reduce model complexity by eliminating redundant parameters, often with minimal loss in accuracy. It is particularly attractive because it can be applied post-training, preserves the original architecture, and complements other techniques such as quantization. The theoretical basis for pruning is strengthened by the Lottery Ticket Hypothesis (LTH) [3], which proposes that overparameterized networks contain smaller, trainable subnetworks—referred to as “winning tickets”—capable of reaching comparable accuracy to the full model when trained in isolation. As Frankle and Carbin describe: 59 “A randomly-initialized, dense neural network contains a subnetwork that is initialized such that—when trained in isolation—it can match the test accuracy of the original network after training for at most the same number of iterations.” [3] This insight highlights the inherent redundancy in large neural networks and motivates pruning as a principled strategy for model compression. However, identifying such subnetworks remains nontrivial, especially in deeper architectures, where brute-force or heuristic pruning methods become computationally prohibitive. This challenge underscores the need for more scalable and intelligent pruning approaches—such as those guided by reinforcement learning and graph-based representations. In practice, pruning techniques are generally classified into two categories: unstructured and structured. Unstructured pruning removes individual weights, often resulting in sparse models with high compression rates. However, these irregular patterns typically require specialized hardware for efficient execution. Structured pruning, in contrast, removes entire filters, channels, or blocks, producing smaller, dense models that are more compatible with conventional hardware and easier to deploy [4,5]. Despite their practical advantages, traditional pruning methods often rely on fixed heuristics to determine layer-wise pruning ratios, which may not generalize well across architectures or datasets. To overcome these limitations, recent work has turned to reinforcement learning (RL) to automate the pruning process. Notable examples include ABCPruner [6] and CCPruner [7], which use RL agents to learn pruning policies that balance model efficiency and accuracy. While these methods have shown promise, they often treat layers independently and fail to account for the structural dependencies across the network. Our work addresses this gap by modeling the neural network as a computational graph, enabling a more holistic view of the architecture. By leveraging graph embeddings, we capture global structural information that informs pruning decisions, leading to more coherent and scalable model compression. 3 Building a Computational Graph from a Neural Network As illustrated in Figure 1, our compression pipeline consists of the following stages: 1. Initialization: The process begins with the initialization of a deep neural network, either from scratch or using a pretrained model. At this stage, the initial weight values are preserved to enable potential rewinding after pruning, as part of the iterative compression strategy. 2. Training: The initialized model is trained on the target task until it achieves satisfactory performance. This yields a fully trained network that serves as the baseline for subsequent pruning. 3. Pruning: Reinforcement learning is employed to determine the optimal pruning rates for each layer. A Graph Convolutional Network (GCN) is used to encode the computational graph of the trained model into a latent state representation. This representation is then processed by a policy network, which generates pruning decisions. These decisions are evaluated based on a reward signal reflecting the trade-off between compression and model accuracy. 4. Rewinding: After pruning, the model is reverted to its initial weights saved during the initialization phase. This step, known as rewinding, facilitates the identification of subnetworks—often referred to as “lottery tickets”—that can be retrained from the original initialization to achieve strong performance. 5. Retraining: The selected subnetwork is retrained from its original initialization. The objective is to recover the model’s performance to a level comparable to the original unpruned network, thereby achieving an efficient and accurate compressed model. 60 Figure 1: Pipeline illustrating CNN pruning guided by reinforcement learning, using graph embeddings to encode architectural information. The model automates the pruning phase, removing the need for human input. A DNN’s computational graph maps operations like addition and multiplication during this phase to produce outputs. Nodes perform computations, while edges guide the data flow through the network’s layers [8]. Algorithm 1outlines the procedure for constructing a subgraph for a single layer, which involves initializing input and output nodes and creating edges to represent the connections between them. Algorithm 1: Subgraph Construction for a CNN Layer Data: n: Number of input channels Input: N: Number of output channels E: List of edges (empty for the first layer) Result: E: Updated list of edges Output: Input node of the next layer 1Output ←Input +N+ 1; 2for i←1to Ndo 3E←Insert(E, (Input, Input +i)) ; // Insert edge eik [input] 4E←Insert(E, (Input +i, Output)) ; // Insert edge eik [output] 4 Graph Embedding Graph embedding transforms structured data into low-dimensional vector representations, facilitating efficient learning and decision-making in downstream tasks. In this work, we represent convolutional neural networks (CNNs) as computational graphs, where nodes correspond to layers or operations and edges capture the flow of information. This graph-based formulation enables the use of Graph Convolutional Networks (GCNs) to encode the topological and functional properties of the CNN architecture. By leveraging GCNs, we obtain a fixed-size embedding of the network that preserves both structural dependencies and feature hierarchies, which are essential for informing pruning decisions [9]. This compact representation is then passed to the reinforcement learning (RL) agent, which uses it to select optimal pruning strategies. The use of GCNs ensures that similar CNN architectures yield similar embeddings, improving the generalization of the pruning policy across different models. As input to the GCN, we define the node features using the number of nodes (|V|) and their attributes (Fin), along with edge indices (2,|E|), and output node features (|V|, Fout). Instead of focusing solely on individual node embeddings, we aim to compute a holistic representation of the entire graph. To achieve this, we first apply a GCN encoder that maps the graph Gto a set of node embeddings H∈RN×d, as shown in Equation 1: 61 H= GCNencoder(G)∈RN×d(1) The resulting node embeddings are then aggregated using a global pooling operation. Specifically, we use a GlobalMeanPool that computes the average over all node embeddings, producing the final graph-level embedding g, as defined in Equation 2: g=1 N N X n=1 hi(2) In Equation 2,hidenotes the embedding of the i-th node, Nis the total number of nodes in the graph, and dis the dimensionality of the embedding space. This aggregated representation gcaptures the structural and semantic properties of the CNN architecture, and is used as input to the reinforcement learning agent. 5 Criteria for Choosing the Compression Ratio Our goal is to optimize convolutional neural networks (CNNs) for deployment on low-capacity devices by applying structured pruning techniques. This reduces inference time, memory usage, and model size. To guide pruning, we focus on two key efficiency criteria: •Model Parameters: Reducing the number of parameters directly decreases the computational and memory overhead, enhancing the efficiency of the model [10]. •FLOPs (Floating Point Operations): FLOPs quantify the computational effort required per inference and are especially relevant in CNNs due to weight sharing [11]. Minimizing FLOPs has a direct impact on inference speed and energy consumption. FLOPs are calculated per layer type as follows: •Convolutional Layers: FLOPs = 2 ×CO×CI×K×O(3) where COis the number of output channels, CIis the number of input channels, Kis the kernel size, and Ois the number of output elements. •Fully Connected Layers: FLOPs = 2 ×I×O(4) where Iand Oare the number of input and output units, respectively. 6 Pruning Methods Selection To enable the deployment of convolutional neural networks (CNNs) on low-capacity devices, it is essential to reduce their inference time, memory footprint, and overall model size. Pruning—i.e., removing less important components of the network—is a widely used approach to achieve such compression. In this work, we evaluate both unstructured and structured pruning techniques for their effectiveness in compressing CNN architectures. •Unstructured Pruning: This method eliminates individual weights from convolutional kernels based on their magnitude or contribution. While it can result in highly sparse models with minimal impact on accuracy, it often fails to yield practical improvements in inference time or memory usage. This is largely due to the lack of hardware-level support for irregular sparsity, which limits the efficiency gains on general-purpose devices. •Structured Pruning: In contrast, structured pruning removes entire filters, channels, or even layers. This produces a more compact and regular architecture, leading to measurable reductions in computation and memory requirements. However, structured pruning carries a higher risk of accuracy degradation if critical components are pruned without adequate guidance. 62 In addition to the pruning strategy itself, the choice of how to reinitialize and retrain the pruned network plays a crucial role in recovering or maintaining performance [12]. We evaluate the following post-pruning training approaches: •Rewinding (100%): After pruning, the remaining weights are reset to their initial values recorded at the start of training. This approach is motivated by the Lottery Ticket Hypothesis, which suggests that certain subnetworks can achieve competitive performance when trained from their original initialization. •Random Initialization: As a baseline, we reinitialize the surviving weights with new random values after pruning. This serves to evaluate whether rewinding provides a significant advantage over fresh initialization. •Fine-Tuning: This method retains the final values of the remaining weights and continues training with a reduced learning rate. The goal is to refine the pruned model without substantially altering its learned representations. Fine-tuning is commonly used in practice due to its simplicity and effectiveness. 7 Implementation of Reinforcement Learning To automate the pruning process, we employ reinforcement learning (RL) with the Proximal Policy Optimization (PPO) algorithm. The RL agent is trained to predict optimal pruning ratios for each layer of a convolutional neural network, with the objective of maximizing compression while preserving classification accuracy. Experiments were conducted on the VGG-19 and ResNet-34 architectures using the CIFAR-10 dataset, with a global compression target of 80%. The agent receives as input a graph-level embedding of the CNN architecture, obtained via a Graph Convolutional Network (GCN) encoder. Based on this representation, the agent generates pruning decisions across layers. Following pruning, we apply several retraining strategies—including random initialization, weight rewinding, and fine-tuning—to restore or improve performance. This framework enables a systematic exploration of the trade-off between model compactness and accuracy, leading to efficient CNNs suitable for deployment on resource-constrained devices. 7.1 Proximal Policy Optimization (PPO) Proximal Policy Optimization (PPO) is a widely used reinforcement learning algorithm designed to improve the policy—the strategy for selecting actions—while maintaining stability and efficiency. It achieves this by carefully balancing exploration (trying new actions) and exploitation (choosing the bestknown actions) [13]. PPO optimizes a clipped surrogate objective that limits abrupt changes to the policy and incorporates additional terms to enhance learning. The total objective function consists of three main components: •Policy Surrogate Loss: This term encourages beneficial updates to the policy based on the advantage of actions taken. It uses a ratio of the new and old policy probabilities: rt(θ) = πθ(at|st) πθold(at|st)(5) To prevent large, destabilizing policy updates, PPO applies a clipping mechanism: L1(θ) = min rt(θ)ˆ At,clip(rt(θ),1−ϵ, 1 + ϵ)ˆ At(6) •Value Function Loss: This component minimizes the error between the predicted value of a state and its target return, computed using the mean squared error: L2(θ) = (Vθ(st)−Vtarg,t)2(7) •Entropy Bonus: To promote exploration and prevent premature convergence to deterministic policies, PPO adds an entropy regularization term: L3(θ) = S[πθ](st) (8) 63 The overall loss function used to update the policy parameters is a weighted sum of the three components: LTotal(θ) = ˆ Et[L1+c1L2+c2L3] (9) where c1and c2are coefficients that balance the contributions of the value loss and entropy bonus, and ˆ Et[·] denotes the empirical average over a finite batch of experiences. Exploration Noise: During exploration, the agent samples actions from a Gaussian policy. The probability density function of a Gaussian distribution is given by: f(x) = 1 σ√2πexp −1 2x−µ σ2!(10) Initially, a fixed standard deviation σcontrols the amount of randomness in action selection. Over time, this noise is gradually reduced to encourage exploitation of learned policies as training progresses. 7.2 Memory and Experience Replay To enhance learning efficiency and generalization, the PPO agent maintains a memory buffer that stores past interactions with the environment, including states, actions, action probabilities, and rewards. Rather than updating the policy network using only the most recent data, the agent samples minibatches of past experiences. This experience replay strategy reduces overfitting and provides a more diverse set of training samples, enabling the agent to learn from a broader distribution of environment interactions. 7.3 Reinforcement Learning Environment The reinforcement learning (RL) environment models the deep neural network as a computational graph in a simulated setting that reflects structural changes during pruning. The environment provides the RL agent with the graph representation based on the CNN’s topology and dynamically tracks key metrics, including the number of parameters and FLOPs. The pruning process proceeds step by step until the agent satisfies the compression constraints, at which point the search is terminated. This mimics the RL paradigm, where the environment evolves with each action and ends an episode once a terminal condition—such as reaching a pruning goal—is met. 7.4 Timestamps and Agent Interaction The RL agent interacts with the environment incrementally to prune the network towards a target compression ratio. At each time step, relevant model attributes are updated, including pruning ratios and input/output channel sizes. Upon reaching the compression goal, the pruned model is evaluated in terms of classification accuracy, and a reward is assigned. If the agent fails to meet the target FLOPs within the episode, it is penalized accordingly. To maintain meaningful compression while avoiding excessive information loss, pruning ratios for each layer are constrained within the range [0.02, 0.9]. 8 Experiments and Results Following the formulation of our approach, we conducted experiments on the VGG-19 and ResNet-34 architectures using the CIFAR-10 dataset. The implementation was carried out in Python with support from various libraries: NumPy for numerical computations, Matplotlib for visualizations, and PyTorch and PyTorch Geometric for deep learning and graph-based modeling, respectively. Torchvision was used for computer vision tasks, while Weights Biases (W&B) provided experiment tracking. Execution and collaboration were facilitated via Google Colab, Amazon EC2, and Google Drive. The experimental pipeline focused on evaluating pruning capabilities through four stages: initial model training, reinforcement learning-based pruning (targeting 80% compression), application of postpruning retraining methods (random initialization, rewinding, fine-tuning), and final evaluation of model performance across multiple metrics including accuracy, parameter count, FLOPs, and model size. 64 8.1 VGG-19 Evaluation 8.1.1 Unstructured Pruning Table 1presents the results of unstructured pruning experiments on the VGG-19 architecture using different post-pruning strategies. Table 1: VGG-19 performance after unstructured pruning under various post-pruning strategies. Method Accuracy (%) Error (%) FLOPs (%) Params (%) Size (MB) Without pruning 92.98 -7.02 100.0 100.0 548 Rewinding 93.17 -6.83 19.9 19.0 500 Random initialization 92.92 -7.08 19.9 19.0 500 Fine-tuning 92.85 -7.15 19.9 19.0 500 The application of unstructured pruning to the VGG-19 model achieved over 80% compression in terms of FLOPs and parameters. Among the retraining strategies, the rewinding method yielded the highest accuracy, slightly outperforming both random initialization and fine-tuning, and even exceeding the baseline model’s original accuracy. Moreover, it demonstrated slightly faster convergence. However, due to the nature of unstructured pruning—where individual weights are removed rather than entire structures—the overall model size remained largely unchanged. This is because the weight matrices retain their original dimensions, limiting the benefits in memory reduction. 8.1.2 Structured Pruning Table 2summarizes the performance of the VGG-19 model following structured pruning, evaluated with the same three post-pruning strategies. Table 2: VGG-19 performance after structured pruning under various post-pruning strategies. Method Accuracy (%) Error (%) FLOPs (%) Params (%) Size (MB) Without pruning 92.98 -7.02 100.0 100.0 548 Rewinding 91.59 -8.41 19.82 19.82 24.1 Random initialization 90.94 -9.06 19.82 19.82 24.1 Fine-tuning 89.82 -10.18 19.82 19.82 24.1 As shown in Table 2, structured pruning achieves a compression factor of over 22 times in model size compared to the unpruned baseline. Among the retraining techniques, rewinding consistently outperforms random initialization and fine-tuning in terms of accuracy, though it does not fully recover the original model’s performance. Rewinding also converges faster than the other methods. This performance gap is attributed to the nature of structured pruning, which removes entire filters rather than individual weights. While this simplifies network structure and significantly reduces memory footprint, it can result in a slight drop in classification accuracy. The overall trend is further illustrated in Figure 2. Figure 2: Accuracy comparison of post-pruning strategies for structured pruning on VGG-19. 65 8.1.3 Inference Time Analysis Figure 3presents a comparison of inference times across varying batch sizes for the original network, unstructured pruning, and structured pruning on the VGG-19 model. Structured pruning exhibits the lowest inference latency, followed by unstructured pruning, and finally the baseline unpruned model. The improved efficiency of structured pruning is attributed to the reduction in matrix size, leading to fewer operations. In contrast, unstructured pruning results in moderate gains due to the presence of zeroed weights, which marginally reduce computation. Figure 3: Inference time comparison of VGG-19 under different pruning strategies across batch sizes. 8.2 ResNet-34 Evaluation 8.2.1 Unstructured Pruning Table 3reports the performance of ResNet-34 following unstructured pruning using three retraining strategies. Table 3: ResNet-34 performance after unstructured pruning with various post-pruning strategies. Method Accuracy (%) Error (%) FLOPs (%) Params (%) Size (MB) Without pruning 87.20 -12.80 100.0 100.0 83.1 Rewinding 87.24 -12.76 17.46 17.75 83.1 Random Initialization 86.30 -13.70 17.46 17.75 83.1 Fine-tuning 86.61 -13.39 17.46 17.75 83.1 The unstructured pruning results for ResNet-34 closely resemble those observed for VGG-19. Among the post-pruning strategies, rewinding consistently delivers the highest accuracy, outperforming both random initialization and fine-tuning, as shown in Table 3and Figure 4. However, it should be noted that rewinding experienced convergence challenges during the initial 60,000 training iterations. Figure 4: Accuracy trends of ResNet-34 under unstructured pruning using different retraining strategies. 66 8.2.2 Structured Pruning Table 4outlines the results of structured pruning on ResNet-34, which follow trends similar to those observed in the VGG-19 experiments. Table 4: ResNet-34 performance after structured pruning with various post-pruning strategies. Method Accuracy (%) Error (%) FLOPs (%) Params (%) Size (MB) Without pruning 87.20 -12.80 100.0 100.0 83.1 Rewinding 85.94 -14.06 18.95 18.95 6.2 Random Initialization 85.12 -14.88 18.95 18.95 6.2 Fine-tuning 85.35 -14.65 18.95 18.95 6.2 As summarized in Table 4, structured pruning on ResNet-34 results in a significant reduction in model size—over 13 times smaller than the original—while maintaining competitive accuracy. Rewinding again proves to be the most effective retraining strategy, yielding the highest accuracy and fastest convergence, as illustrated in Figure 5. The observed accuracy drop is attributed to the aggressive nature of structured pruning, where entire filters are removed, simplifying the model structure without considering individual weights. Figure 5: Accuracy trends of ResNet-34 under structured pruning using different retraining strategies. 8.2.3 Inference Time Analysis To further evaluate the efficiency of pruning strategies, we compare inference times for ResNet-34 across different batch sizes. As shown in Figure 6, GPU usage did not significantly impact inference time, so CPUs were used to better highlight contrast. Structured pruning once again yielded the fastest inference, followed by unstructured pruning, and lastly the original model. This trend aligns with the observations for VGG-19. Figure 6: Inference time comparison of ResNet-34 under different pruning strategies across batch sizes. 67