[Paper Review] When does the student surpass the teacher? Federated Semi-supervised Learning with Teacher-Student EMA

Jessica Zhao, Sayan Ghosh+2 more|arXiv (Cornell University)|Jan 24, 2023

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes FedSwitch, a novel federated semi-supervised learning framework that dynamically switches between teacher and student models for pseudo-label generation using a KL-divergence-based criterion. By enabling local teacher adaptation and adaptive switching, FedSwitch achieves state-of-the-art accuracy on CIFAR-10 and Fashion MNIST with improved robustness to non-IID data and minimal communication overhead, while preserving client privacy and statelessness.

ABSTRACT

Semi-Supervised Learning (SSL) has received extensive attention in the domain of computer vision, leading to development of promising approaches such as FixMatch. In scenarios where training data is decentralized and resides on client devices, SSL must be integrated with privacy-aware training techniques such as Federated Learning. We consider the problem of federated image classification and study the performance and privacy challenges with existing federated SSL (FSSL) approaches. Firstly, we note that even state-of-the-art FSSL algorithms can trivially compromise client privacy and other real-world constraints such as client statelessness and communication cost. Secondly, we observe that it is challenging to integrate EMA (Exponential Moving Average) updates into the federated setting, which comes at a trade-off between performance and communication cost. We propose a novel approach FedSwitch, that improves privacy as well as generalization performance through Exponential Moving Average (EMA) updates. FedSwitch utilizes a federated semi-supervised teacher-student EMA framework with two features - local teacher adaptation and adaptive switching between teacher and student for pseudo-label generation. Our proposed approach outperforms the state-of-the-art on federated image classification, can be adapted to real-world constraints, and achieves good generalization performance with minimal communication cost overhead.

Motivation & Objective

• To address privacy and scalability limitations in existing federated semi-supervised learning (FSSL) methods that compromise client privacy or assume stateful clients.
• To investigate the challenges of integrating Exponential Moving Average (EMA) updates into federated settings, particularly regarding communication cost and model synchronization.
• To design a practical FSSL framework that maintains high generalization performance under real-world constraints such as data non-IIDness and low communication overhead.
• To develop a dynamic switching mechanism between teacher and student models for pseudo-label generation that adapts to client-specific data distributions.

Proposed method

• FedSwitch employs a federated teacher-student EMA framework where both global and local student and teacher models are maintained across clients and server.
• Local teacher models are adapted to client data through EMA updates, enabling them to track the student model’s distribution shifts over time.
• A KL-divergence metric is used to evaluate the divergence between the student’s and teacher’s output distributions on unlabeled data.
• Based on this metric, FedSwitch dynamically switches between using the teacher or student model to generate pseudo-labels for each client, improving label quality under non-IID conditions.
• The method ensures privacy by avoiding model sharing between clients and maintains client statelessness by not requiring persistent client-side model storage across rounds.
• Communication cost is minimized by only transmitting the student model weights to the server, while local EMA updates are performed client-side.

Experimental results

Research questions

• RQ1 How can EMA-based teacher-student training be effectively integrated into federated semi-supervised learning without incurring prohibitive communication costs or violating privacy constraints?
• RQ2 What are the performance trade-offs between centralizing teacher updates (server EMA) versus distributing them across clients (client EMA) in non-IID federated settings?
• RQ3 Can dynamic switching between teacher and student models for pseudo-label generation improve generalization under high data non-IIDness?
• RQ4 How does FedSwitch compare to state-of-the-art FSSL methods like FedMatch and FedProx-FixMatch in terms of accuracy, convergence stability, and robustness to data imbalance?
• RQ5 To what extent does local teacher adaptation enhance pseudo-label quality and model generalization in federated semi-supervised learning?

Key findings

• FedSwitch achieves state-of-the-art test accuracy on CIFAR-10 under non-IID data settings, outperforming FedProx-FixMatch, FedMatch, and FedRGD.
• The method demonstrates superior robustness to increasing levels of client data non-IIDness, maintaining high performance even when the Dirichlet α parameter is reduced to 0.1.
• FedSwitch shows a smoother convergence curve and reduced model variance compared to baseline methods, indicating more stable training dynamics.
• The KL-divergence analysis confirms that FedSwitch maintains pseudo-label distributions closer to the true client data distributions than TS-Server EMA and TS-Client EMA, especially under high non-IIDness.
• FedSwitch achieves high performance with minimal communication cost, as it only transmits student model weights to the server, avoiding the need to share teacher models or update them centrally.
• The dynamic switching mechanism enables the system to adaptively use the most reliable model (teacher or student) for pseudo-labeling, improving label quality under distribution shift.