FEAST: Federated Shared-Space Training for Resource-Heterogeneous Clients

FEAST: Federated Shared-Space Training for Resource-Heterogeneous Clients

Summary

This paper introduces FEAST (FEderAted Shared-space Training), a federated learning framework designed to train a single elastic supernet model that can serve clients with heterogeneous computational capabilities. The core problem addressed is that in federated supernet training, when client inference budgets differ, parameters exclusive to high-cost subnetworks are reachable only by a minority of high-budget clients, creating an optimization imbalance.

Problem Formulation

The paper formulates budget-constrained federated supernet training as a federated shared model space (FSMS), defined by one elastic parameter tensor, overlapping subnetworks, client-specific affordable sets, and budget-dependent parameter reachability. The authors note that parameters shared by low-cost subnetworks are reachable by nearly all clients, whereas parameters exclusive to high-cost subnetworks are reachable only by the high-budget minority.

FEAST Method

FEAST combines several key mechanisms:

1. Local Multi-Variant Co-Training: On each mini-batch, a client jointly trains the globally smallest, an affordable random intermediate, and its largest affordable cached variant, using the local maximum to distill the smaller variants. The local maximum is evaluated first and optimized only with supervised loss; its detached distribution supervises the global minimum and random intermediate through a combined loss: Lstu(z) = (1 - ρkd)Lsup(z) + ρkd LKD(z, p̂max) with ρkd = 0.5.

2. Budget-Tailored Sub-Supernet Routing: Budget-tailored routing transmits only the parameter envelope spanning the client's affordable variants. The server derives coordinate-wise depth and width bounds from each client's affordable set, instantiates a reduced supernet, and records a mapping from local to global positions.

3. Sparse Sample-Weighted Aggregation: Position-wise sparse, sample-count-weighted aggregation maps returned slices into the global tensor. Each position is averaged only over selected clients whose routed envelopes contain it; positions not routed retain their previous values.

γ-Allocation Protocol

The paper identifies that independently assigning clients' training-data volumes and inference budgets can distort accuracy–inference-cost comparisons in heterogeneous FL simulations. They introduce a one-parameter allocation family where qi(γ) = (balloc i) γ / Σ j (balloc j) γ, with balloc i = min(bi, Bcap). γ = 0 recovers budget-independent allocation and γ = 1 makes expected data share proportional to capped budget. The paper emphasizes this is a consequential evaluation choice, rather than a neutral implementation detail, showing that changing only this relationship raises accuracy at each method's largest evaluated point.

Key Results

1. Supernet Training Comparison: "the SuperFedNAS and DeepFedNAS supernet training procedures remain near chance at 25M and reach at most 17.09% at 596M inference MACs; FEAST reaches 71.06% at 596M, 2.4 points above the strongest model-heterogeneous weight-sharing baseline at its largest tier."

2. Population-Averaged Accuracy: Across CIFAR-100, CINIC-10, and TinyImageNet200, FEAST achieves the highest population-averaged accuracy among the evaluated weight-sharing methods when each client receives its largest affordable subnetwork. Specifically, FEAST leads the strongest weight-sharing baseline by 12.61, 26.09, and 9.44 points on CIFAR-100, CINIC-10, and TinyImageNet, respectively.

3. Communication Reduction: Sub-supernet routing reduces aggregate model-parameter traffic by 6.8× relative to full-supernet transmission. Per-client-round model-parameter reductions average 43.8× (median 22.9×, maximum 111×).

4. Training Rule Ablation: "Endpoint-only rules perform poorly away from their trained endpoints. Relative to co-training both endpoints with KD, full FEAST raises the 50M, 229M, 408M, and 596M variants by 5.17, 9.43, 8.22, and 0.97 points while remaining within 0.13 points at 25M."

5. γ Sensitivity: the CIFAR-100 largest-budget-minus-smallest-budget accuracy gap changes from −4.29 to +6.17 percentage points as γ increases from 0 to 1. The 596M result progresses from 58.96% at γ = 0 to 71.06% at γ = 1.

6. Post-Hoc Extraction: Fourteen uncached CIFAR-100 variants in the 25M–52M region have a mean absolute gap of 1.14 points from the cached curve, and the 30M variant reaches 64.35%, without further federated training.

