Saved in:
Bibliographic Details
Main Authors: Du, Junye, Li, Zhenghao, Feng, Yushi, Feng, Long
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2605.06446
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915989646475264
author Du, Junye
Li, Zhenghao
Feng, Yushi
Feng, Long
author_facet Du, Junye
Li, Zhenghao
Feng, Yushi
Feng, Long
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_06446
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle FedFrozen: Two-Stage Federated Optimization via Attention Kernel Freezing
Du, Junye
Li, Zhenghao
Feng, Yushi
Feng, Long
Machine Learning
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.
title FedFrozen: Two-Stage Federated Optimization via Attention Kernel Freezing
topic Machine Learning
url https://arxiv.org/abs/2605.06446