Capability-Aware Distributed Learning and Inference for Heterogeneous 6G Edge Environments

Abstract

The rapid evolution of sixth-generation (6G) wireless networks is increasing the demand for AI-enabled edge intel ligence that can operate across devices with different capabil ities, resource constraints, and availability. Existing distributed learning paradigms, including Federated Learning (FL) and Split Learning (SL), either impose uniform computational burdens regardless of device capability or rely on sequential execution pipelines that are sensitive to node failures. To address these challenges, we propose a Capability-Aware Orchestrator (CAO) for joint distributed learning and inference in heterogeneous 6G environments. CAO profiles both node resources and model component requirements, then assigns workloads according to the capability of each node. To improve reliability, model components are assigned with overlapping coverage, allowing training to continue even when nodes become unavailable. A layer-aware aggregation mechanism is used to combine model updates and produce stable global models. Experiments on CIFAR-10 using ResNet-18 across eight heterogeneous edge nodes show that CAO achieves over 96% of the accuracy of FL and SL while reducing training time by up to 81%, communication overhead by up to 66%, and trained parameters by up to 53%. Unlike SL, which stalls when nodes fail, CAO maintains successful end-to end training under both permanent and transient node failures. In addition, the capability-aware allocation strategy reduces distributed inference latency by up to 89% compared with random deployment, demonstrating its suitability for practical 6G edge environments.