Recurrent Early Exits for Federated Learning with Heterogeneous Clients

TL;DR

ReeFL introduces a recurrent early exit method for federated learning that efficiently combines features from sub-models using a transformer-based module, improving performance across heterogeneous clients without relying on multiple classifiers.

Contribution

The paper proposes ReeFL, a novel recurrent early exit approach that fuses sub-model features with a shared transformer-based module, eliminating the need for multiple classifiers and heuristic knowledge distillation.

Findings

ReeFL outperforms previous methods on image and speech benchmarks.

ReeFL effectively handles client heterogeneity in federated learning.

The method improves task-specific prediction accuracy.

Abstract

Federated learning (FL) has enabled distributed learning of a model across multiple clients in a privacy-preserving manner. One of the main challenges of FL is to accommodate clients with varying hardware capacities; clients have differing compute and memory requirements. To tackle this challenge, recent state-of-the-art approaches leverage the use of early exits. Nonetheless, these approaches fall short of mitigating the challenges of joint learning multiple exit classifiers, often relying on hand-picked heuristic solutions for knowledge distillation among classifiers and/or utilizing additional layers for weaker classifiers. In this work, instead of utilizing multiple classifiers, we propose a recurrent early exit approach named ReeFL that fuses features from different sub-models into a single shared classifier. Specifically, we use a transformer-based early-exit module shared among sub-models to i) better exploit multi-layer feature representations for task-specific prediction and ii) modulate the feature representation of the backbone model for subsequent predictions. We additionally present a per-client self-distillation approach where the best sub-model is automatically selected as the teacher of the other sub-models at each client. Our experiments on standard image and speech classification benchmarks across various emerging federated fine-tuning baselines demonstrate ReeFL's effectiveness over previous works.