FedFrozen: Two-Stage Federated Optimization via Attention Kernel Freezing

TL;DR

FedFrozen introduces a two-stage federated learning method that enhances Transformer robustness in heterogeneous environments by freezing attention kernels after initial warm-up training.

Contribution

This work proposes FedFrozen, a novel two-stage federated optimization framework that selectively freezes attention components to improve stability and performance in heterogeneous federated learning.

Findings

FedFrozen improves model stability under client heterogeneity.

The warm-up stage acts as an inexact descent on a regularized kernel-profile objective.

Freezing the attention kernel after warm-up enhances federated training effectiveness.

Abstract

Federated learning with heterogeneous clients remains a significant challenge for deep learning, primarily due to client drift arising from inconsistent local updates. Existing federated optimization methods typically address this issue through objective-level regularization or update-correction mechanisms. Recent studies, however, suggest that Transformer-based architectures may be inherently more robust than conventional models under heterogeneous federated training. Motivated by this observation, we investigate how different parameter components within the attention mechanism influence federated optimization. Specifically, we decompose the attention module into a query/key block, which determines the attention kernel, and a value block, which performs semantic transformation under the induced kernel. Based on this perspective, we propose FedFrozen, a two-stage federated optimization framework that first performs full-model warm-up training and then freezes the query/key block while continuing to optimize the value block. Under a linear-attention formulation, we show that the warm-up stage can be interpreted as an inexact descent procedure on a regularized kernel-profile objective, while the frozen stage reduces to a restricted value-block optimization problem under a fixed attention kernel. Our analysis further reveals an explicit trade-off that governs the choice of warm-up length. Simulations validate the predicted bias-drift behavior, and real-data experiments demonstrate that FedFrozen improves both the stability and effectiveness of Transformer models in heterogeneous federated learning.