FedAPTA: Adaptive Pruning in CPNs

• FedAPTA is a federated multi-task learning framework that employs adaptive layer-wise pruning to match device capabilities and task complexity.
• It incorporates a model recovery strategy that refills pruned parameters using a global reference model, ensuring consistent architecture for aggregation.
• It uses task-aware clustering with cosine similarity to aggregate updates from similar tasks, reducing negative transfer and improving resource efficiency.

FedAPTA (Federated Adaptive Pruning and Task-Aware Aggregation) designates a federated multi-task learning framework specifically tailored for Computing Power Networks (CPNs), which addresses the inefficiencies associated with heterogeneous task deployment and resource utilization among edge devices. Unlike traditional Federated Learning (FL) systems that apply a uniform global model to all client devices—thereby incurring substantial resource wastage on devices assigned simple tasks—FedAPTA adaptively prunes models and clusters updates to facilitate resource-efficient and task-coherent aggregation. The framework combines adaptive layer-wise pruning, model recovery, and cluster-based aggregation to optimize federated learning in scenarios characterized by device, data, and task heterogeneity (Yuan et al., 4 Aug 2025).

1. Motivation and System Overview

FedAPTA was developed in response to the limitations of classical FL within CPNs, where simplifying assumptions about task uniformity and device capability result in resource under-utilization and negative transfer. In a typical CPN, devices run diverse tasks with distinct complexities and computational budgets. Uniform deployment of large global models results in excessive computation for low-complexity tasks and diminishes overall system efficiency.

FedAPTA introduces three core innovations:

• Per-device adaptive model pruning, tuned to device capability and task complexity.
• Recovery of full model structures prior to aggregation, ensuring architectural consistency.
• Task-aware aggregation by clustering updates according to task similarity, thus limiting negative task interference (Yuan et al., 4 Aug 2025).

The system architecture consists of a central server coordinating a set of devices, each with a private dataset $D_i$ and computing capability. During each FL round:

1. The server broadcasts a global model $w$ to all devices.
2. Each device $i$ selects a pruning ratio $\rho_i$, prunes its model, trains locally, and returns the pruned model $\hat w_i$ with a binary pruning mask $M_i$.
3. The server recovers full-structured models from the pruned versions and masks, clusters devices into task groups, aggregates per-task, and redistributes updated models.

2. Adaptive Layer-wise Pruning Mechanism

FedAPTA's pruning approach is layer-wise and jointly considers device characteristics and data heterogeneity. For a local model with $n$ convolutional layers, layer $k$'s parameter count is $N_k$ and its L1-norm importance is $I_k = \|w_i^{(k)}\|_{1}$. Per-layer pruning ratios $w$0 are determined by solving the following constrained optimization:

$w$1

This guarantees total pruned parameters matches the device-level target $w$2, while less important layers receive higher pruning. Device-specific $w$3 is selected based on compute profile and task complexity: higher for simpler tasks and lower for more capable devices (Yuan et al., 4 Aug 2025).

Pseudocode for the procedure involves ranking layers by importance, solving for optimal $w$4, and zeroing the least important channels until each target fraction is achieved. The resultant mask $w$5 is used to derive the sparse model $w$6.

3. Heterogeneous Model Recovery

Because local models $w$7 may have different sparse structures after pruning, direct aggregation as in FedAvg is infeasible. FedAPTA introduces a recovery strategy to reconstruct each device model in a consistent, full-structured format. This is accomplished by "infilling" missing (zeroed) parameters from the task-global reference model $w$8:

$w$9

Retained channels in $i$0 capture device-specific fine-tuned information, while pruned channels are filled in with server-side priors from $i$1. This process ensures all recovered models $i$2 share the same architecture and are suitable for aggregation (Yuan et al., 4 Aug 2025).

4. Task-aware Model Aggregation and Clustering

FedAPTA avoids aggregating updates from heterogeneous tasks. Instead, device models are clustered according to task using a similarity-driven approach:

• Updates $i$3 for each device are compared via cosine distance on the last-layer weights:

$i$4

• The server constructs the distance matrix $i$5 and applies HDBSCAN to cluster devices into tasks (cluster number need not be predefined).

Per-cluster (per-task) aggregation is then performed using a data-size-weighted mean: $i$6 where $i$7 is the set of devices in task cluster $i$8. Only recovered, full-structured models are aggregated, and each per-task model is broadcast back to its respective devices. This process sharply contrasts with FedAvg by supporting multiple concurrent tasks and strictly limiting negative transfer (Yuan et al., 4 Aug 2025).

5. Empirical Evaluation and Results

FedAPTA was evaluated in a realistic multi-task FL setting:

• 50 heterogeneous devices, 5 classification tasks (datasets: MNIST, FashionMNIST, SVHN, CIFAR-10, EMNIST), with both i.i.d. and non-i.i.d. data partitions.
• Models: ResNet18 and ShuffleNetV2, pre-trained on ImageNet, with frozen initial layers and group normalization.
• Baselines: FedAvg, Ditto, FedProx, FedGen, MOON, FedBABU, FedNTD, FedLC, FedLPS.

Selected results (accuracy; i.i.d./non-i.i.d.):

Model	Best Baseline	FedAPTA	Improvement
ResNet18	88.80 / 83.37	89.54 / 85.87	+0.74 / +2.50 points
ShuffleNetV2	77.19 / 70.98	79.96 / 75.30	+2.77 / +4.32 points

Additional findings:

• Model size reduction (ResNet18): 11.01M to 2.20M parameters at $i$9.
• Communication cost per round drops in direct proportion to pruning sparsity.
• Across all tested pruning ratios ($\rho_i$0), FedAPTA outperformed prior pruning-based federated methods on the accuracy-resource tradeoff.

Similarity metric studies indicated that cosine distance best separated devices by task compared to alternatives (L1, L2, inner product) (Yuan et al., 4 Aug 2025).

6. Practical Considerations, Insights, and Limitations

FedAPTA's key advantages are:

• Layer-wise pruning, preserving critical filters and tuning model size to device-task pairs for maximal efficiency.
• Model recovery that allows both pruned and non-pruned parameters to contribute to the aggregated knowledge, enhancing transfer.
• Task-aware clustering that restricts aggregation to genuinely related updates, controlling negative transfer.

Trade-offs and limitations include:

• Aggressive pruning improves efficiency but increases reliance on high-quality global priors for recovery and may degrade task-specific adaptation.
• The server experiences $\rho_i$1 computational overhead for pairwise distance calculations and mask manipulations.
• Clustering is vulnerable to update signature overlap among tasks; future work may integrate supervised task IDs or learned partitioners.
• Application scope has so far been restricted to image classification. Extension to NLP, regression, and reinforcement learning tasks is noted as a future direction.
• A plausible implication is that joint optimization of device-level pruning ratios in a bi-level FL setup could further enhance efficiency and accuracy (Yuan et al., 4 Aug 2025).

FedAPTA addresses both resource waste in heterogeneous CPNs and multi-task negative transfer, achieving demonstrably superior accuracy and communication efficiency compared to leading FL methods in multi-task, heterogeneous device settings.