# FedFrozen: Two-Stage Federated Optimization via Attention Kernel Freezing

arXiv: https://arxiv.org/abs/2605.06446 · HTML: https://mlsanity.com/paper/2605.06446/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 26 references that are in the map.
- First posted on arXiv 7 May 2026.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [LEAF: A Benchmark for Federated Settings](https://mlsanity.com/paper/1812.01097/) (2018)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [Federated Learning with Personalization Layers](https://mlsanity.com/paper/1912.00818/) (2019)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [FedBABU: Towards Enhanced Representation for Federated Image Classification](https://mlsanity.com/paper/2106.06042/) (2021)
- [Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning](https://mlsanity.com/paper/2106.06047/) (2021)
- [Heterogeneity-Aware Memory Efficient Federated Learning via Progressive Layer Freezing](https://mlsanity.com/paper/2408.09101/) (2024)
- [Why Go Full? Elevating Federated Learning Through Partial Network Updates](https://mlsanity.com/paper/2410.11559/) (2024)

## Abstract (arXiv)

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.