7. Equal-Total-Training-Compute Control: When matching total training computation, FEAST's 67.74% population average remains ahead of HeteroFL (44.44%), ScaleFL (39.23%), and FIARSE (55.42%) by 23.3, 28.5, and 12.3 points with 100/100 client wins.

Theoretical Contributions

The paper establishes nested budget reachability (Proposition 1): bi ≤ bj implies Ai ⊆ Aj, Ui ⊆ Uj, ri(u) ≤ rj(u) for all u. It also provides a conditional finite-horizon stationarity bound whose dependence on the minimum routed computation across served parameters makes the cost of unequal coverage explicit (Theorem 1). The bound shows that a smaller minimum routed computation scale loosens the finite-horizon guarantee.

Conclusion

The paper concludes that "Federated supernet training offers a train-once, deploy-many route for client populations with heterogeneous deployment budgets, but our cross-setting stress test shows that the evaluated predecessor training procedures do not extend to a wider, budget-conditioned shared space. FEAST addresses the resulting challenges through local multi-variant co-training, in-place distillation, budget-tailored sub-supernet routing, and position-wise sample-count-weighted aggregation. Future work should reduce the overhead of local multi-variant co-training and test the formulation on larger architectures and additional modalities."

Improvements to AI Systems:

1. Elastic Supernet Training with Distillation-Based Co-Training: Implement FEAST's local multi-variant co-training, where each client trains the smallest, a random intermediate, and its largest affordable subnetwork simultaneously, using the largest variant's detached outputs to distill knowledge into smaller variants. This improves accuracy across all deployment budgets, especially for low-resource clients, by preventing optimization imbalance and ensuring high-cost parameters are not neglected.

2. Budget-Tailored Parameter Routing for Communication Efficiency: Adopt FEAST's sub-supernet routing, where the server transmits only the parameter envelope covering a client's affordable variants, reducing aggregate model-parameter traffic by up to 6.8× (per-client reductions averaging 43.8×). This enables federated training on bandwidth-constrained networks without sacrificing model quality.

3. Sparse Sample-Weighted Aggregation for Heterogeneous Participation: Use position-wise sparse aggregation that averages each parameter only over clients whose routed envelopes contain it, while preserving previous values for unrouted positions. This stabilizes training when clients have vastly different compute budgets and prevents low-budget clients from dominating updates to high-cost parameters.

4. γ-Allocation Protocol for Fair Evaluation and Training: Introduce a one-parameter allocation family (γ) that controls the relationship between client data volume and inference budget. Setting γ=1 makes expected data share proportional to capped budget, which corrects distorted accuracy–inference-cost comparisons in heterogeneous FL. This allows fair benchmarking and better generalization across client populations.

5. Post-Hoc Subnetwork Extraction Without Retraining: Leverage FEAST's trained elastic supernet to extract uncached subnetworks (e.g., 30M variant achieves 64.35% accuracy on CIFAR-100) with minimal accuracy loss (mean absolute gap of 1.14 points). This enables instant deployment of new model sizes for emerging devices without additional federated rounds.

6. Nested Budget Reachability for Guaranteed Coverage: Apply Proposition 1's nested reachability property to design client selection and routing algorithms that ensure all parameters are adequately trained, even when budgets are hierarchical. This improves worst-case performance for the smallest-budget clients and reduces the finite-horizon stationarity gap.

What the Improved AI System Can Do:

• Train a single elastic model that serves clients ranging from edge IoT devices (25M MACs) to high-end servers (596M MACs) with population-averaged accuracy improvements of 12.61, 26.09, and 9.44 points over prior weight-sharing baselines on CIFAR-100, CINIC-10, and TinyImageNet, respectively.

• Operate under extreme communication constraints by transmitting only budget-relevant parameter slices, cutting per-client communication by up to 111×, while maintaining state-of-the-art accuracy.

• Adapt to heterogeneous data–compute relationships via the γ-parameter, enabling accurate performance prediction and fair model selection across diverse real-world deployments (e.g., mobile users with varying battery and network speeds).

• Deploy new model variants instantly by extracting any subnetwork from the trained supernet, eliminating the need for retraining or fine-tuning when new hardware emerges.

• Provide theoretical guarantees on convergence and parameter coverage, ensuring that even the smallest clients contribute meaningfully to training and receive well-optimized models, with a conditional finite-horizon stationarity bound that quantifies the cost of unequal compute coverage.