Source-linked AI summary
Resource-Aware Federated Mixture-of-Experts with Adaptive Pruning for Onboard Learning in LEO Satellite Constellations
Mohamed Shaaban, Mohamed Elmahallawy, Marius Bernahrndt, Tobias Hecking
TL;DR
Resource-constrained LEO satellites face onboard learning challenges from limited compute, memory, power, and intermittent communications. COSMIC-FL combines semantic expert routing, complementary MoE architectures, staged training, and structured pruning. Evaluations report an accuracy–efficiency trade-off and validate onboard execution on a Jetson AGX Orin.
Problem
Onboard FL for LEO satellites is challenging because satellites have strict compute, memory, and power limits, intermittent connectivity, and short, unpredictable ground-station contact windows.
Source: Introduction · p. 1Method
COSMIC-FL combines Sliced and Modular MoE designs, semantic class-to-expert routing, staged training, and three structured pruning strategies for resource-constrained satellite learning.
Source: Introduction · pp. 1–2Results
Evaluations report an accuracy–efficiency trade-off across trained, communicated, and retained parameters, and embedded experiments on NVIDIA Jetson AGX Orin validate onboard execution feasibility.
Source: Conclusion · p. 8Takeaways & Limitations
Structured pruning reduces memory and power consumption, while runtime improvements depend on the expert architecture.
Source: Conclusion · p. 8The current implementation prunes structured convolutional channels, leaving expert heads unpruned and contributing to retained model footprint and runtime.
Source: Federated Pruning Strategies · p. 6
Abstract
from arXiv · showhide
Low-Earth-orbit (LEO) satellites are increasingly expected to perform onboard learning for applications such as disaster response and environmental monitoring. However, conventional federated learning (FL) is ill-suited to onboard satellite learning, as it assumes computational, memory, and communication resources beyond the capabilities of resource-constrained LEO platforms, often necessitating the transmission of raw imagery to ground stations. We present COSMIC-FL, a resource-aware FL framework for efficient onboard learning in LEO satellite constellations. COSMIC-FL introduces two complementary Mixture-of-Experts (MoE) architectures: a Sliced design that shares backbone representations while activating task-specific channel subsets, and a Modular design that employs lightweight gating to route inputs to physically separated expert networks. A semantic class-to-expert mapping enables each satellite to train, update, and communicate only the expert paths relevant to its local data. To further improve efficiency, COSMIC-FL integrates staged optimization with three structured pruning strategies: server-side pruning, client-side fixed-ratio pruning with mean-vote aggregation, and adaptive client-side per-layer pruning based on aggregated importance and a MAD-based gap criterion. Combined with semantic expert routing, these techniques jointly adapt computation and model sparsity to both data semantics and layer importance, yielding a favourable accuracy--efficiency trade-off for heterogeneous space platforms. Experiments on six image classification benchmarks under highly non-i.i.d. settings show that COSMIC-FL maintains competitive accuracy while reducing communication, computation, and energy consumption by up to 80% over SOTA FL methods. We further validate COSMIC-FL on an NVIDIA Jetson AGX Orin, confirming its efficiency gains under realistic embedded deployment constraints.
Introduction
LEO satellites face limited compute, memory, power, and communications, making conventional FL workloads difficult to deploy onboard. COSMIC-FL combines expert-based learning and structured pruning to adapt model updates and computation to these constraints.
Introduction
- COSMIC-FL addresses the gap between satellite FL frameworks focused on connectivity and the compute, memory, and power limits that make conventional deep-learning workloads difficult onboard.Existing satellite-FL approaches improve convergence under intermittent links but mostly rely on monolithic models.Source: Introduction · p. 1
- COSMIC-FL combines two MoE designs with semantic class-to-expert mapping, so satellites train one relevant expert at a time and use a lightweight gate at inference.Sliced-MoE partitions channel slices across satellites; Modular-MoE uses experts specializing in subtasks or regions.Source: Introduction · p. 1
- Three structured pruning strategies remove filters or channels: centralized weight-based pruning, client-informed fixed-ratio pruning, and client-informed adaptive per-layer pruning.The strategies use global importance scores, aggregated layer-wise importance summaries, and per-layer pruning strength, respectively.Source: Introduction · pp. 1–2
- The framework stages expert specialization, gate fine-tuning, and head fine-tuning, and evaluates two MoE variants and pruning on six datasets under i.i.d. and non-i.i.d. partitions.The evaluation also runs the proposed training pipeline on an NVIDIA Jetson AGX Orin as an embedded feasibility study.Source: Introduction · p. 2
Related Work
Prior satellite-FL work focuses mainly on connectivity and orchestration, while MoE-FL targets personalization and pruning work targets terrestrial edge settings. COSMIC-FL combines sparse expert collaboration with structured compression for resource-constrained satellite constellations.
Related Work
- Satellite-FL frameworks address connectivity and orchestration but generally assume models feasible for resource-constrained satellites.MoE-FL methods target personalization under non-i.i.d. data, while pruning methods cited here target terrestrial edge settings.Source: Related Work · p. 2
- COSMIC-FL differs by jointly enabling sparse expert collaboration and structured compression for satellite constellations through federated MoE training and adaptive pruning.Source: Related Work · p. 2
Methodology
COSMIC-FL combines complementary Mixture-of-Experts designs, semantic routing, staged federated training, and structured pruning for onboard LEO learning. Satellites update expert paths relevant to their local data, while pruning compresses those paths into smaller subnetworks.
Methodology
- COSMIC-FL combines two MoE architectures, semantic class-to-expert routing, staged federated training, and three pruning strategies for onboard satellite learning.The stages cover expert specialization, gate fine-tuning, head fine-tuning, and structured pruning.Source: Methodology · p. 2
Sliced Mixture-of-Experts.
- Sliced-MoE partitions shared convolutional channels into expert-specific paths, reducing memory and communication without replicating the full model.Satellites update gradients only for the channels assigned to their expert.Source: Sliced Mixture-of-Experts · pp. 2–3
Modular Mixture-of-Experts.
- Modular-MoE uses physically separate expert towers, increasing static parameters while enabling independent expert training, communication, and pruning.The design remains compatible with standard dense training and inference operators.Source: Modular Mixture-of-Experts · p. 4
Semantic Class-to-Expert Mapping and Gate Initialization
- Class centroids from an auxiliary dataset are clustered into semantically coherent expert groups, providing a class mapping and gate initialization.The gate and shared backbone are pre-trained while expert parameters remain frozen.Source: Semantic Class-to-Expert Mapping and… · p. 4
Federated Expert Specialization via Rotation.
- During expert rotation, each satellite updates only one active expert and its head, while the backbone and gate remain fixed; rotation duration is independent of expert count.Satellites train on samples assigned by the class-to-expert mapping and upload only updated expert parameters.Source: Federated Expert Specialization via… · p. 4
Federated Gate Fine-Tuning / Federated Head Fine-Tuning.
- Gate fine-tuning updates only gate parameters with experts and heads frozen; head fine-tuning then updates only classification heads with routing and feature extraction frozen.The ground station aggregates the updated gate and head parameters using federated averaging.Source: Federated Gate Fine-Tuning · p. 4, Federated Head Fine-Tuning · p. 5
Structured (Physical) Expert Pruning
- Structured pruning physically removes complete channels and reconstructs smaller dense layers, reducing parameters, FLOPs, communication, and memory without specialized sparse kernels.Pruning follows an initial dense-training phase and preserves residual-addition shape compatibility.Source: Structured (Physical) Expert Pruning · p. 5
Federated Pruning Strategies.
- The three pruning strategies use global weights, mean-aggregated client votes, or median-aggregated importance with a MAD-based gap rule to choose filters for removal.Adaptive pruning skips layers without clear importance separation and uses minimum-retention and per-step safety constraints; pruning currently targets convolutional channels, not expert heads.Source: Federated Pruning Strategies · pp. 5–6
Performance Evaluation
COSMIC-FL is evaluated across satellite counts, data heterogeneity, pruning strategies, and embedded deployment. Results show trade-offs among accuracy, communication, computation, memory, and runtime.
Performance Evaluation
- The evaluation simulates non-overlapping data across 10, 50, or 100 clients and varies MoE expert counts from 3 to 10.Reported metrics include accuracy, parameters, FLOPs, phase-specific trainable and communicated parameters, and embedded latency, throughput, memory, power, and temperature.Source: Performance Evaluation · p. 6
Accuracy Comparison of COSMIC-FL vs. Baselines under i.i.d. Settings.
- Under i.i.d. settings, CentMoE has the highest accuracy in most cases, while COSMIC-FL reduces phase-specific training and communication costs with marginal accuracy loss.At 50 clients on RESISC45, Modular and Sliced reach 66.02% and 63.67%, versus 58.37% for HeteroFL.Source: Accuracy Comparison of COSMIC-FL vs.… · pp. 6–7, Table 1 · p. 6
Accuracy Comparison of COSMIC-FL vs. Baselines (non-i.i.d) Settings.
- On EuroSAT with 50 clients, COSMIC-FL reaches 67.26% (Modular) and 67.67% (Sliced) at α=0.1, while all methods improve as α rises to 10.FedAvg, FedProx, and CentMoE reach about 81–87% at α=0.1 and 94–95% at α=10; similar heterogeneity trends appear on RESISC45.Source: Table 2 · p. 6, Accuracy Comparison of COSMIC-FL vs.… · p. 7
Sparse Expert Participation.
- At α=0.1 with 50 clients, sparse participation cuts per-round uploads by 72.7% (Modular) and 69.3% (Sliced), while accuracy improves by 7.3% and 8.2%, respectively.Average active experts fall from 3.0 to 0.82 (Modular) and 0.92 (Sliced); communication savings reach 89.3% and 87.7% at 100 clients.Source: Sparse Expert Participation · p. 7
Effect of Structured Expert Pruning.
- All three pruning strategies substantially reduce parameters with modest accuracy loss, but the best pruning level depends on architecture.Modular-MoE remains competitive at 70–80% pruning, whereas Sliced-MoE has its best trade-off at 30–50%; adaptive pruning may retain slightly more parameters.Source: Effect of Structured Expert Pruning · p. 7, Table 3 · p. 7
Embedded Training Feasibility
- On Jetson AGX Orin, the sliced EuroSAT pipeline reaches 87.78% accuracy in 1,195.12 minutes, including 1,103.19 minutes of training.Embedded comparisons show resource trade-offs rather than universally faster models: Sliced pruning reduces memory relative to PM-MoE at an accuracy cost.Source: Embedded Training Feasibility · p. 7, Table 4 · p. 8
Communication and Inference Efficiency of COSMIC-FL vs. Baselines.
- Figure 2 compares inference cost in GFLOPs per sample across methods and three datasets, but analytical FLOPs do not necessarily predict embedded latency.Sliced execution can run slower on Jetson AGX Orin because of kernel overhead, tensor indexing, and memory behavior.Source: Communication and Inference Efficiency… · p. 7, Figure 2 · p. 8
Conclusion
The conclusion presents COSMIC-FL as a resource-aware satellite FL framework combining staged expert training, complementary MoE designs, semantic routing, and structured pruning. Its evaluations show accuracy–efficiency trade-offs and feasible onboard execution, while runtime gains depend on architecture.
Conclusion
- COSMIC-FL combines staged expert training, complementary MoE architectures, semantic class-to-expert routing, and federated structured pruning for LEO constellations.Embedded experiments on NVIDIA Jetson AGX Orin validate onboard execution feasibility.Source: Conclusion · p. 8
- Evaluations show trade-offs across trained, communicated, and retained parameters; structured pruning reduces memory and power, while runtime improvements depend on expert architecture.FLOPs, parameter count, and embedded runtime provide complementary efficiency perspectives.Source: Conclusion · p. 8
- Future work will investigate operator-level optimization of sliced execution and adaptive expert scheduling under dynamic satellite connectivity.These are identified as future directions for the framework.Source: Conclusion · p. 8
Dataset & Architecture
The evaluation uses six image-classification benchmarks with simulated satellite clients holding non-overlapping data. The architecture pairs a shared stem and expert paths with gating and expert-specific heads.
Datasets
- The framework is evaluated under federated learning on five standard datasets—AID, EuroSAT, UCMerced, CIFAR-10, and RESISC45—with CIFAR-10 included for comparison to general federated-learning and MoE methods.Source: Datasets · p. 11
- Each simulated satellite holds a non-overlapping data portion; client counts are 10, 50, or 100, and MoE expert counts are 3, 5, 7, or 10.Each dataset uses a 70%–15%–15% train/validation/test split, with 5% of training data reserved for server-side gate pretraining.Source: Datasets · p. 11
Architecture
- The architecture uses a shared convolutional stem, expert convolutional paths, a gating network, and a dedicated classification head for each expert.Semantic class-to-path assignment uses frozen ImageNet-pretrained ResNet18 features, per-class means, and k-means with k equal to the number of experts.Source: Architecture · p. 11
Hyperparameters
The experiments use staged COSMIC-FL training and specified non-i.i.d. heterogeneity settings. Server-side simulations run on a workstation, while embedded measurements use Jetson AGX Orin.
Hyperparameters
- Unless otherwise stated, COSMIC-FL uses 50 expert-specialization rounds, 20 gate-finetuning rounds, and 20 head-finetuning rounds.Non-i.i.d. heterogeneity uses Dirichlet α values of 0.1, 0.3, 1.0, or 10; the sparse-participation study instead uses 15 specialization rounds.Source: Hyperparameters · p. 11
Simulation Environment
- Server-side experiments use a workstation with an NVIDIA GeForce RTX 4090, while embedded measurements use NVIDIA Jetson AGX Orin.Server-side results are single runs with fixed seed 42; embedded measurements average two independent runs.Source: Simulation Environment · p. 11
Pruning Protocol
The pruning ablation starts from dense baselines, then applies each pruning strategy over four pruning rounds followed by six recovery rounds.
Pruning Protocol
- Each pruning run starts from the corresponding dense baseline trained for 10 additional epochs, which serves as the “No Pruning” reference.Runs initialize from the dense baseline obtained in the preceding experiment.Source: Pruning Protocol · p. 11
- Each pruning strategy uses four pruning rounds followed by six recovery rounds.Source: Pruning Protocol · p. 11
- The per-round pruning rate is set to reach a desired final pruning ratio of 10%, 30%, 50%, 70%, or 80%.Source: Pruning Protocol · p. 11
MAD Threshold Sensitivity
Adaptive pruning uses a MAD-based gap criterion, so requested pruning targets are compression objectives rather than guaranteed achieved ratios.
MAD Threshold Sensitivity
- The adaptive method prunes a layer only when its sorted importance scores show a sufficiently pronounced gap, so achieved ratios can differ from requested targets.Source: MAD Threshold Sensitivity · p. 11
- Low and moderate MAD thresholds often produce identical pruning decisions because detected score gaps are sufficiently pronounced.Source: MAD Threshold Sensitivity · p. 11
- Higher MAD thresholds make adaptive pruning more conservative, reducing achieved pruning ratios; the requested target is an upper-level compression objective.Source: MAD Threshold Sensitivity · p. 13
Baseline descriptions
The baselines span standard federated averaging, proximal optimization, mixture-of-experts, personalization, and sub-model training approaches.
Baseline descriptions
- FedAvg averages client updates, while FedProx adds a proximal term to limit client drift.Source: Baseline descriptions · p. 13
- CentMoE trains a central mixture-of-experts model with federated aggregation, while PM-MOE combines personalized heads before fine-tuning a mixture of experts.Source: Baseline descriptions · p. 13
- FedDropout trains randomly selected sub-models, whereas HeteroFL assigns width-scaled sub-models and aggregates overlapping parameters.For both methods, r is the fraction of the full model trained per client; r=1.0 recovers full-model training.Source: Baseline descriptions · p. 13
More Results
Across the reported studies, COSMIC-FL’s results depend on expert count, training stage, participation, heterogeneity, communication design, and pruning strategy.
Effect of Number of Experts
- For Modular, E=3 yields the best accuracy across the studied datasets; on CIFAR-100, accuracy falls from 29.1% to 20.0% as E grows from 3 to 10.The expert count remains decoupled from class count: E=3 remains optimal across 10 to 100 classes, and doubling rotation rounds does not reverse the ranking.Source: Effect of Number of Experts · p. 13
Accuracy over Expert-Specialization Rounds
- Across EuroSAT, AID, and RESISC45, CentMoE and single-tower baselines converge fastest and to the highest accuracy during the 50-round expert-specialization phase.PM-MoE and COSMIC-FL start lower and converge more slowly, with COSMIC-FL trading some accuracy for personalization and efficiency properties.Source: Accuracy over Expert-Specialization… · pp. 13–14
Contribution of Each Training Stage
- Expert specialization adds the largest stage-wise gain (+45 to +68 points), while gate fine-tuning adds +3.0–4.3 points and head fine-tuning adds most under strong heterogeneity (+10.6 at α=0.1).These are sequential gains under a fixed schedule, not a leave-one-out ablation; the same qualitative pattern appears for Modular at α=0.1.Source: Contribution of Each Training Stage · p. 14
Sparse Expert Participation across Constellation Sizes
- Sparse participation’s communication savings grow from 40.0% to 89.3% for Modular and 43.3% to 87.7% for Sliced as client count rises from 10 to 100, while validation accuracy remains at or above full training.At 50 clients under α=0.1, active experts per client fall to 0.82 for Modular and 0.92 for Sliced, with uploads reduced by 72.7% and 69.3%.Source: Sparse Expert Participation across… · p. 14, Table 9 · p. 14
Data Heterogeneity on EuroSAT and RESISC45
- On EuroSAT and RESISC45, all methods improve with larger α; PM-MoE collapses at α=0.1, while COSMIC-FL tracks full-model baseline behavior at a fraction of staged training and communication cost.Table 10 reports the heterogeneity study on both datasets, with EuroSAT shown in the main paper.Source: Data Heterogeneity on EuroSAT and… · p. 14, Table 10 · p. 15
Per-Round Communication Across Training Phases
- Across EuroSAT, AID, and RESISC45, Sliced-MoE transfers 4.9M parameters per round, 34% less than CentMoE and 46% less than Modular-MoE.During expert specialization, per-round parameters are 6.6M for Modular and 2.5M for Sliced versus approximately 7.4–7.5M for CentMoE; gate fine-tuning adds about 0.02M.Source: Per-Round Communication Across… · p. 14, Figure 4 · p. 15
Total Uplink Communication across Training Phases
- At 100 clients, Modular-sparse and Sliced-sparse reduce total uplink by 47% and 68% versus FedAvg, respectively, although full staged training sends 3.6× and 1.4× FedAvg.Table 11 reports total uplink across all phases under the reference schedule; sparse gains grow with constellation size.Source: Table 11 · p. 15, Total Uplink Communication across… · p. 16
Sub-Model Baselines under Varying Capacity Fractions
- FedDropout accuracy degrades sharply at r=0.25 (5–33%), while HeteroFL degrades more gracefully; both improve monotonically with r and r=1.0 closely matches FedAvg.The study sweeps r ∈ {0.25, 0.5, 0.75, 1.0} under IID partitioning.Source: Sub-Model Baselines under Varying… · p. 16, Table 12 · p. 16
Additional Jetson Hardware Results
Additional Jetson AGX Orin inference results compare dense and adaptively pruned checkpoints, showing memory and power reductions but implementation-dependent latency effects. The measurements are checkpoint-level benchmarks, not a simulation of a federated training round.
Additional Jetson Hardware Results
- Final dense and adaptively pruned checkpoints were benchmarked independently on one Jetson AGX Orin, without loading all client models or simulating a federated training round.Measurements used batch size one, 50 warm-up batches, 500 timed batches, and two independent runs per checkpoint; reported measures include latency, throughput, peak allocated CUDA memory, power, and temperature.Source: Embedded inference protocol · p. 18, Additional Jetson Hardware Results · p. 18
Embedded inference protocol
- 59.75% on EuroSAT and 57.74% on RESISC45 were the peak allocated CUDA memory reductions for sliced models; board power fell about 11%, while latency improved 7–8% for sliced models but rose 5–6% for modular models.These are relative effects of adaptive pruning versus the corresponding dense implementations.Source: Embedded inference protocol · p. 18
- On EuroSAT, modular and sliced baselines used 83.45 MB and 105.05 MB of peak allocated CUDA memory; on RESISC45, they used 84.78 MB and 105.26 MB.The sliced baseline used more peak memory despite its smaller static parameter count; runtime footprint can also depend on activations, temporary tensors, indexing, and CUDA workspaces.Source: Embedded inference protocol · p. 18
- The CUDA-event diagnostic left substantial sliced forward time outside coarse instrumented module blocks, so it cannot isolate a single causal operation without lower-level operator tracing.The observation suggests possible overhead from slicing, indexing, routing, tensor reorganization, or temporary allocation.Source: Embedded inference protocol · p. 18, Additional Jetson Hardware Results · p. 18