Divergence-Based Adaptive Aggregation for Byzantine Robust Federated Learning

Extracted source text · arXiv 2601.06903

% \documentclass[lettersize,journal]{IEEEtran}
% \documentclass[draftclsnofoot,onecolumn]{IEEEtran}

% \usepackage{amsfonts}       % blackboard math symbols
% \usepackage{nicefrac}       % compact symbols for 1/2, etc.
% \usepackage{graphicx}
% \usepackage{microtype}      % microtypography
% \usepackage{xcolor}         % colors
% \usepackage{amsmath}
% \usepackage{amssymb}
% \usepackage{amsfonts}
% \usepackage{amsthm}
% \usepackage{verbatim}
% \usepackage{algorithmic}
% \usepackage{algorithm}
% \usepackage{subfigure}
% \pagenumbering{roman}
% % \usepackage{authblk}

% \usepackage{footnote}

% \allowdisplaybreaks[4]


% \documentclass[journal]{IEEEtran}
\documentclass[lettersize,journal]{IEEEtran} 
% \documentclass[draftclsnofoot,onecolumn]{IEEEtran}
\usepackage{amsmath,amsfonts}
\usepackage{bm}
\usepackage{algorithmic}
\usepackage{algorithm}
\usepackage{array}
\usepackage{textcomp}   
\usepackage{stfloats}
\usepackage{url}
\usepackage{verbatim}
\usepackage{graphicx}
\usepackage{cite}
\usepackage{xcolor}
% \usepackage{subfigure}
\newtheorem{remark}{Remark}
\usepackage{amsthm}
\newtheorem{theorem}{Theorem}
\newtheorem{definition}{Definition}   
\newcommand*{\red}{\color{red}}
\newcommand*{\blue}{\color{blue}}
% \usepackage[caption=false,font=normalsize,labelfont=sf,textfont=sf]{subfig}
\usepackage[caption=false, font=footnotesize]{subfig}

\usepackage[normalem]{ulem}

\usepackage{empheq}  


\allowdisplaybreaks[4]

\hyphenation{op-tical net-works semi-conduc-tor IEEE-Xplore}

\usepackage{xpatch}


\title{
% DRAG and BR-DRAG: 
Divergence-Based Adaptive Aggregation for Byzantine Robust Federated Learning
}



\begin{document}

\author{Bingnan Xiao,
        Feng~Zhu,
        Jingjing Zhang,~\IEEEmembership{Member,~IEEE}, 
        Wei~Ni,~\IEEEmembership{Fellow,~IEEE},
        and
        Xin~Wang,~\IEEEmembership{Fellow,~IEEE} 
\thanks{B. Xiao, J. Zhang, and X. Wang are with the Key Laboratory of EMW Information (MoE), College of Future Information Technology, Fudan University, Shanghai 200433, China (e-mail: 22110720061@m.fudan.edu.cn, \{jingjingzhang, xwang11\}@fudan.edu.cn).

F. Zhu is with the Department of Electrical and Computer Engineering, North Carolina State University, Raleigh, NC, USA. (email: fzhu5@ncsu.edu).

W. Ni is with the School of Engineering, Edith Cowan University, Perth, WA 6027, and the School of Computing Science and Engineering, the University of New South Wales, Sydney, NSW 2052, Australia (e-mail: wei.ni@ieee.org).
}

}


\maketitle

\begin{abstract}

Inherent client drifts caused by data heterogeneity, as well as vulnerability to Byzantine attacks within the system, hinder effective model training and convergence in federated learning (FL).
This paper presents two new frameworks, named DiveRgence-based Adaptive aGgregation (DRAG) and Byzantine-Resilient DRAG (BR-DRAG), to mitigate client drifts and resist attacks while expediting training. DRAG designs a reference direction and a metric named divergence of degree to quantify the deviation of local updates. Accordingly, each worker can align its local update via linear calibration without extra communication cost. 
BR-DRAG refines DRAG under Byzantine attacks by maintaining a vetted root dataset at the server to produce trusted reference directions. The workers' updates can be then calibrated to mitigate divergence caused by malicious attacks.
We analytically prove that DRAG and BR-DRAG achieve fast convergence for non-convex models under partial worker participation, data heterogeneity, and Byzantine attacks. 
Experiments validate the effectiveness of DRAG and its superior performance over state-of-the-art methods in handling client drifts, and highlight the robustness of BR-DRAG in maintaining resilience against data heterogeneity and diverse Byzantine attacks.





\end{abstract}


\begin{IEEEkeywords}
    Federated learning, client drift, Byzantine attack, convergence analysis
\end{IEEEkeywords}


\section{Introduction}
% As the complexity of machine learning tasks increasingly grows and the volume of data also multiplies at an exponential rate, distributed implementation of working nodes, i.e., federated learning (FL), has gained rising attention \cite{dean2012large,smith2017federated,lim2020federated,yang2019federated}. Parameter-server (PS) setting is one of the most widely employed paradigms in FL, where the PS broadcasts the latest global model to the workers for computation and the workers send their computed local models back to the PS for aggregation and update \cite{lian2017can,li2014scaling,gupta2016model}. By developing a novel method named divergence-based adaptive aggregation (DRAG), this paper aims to deal with the following two issues. In this work, we address the challenges posed by the growing complexity and scale of machine learning tasks, with a specific focus on FL involving heterogeneous data. Our primary aim is to overcome the following two critical challenges:

With the growing complexity of machine learning (ML) tasks and exponential increase in data volumes, federated learning (FL) has received growing attention~\cite{pmlr-v54-mcmahan17a,9069945}. 
% Based on a parameter server (PS)-worker setup, the geographically distributed workers train a global deep neural network (DNN) model with local data, and only upload the local updates instead of raw data to the PS, for aggregation~\cite{9446488}.
% % Among them, the parameter-server (PS) architecture is one of the most widely adopted. In this setup, the server broadcasts the current global model to a set of workers (agents), each of which performs local computations and sends updated local models back to the server for aggregation and update~\cite{lian2017can,li2014scaling,gupta2016model}.
% To mitigate the communication bottlenecks when the models have high dimensions~\cite{zhang2015deep,alistarh2017qsgd}, the typical Federated Averaging (FedAvg)~\cite{pmlr-v54-mcmahan17a} enables each worker to perform multiple local stochastic gradient descent (SGD) before synchronizing with the PS, facilitating model training with user privacy and communication efficiency.
% 
% % Yet, a major bottleneck in large-scale FL is communication overhead~\cite{zhang2015deep,alistarh2017qsgd}. When model dimensions are large, it becomes impractical for agents to communicate with the server at every iteration. To address this, the widely-used Federated Averaging (FedAvg) algorithm was proposed in~\cite{pmlr-v54-mcmahan17a}. 
% % In FedAvg, each agent performs a fixed number of \textit{local stochastic gradient descent (local SGD)} updates before synchronizing with the server. During communication rounds, the server aggregates the latest local models and broadcasts the updated global model to all agents for the next local training phase.
% 
Meanwhile, FL with local stochastic gradient descent (SGD) introduces a series of challenges different from centralized learning.

One challenge is client drift due to the lack of persistent communication between workers and the parameter server (PS)~\cite{zhao2018federated}. 
% Since the imbalanced distribution of locally generated data across workers is characterized by data heterogeneity, i.e., non-independent and identically distributed (non-IID) data, their local updates without timely synchronization cause the models to diverge toward distinct local optima
Since the local data of workers can be imbalanced and heterogeneous, i.e., non-independent and identically distributed (non-IID), untimely synchronization of their local updates can cause their models to drift apart~\cite{zhao2018federated}. 
Naively averaging the models, e.g., using Federated Averaging (FedAvg)~\cite{pmlr-v54-mcmahan17a}, fails to mitigate such drifts, thwarting convergence.
% The local SGD paradigm in federated learning (FL) introduces a set of challenges not present in centralized learning systems. A prominent issue is the phenomenon of \emph{client drift}, which arises due to the lack of frequent communication between clients and the central server. In particular, when clients maintain heterogeneous data distributions, performing multiple local updates without synchronization causes their models to diverge toward distinct local optima. Consequently, naively averaging the local models at the server fails to fully mitigate this drift, thereby hindering convergence, as observed in~\cite{zhao2018federated}.

Another challenge in FL stems from susceptibility to attacks launched by malicious workers due to its distributed nature. 
These attacks, commonly known as Byzantine attacks~\cite{so2020byzantine,10574838}, disrupt the global model training by tampering with the local data or model updates of the attacked workers, which seriously deteriorate training performance and system security.


% These attacks, widely known as \textit{Byzantine attacks}~\cite{so2020byzantine,cao2020fltrust,prakash2020mitigating}, can potentially compromise the integrity of the learning process and the security of the overall system. 
% It is essential to develop defense mechanisms to counter such attacks and ensure the robustness and reliability of FL in the presence of adversarial clients.

% Another challenge in FL stems from susceptibility to attacks initiated by malicious clients due to its distributed nature. These attacks, widely known as \textit{Byzantine attacks}~\cite{so2020byzantine,cao2020fltrust,prakash2020mitigating}, can potentially compromise the integrity of the learning process and the security of the overall system. It is essential to develop defense mechanisms to counter such attacks and ensure the robustness and reliability of FL in the presence of adversarial clients.



%Under heterogeneous data distribution, FedAvg achieves remarkable performance compared to parallel SGD which does not employ local updates, which is also theoretically substantiated by rigorous convergence analysis \cite{zhou2018convergence,woodworth2018graph,wang2019cooperative,khaled2020tighter}. 


%To overcome this challenge, the local stochastic gradient descent (SGD) method~\cite{zinkevich2010parallelized} has been introduced, where each worker runs \textit{local updates} from the same initialized model, and the server performs a \textit{one-shot aggregation} at the end of the training process as the final output. Nevertheless, this method is proved to perform worse than distributed SGD under worst-case scenarios, particularly when dealing with non-convex or non-smooth objective functions \cite{arjevani2015communication,zhang2012communication}. In response to this limitation, 
%Due to the frequent bidirectional transmissions between the PS and the workers, communication efficiency has become a critical bottleneck in large-scale FL \cite{zhang2015deep, alistarh2017qsgd}.  To overcome this challenge, the local stochastic gradient descent (SGD) method \cite{zinkevich2010parallelized,pmlr-v54-mcmahan17a} has been introduced. This method enables each worker to perform multiple \textit{local updates} before uploading the latest model, significantly enhancing communication efficiency.
% The authors of \cite{zinkevich2010parallelized} first put forth the idea of local SGD to improve communication efficiency. Specifically, each client optimizes its own local model and the PS aggregates the latest local models as the final output. Nevertheless, this method is proved to perform worse than distributed SGD under worst-case scenarios, particularly when dealing with non-convex or non-smooth objective functions \cite{arjevani2015communication,zhang2012communication}. In response to this limitation, the most widely-studied federated averaging (FedAvg) algorithm is proposed in~\cite{pmlr-v54-mcmahan17a}, where the communication frequency between the PS and the workers is increased, leading to improved performance. 




% This paper introduces a novel method called divergence-based adaptive aggregation (DRAG) to address the challenges presented by the growing complexity and scale of machine learning tasks, with a particular focus on FL involving heterogeneous data. The primary aim of DRAG is to overcome the following two critical challenges:
% A more detailed discussion of the related works regarding the client drift issue and Byzantine attacks are then in order.
\subsection{Related Work}

Client drift in FL was first reported in~\cite{zhao2018federated}, where data heterogeneity can cause local model updates to deviate systematically from the true global gradient direction, thereby introducing bias and impeding convergence.
Client drift was further analyzed in~\cite{li2019convergence,malinovskiy2020local}.
% % The phenomenon of client drift in FL was first identified in~\cite{zhao2018federated}, which was proved to introduce a heterogeneity-induced bias and hinder convergence, as observed in~\cite{li2019convergence,malinovskiy2020local,charles2020outsized}. 
% Recent studies have devoted emerging attention to this issue. 
Some works incorporated variance reduction techniques~\cite{defazio2014saga} into local SGD, e.g.,~\cite{liang2019variance,li2019feddane,pathak2020fedsplit,mitra2021linear}. Yet, these methods require full worker participation, limiting their practicality when only partial workers are active.

Another mainstream strategy is to leverage control variates or explicit gradient constraints for update correction. In~\cite{karimireddy2020scaffold}, SCAFFOLD was designed with both local and global control variates to mitigate drifts at the workers. 
Similarly, FedDyn~\cite{acar2021federated} adopts first-order regularization terms on both the worker and PS sides.
MIME~\cite{karimireddy2020mime} utilizes momentum-based gradients to enhance the server-level optimization. As a method based on gradient constraints, FedProx was proposed in~\cite{li2020federated}, which introduced a proximal term into the local objective function to restrict the deviation of local updates from the global model. Built on FedProx, FedDC~\cite{gao2022feddc} decouples the local and global models, and integrates both control variates and regularization terms to jointly mitigate the drift between them.
FedSAM~\cite{pmlr-v162-qu22a} deploys the SAM optimizer~\cite{NEURIPS2023_e095c0a3}, which prompts workers to approximate local flat minima to enhance model generalization.
Unfortunately, these approaches incur additional complexity and memory overhead due to the stored control variates and incorporated regularization terms.
% , increasing the computational burden and limiting their scalability and deployment in resource-constrained FL systems.
% % However, despite their effectiveness, these methods often require additional memory to store control variates and may introduce hyperparameter tuning challenges, limiting their scalability and ease of deployment in practice.
% % However, these approaches often introduce additional optimization complexity and may require careful tuning of regularization parameters to maintain stability and convergence. Furthermore, the added computational burden on clients can pose practical limitations, especially in resource-constrained environments.

Recently, strategies, such as FedExP~\cite{jhunjhunwala2023fedexp} and FedACG~\cite{Kim_2024_CVPR}, have aimed to improve training convergence by refining the aggregation process. FedExP dynamically adjusts the stepsize at the PS side with extrapolation on pseudo-gradients. FedACG stabilizes training by broadcasting a lookahead global gradient to align the local updates.
However, these methods rely on accurate estimation of global or pseudo-gradients, which can be noisy or biased in heterogeneous settings. 
% Moreover, the use of extrapolation or lookahead mechanisms may introduce instability or additional communication overhead, making them less robust or efficient under practical constraints such as limited bandwidth or asynchronous participation.

% First discovered in \cite{zhao2018federated}, the authors identify the phenomenon of ``client-drift'' in local SGD, namely that the local models of different workers will converge to local optima under heterogeneously distributed data, and the simple direct averaging of the local models will cause terrible convergence results.

%Byzantine attack, an intractable threat in conventional distributed network, is discovered to be rather efficacious against FL as well. As FL operates in a distributed manner, it is susceptible to adversarial attacks launched by malicious clients, commonly known as Byzantine attacks \cite{so2020byzantine,cao2020fltrust,bagdasaryan2020backdoor}. It is essential to develop defense mechanisms to counter such attacks and ensure the robustness and reliability of federated learning in the presence of adversarial clients. 

% \newpage
On the other hand, FL with local SGD is vulnerable to Byzantine attacks~\cite{9947081}. Common Byzantine attacks include:
\begin{itemize}
    \item Noise Injection~\cite{9614992}: The malicious workers upload random vectors (e.g., sampled from a Gaussian distribution) instead of genuine updates, which corrupt the global model with unpredictable perturbations during training.
    
    \item Sign Flipping~\cite{li2019rsa}: For the malicious workers, the sign of their local updates is flipped, leading to gradient ascent rather than descent. This actively maximizes the loss and drives the global model away from convergence.
    
    \item Label Flipping~\cite{10054157}: This attack erodes local data by reversing each data label $l$ to $L - l - 1$, where $L$ is the total number of classes, producing semantically misleading updates while remaining stealthy and API-independent.
\end{itemize}
% \begin{itemize}
%     \item Noise Injection: A simple yet disruptive attack where malicious clients upload random vectors (e.g., sampled from a Gaussian distribution) instead of genuine updates. These injected noises corrupt the aggregated model by introducing unpredictable perturbations during training.
%     \item Sign Flipping: In this attack~\cite{prakash2020mitigating}, each adversarial client flips the sign of its gradient vector, effectively performing gradient ascent rather than descent. This actively maximizes the loss and drives the global model away from convergence.
%     \item Label Flipping: Unlike NI and SF, which manipulate gradients directly, label flipping~\cite{fang2020local} corrupts the training data itself. The attacker reverses each label $l$ to $L - l - 1$, where $L$ is the total number of classes, producing semantically misleading updates while remaining stealthy and API-independent.
% \end{itemize}
% Other notable attacks include Krum~\cite{blanchard2017machine} and Trimmed Mean~\cite{fang2020local}, which are often used in tandem with their respective robust aggregation rules to evade detection and suppress benign updates.

Recent works have designed Byzantine-robust aggregation strategies to remove outliers among the workers, i.e., comparing all local updates and suppressing the anomalies~\cite{blanchard2017machine,fang2020local,yin2018byzantine}.
With techniques, e.g., Krum~\cite{blanchard2017machine} and Trimmed Mean~\cite{fang2020local}, FL systems can withstand malicious attacks to some extent. In \cite{li2019rsa}, RSA was designed to penalize the difference between local and global model parameters to defend against Byzantine workers.
These methods lack robustness against Byzantine attacks when malicious workers become widespread. 

In \cite{cao2020fltrust}, FLTrust was proposed, where the PS leverages a root dataset to guide the global update direction. In \cite{9721118}, the geometric median~\cite{minsker2015geometric} was adopted for FedAvg due to its strong tolerance against a high proportion of malicious workers. In~\cite{9153949}, Byrd-SAGA was developed, which augments the SAGA method~\cite{defazio2014saga} with the geometric median to achieve Byzantine robustness. In~\cite{10100920}, BROADCAST was proposed, which integrates gradient difference compression with SAGA to mitigate compression and stochastic noise, enhancing resilience to Byzantine attacks.
However, these algorithms lack convergence guarantees in Byzantine FL systems with heterogeneous data, particularly for non-convex learning models. Recently, in \cite{zuo2024byzantine}, a convergence guarantee was established for the geometric median under data heterogeneity, albeit relying on the debatable assumption of bounded gradients.
% For instance, the Krum aggregation rule \cite{blanchard2017machine} uses the one local update that has the smallest Euclidean distance to the rest of the clients to update the global model. 







% To be specific, these two attack modes try to solve an optimization problem for particular aggregation rules (Krum or Trim) that aims to change the correct global update direction the most through poisoning the local model updates of the malicious clients. The malicious client that initiates a Scaling attack trains a backdoor-poisoned dataset and scales the byzantine local model update by multiplying it with a factor much larger than 1. The reason for doing that is to magnify the influence of the poisoned local update and prevent it from being diluted via averaging.
% One of the approaches to defending against Byzantine attacks is to remove the outliers in the clients through utilizing Byzantine-robust aggregation rules, i.e., comparing the local updates of the clients and getting rid of the anomalies \cite{blanchard2017machine,fang2020local,yin2018byzantine}. 




\vspace{-3 mm}
% \subsection{Our Contributions}
\subsection{Contribution} 
This paper proposes two new aggregation frameworks, namely DiveRgence-based Adaptive aGgregation (DRAG) and Byzantine-Resilient DRAG (BR-DRAG), to tackle data heterogeneity and malicious Byzantine attacks, and accelerate model training in FL. The key contributions include:



% To address the aforementioned challenges, we propose a novel adaptive aggregation framework termed \emph{Divergence-based Adaptive Aggregation} (DRAG), in which the server dynamically adjusts each client's update based on a \emph{reference direction}. To the best of our knowledge, this is the first method that simultaneously achieves strong performance in mitigating client drift and resisting Byzantine attacks. 
\begin{itemize}
\item 
We interpret client drift as gradient misalignment and define the \textit{degree of divergence} (DoD) to quantify the misalignment between each worker’s update and a momentum-based global reference. 
We propose DRAG to steer local updates toward the reference direction to enhance the alignment in a decentralized manner.

\item 
We strengthen the resistance of DRAG against Byzantine attacks by producing trustworthy reference directions based on a small and vetted root dataset. By aligning and scaling the local updates with the reference direction, this BR-DRAG algorithm mitigates the adverse effect of malicious gradients while preserving their meaningful contribution to aggregations.
% retaining their contributions to aggregation.

% We develop BR-DRAG, which anchors the reference direction to a small, vetted PS-side root set to preserve the alignment principle under adversarial workers. By combining the reference direction with a redesigned Byzantine-stable normalization operator, BR-DRAG effectively tolerates malicious gradients and utilizes them for global updates.

\item 
We analyze the convergence of DRAG and BR-DRAG in the non-convex case, accounting for data heterogeneity, partial worker participation, and Byzantine attacks. 
% It is revealed that under both benign and adversarial settings, the proposed algorithms can achieve the convergence rate that the state-of-the-art algorithms achieve under no attacks. We also prove that BR-DRAG maintains convergence while relaxing the existing assumption on the proportion of malicious workers.
We reveal that, under both benign and adversarial settings, our algorithms attain the convergence rate of state-of-the-art methods operating without attacks. Also, BR-DRAG sustains convergence while relaxing the conventional assumption on the proportion of malicious workers.
\end{itemize}


% \begin{itemize}
%     \item We develop a new algorithm to mitigate client drifts in FL systems, which defines a new metric named degree of divergence to quantify the angular deviation between each worker's local update and a designed global reference direction. The local updates are adaptively ``dragged'' toward the reference direction via vector interpolation to alleviate data heterogeneity.
     
%     % introduce the \textbf{DRAG algorithm}, a novel adaptive aggregation framework that defines a new metric, the \emph{degree of divergence}, which quantifies the angular deviation between each client's local update and a global reference direction. The reference direction is constructed in a momentum-like fashion as a weighted sum of historical global updates. Using this metric, each client's update is adaptively ``dragged'' toward the reference direction via vector interpolation before being aggregated at the server.
    
%     \item We expand a Byzantine-resilient variant, BR-DRAG, to handle the potential adversarial workers. The reference direction and gradient modifications are redesigned to preserve the training robustness under attack.
    
%     \item For both DRAG and BR-DRAG, we provide a rigorous convergence analysis in the non-convex case, accounting for the adverse effects from data heterogeneity, partial worker participation, and Byzantine attacks. Our analysis shows that under both benign and adversarial settings, the proposed algorithms achieve superior convergence, validating their effectiveness in client-drift mitigation and Byzantine resilience.
% \end{itemize}


We evaluate the performance of DRAG and BR-DRAG on three broadly adopted datasets: EMNIST, CIFAR-10, and CIFAR-100. 
% It is demonstrated that 
DRAG effectively mitigates client drifts and accelerates convergence, and consistently outperforms state-of-the-art baselines under heterogeneity and partial worker participation. 
Under various Byzantine attacks, BR-DRAG shows its superb resilience to severely skewed data distributions and high proportions of malicious workers, compared to existing defense methods, e.g., FLTrust and geometric median. 
% We also explore the impact of the hyperparameters on the convergence of DRAG and BR-DRAG.



% {\color{red}
% \begin{itemize}
%     \item We introduce a new metric called ``degree of divergence" that utilizes the angle to quantify the deviation between workers' local updates and the weighted sum of historical global updates. This metric forms the basis for our proposed method, divergence-based adaptive aggregation (DRAG), effectively addressing the client-drift phenomenon and accelerating the convergence rate.  
%     \item We establish a rigorous convergence analysis of the DRAG algorithm, demonstrating that it achieves a sublinear convergence rate similar to other local SGD methods.
%     \item The performance of DRAG is evaluated on EMNIST and CIFAR-10 datasets, and we conduct a comprehensive comparison with state-of-the-art algorithms, revealing that our method consistently achieves superior results.
%     \item Finally, we demonstrate that DRAG exhibits resilience against certain Byzantine attacks, such as reversing the local update direction or scaling the local update.
% \end{itemize}

% }

%\section{Related Works}
% In this section, we briefly review previous works on local SGD and on dealing with the client-drift issue caused by data heterogeneity. After that, we also introduce some background information regarding byzantine attacks.
%In this section, we provide a brief review of prior research on strategies employed to address the client-drift issue resulting from data heterogeneity. Furthermore, we offer background information on byzantine attacks.

% As the performance of local SGD is tested on non-convex objective functions, it is straightforward to test it on non-i.i.d. data, which is when the client-drift problem is identified \cite{zhao2018federated}. The authors find that the accuracy reduces significantly with each client training a single class of data, and they propose to improve the performance by assigning a shared subset of data to all the clients. The main cause of this phenomenon is that the local optima of the workers diverge too far from the aggregated update direction. 



% In contrary to all the methods mentioned above, the proposed DRAG is a heuristic and intuitive method that drags each local model towards the reference direction according to the extent of the drift via vector manipulation. The details of the method is elaborated in Section 4.

% Next, some background knowledge about byzantine attacks is reviewed. Due to the distributed manner of federated learning, it is vulnerable to attacks once some of the clients experience some software bugs or are even hacked by people with malicious intentions, which is widely known as the byzantine attacks \cite{so2020byzantine,cao2020fltrust,prakash2020mitigating}.



The rest of this paper is organized as follows. Section II describes the system model. DRAG is developed in Section III, followed by the design of BR-DRAG in Section IV. Section V analyzes the convergence of DRAG and BR-DRAG. Numerical results are provided in Section VI, followed by the conclusions in Section VII.

\textit{Notation}: Calligraphic letters represent sets; boldface lowercase letters indicate vectors; \(\mathbb{R}\) denotes the real number field; \(\mathbb{E}[\cdot]\) denotes expectation; \(\nabla\) stands for gradient. 
$\|\cdot\|_2$ is the $\ell_2$-norm; and $\langle\cdot, \cdot\rangle$ represents the inner product operation.

\section{System Model}
In this section, we elucidate the FL system and the corresponding setup under Byzantine attacks.

\vspace{-3 mm}
\subsection{Federated Learning Framework}

% {\color{red} These datasets are drawn from a global dataset $\mathcal{D}=\{z_i\}_{i=1}^N$,i.e., we have $\mathcal{D}=\bigcup_{m\in\mathcal{M}} \mathcal{D}_m$. 
% }

We study an FL system consisting of a PS and $M$ workers collected by the set $\mathcal{M}:=\{1,\cdots, M\}$; see Fig.~\ref{fig:sys model}. Each worker $m$ maintains a local dataset $\mathcal{D}_m$ with $N_m$ data samples. Let $\mathcal{D}=\bigcup_{m\in\mathcal{M}} \mathcal{D}_m$.
The objective of FL is to 
% solve an empirical minimization problem defined as
minimize the average sum of local functions of individual workers:
\begin{align} \label{problem}
    \min_{\boldsymbol{\theta}\in \mathbb{R}^d} f(\boldsymbol{\theta})&=\frac{1}{M}\sum_{m=1}^{M} F_m(\boldsymbol{\theta}), 
    % \\ \text{with}\quad F_m(\boldsymbol{\theta})&:=\mathbb{E}\left[F_m(\boldsymbol{\theta}; z_m)\right], m\in\mathcal{M},
\end{align}
where $\boldsymbol{\theta} \in \mathbb{R}^d$ is the $d$-dimensional model parameter, and $F_m(\boldsymbol{\theta}) \buildrel \Delta \over =  \mathbb{E}_{z_m \sim \mathcal{D}_m}\left[F_m(\boldsymbol{\theta}; z_m)\right]$ is the smooth local loss function concerning the dataset $\mathcal{D}_m$ of worker $m$.


\begin{figure}[!t]
	\centering
	\includegraphics[width=0.9\linewidth]{new_figs/drag_system_model.png}
	\caption{The architecture of a FL system with $S$ selected workers in round~$t$. For Byzantine FL systems, the attacked workers can upload arbitrary local updates to affect global model aggregation.}
	\label{fig:sys model}
\end{figure} 

% with dimension $d$ is the variable to be optimized and $\{F_m(\boldsymbol{\theta}),m\in\mathcal{M}\}$ are smooth functions, and $z_m$ is a sample randomly selected from the local dataset $\mathcal{D}_m$ of worker $m$.

With the objective function in (\ref{problem}), we attempt to solve it iteratively using local SGD. In any training round $t$, the following steps are executed:
\begin{itemize}
    \item Worker selection and model transmission: The PS first selects a subset $\mathcal{S}^t \subset \mathcal{M}$ with $S$ workers uniformly at random (UAR) without replacement, and broadcasts the global model $\boldsymbol{\theta}^t$ to the selected workers $m \in \mathcal{S}^t$.

    \item Local update: With $\boldsymbol{\theta}_m^{t,0}=\boldsymbol{\theta}^t$, each selected worker $m\in \mathcal{S}^t$ performs $U$ local updates via SGD updating rule:
    % \begin{align}
    % \boldsymbol{\theta}_{m}^{t,u+1}=\boldsymbol{\theta}_{m}^{t,u}-\frac{\eta}{B} \sum_{b=1}^B \nabla F(\boldsymbol{\theta}_{m}^{t,u}; z_{m,b}^{t,u}), \label{update}
    % \end{align}
    \begin{align} \label{update}
    \boldsymbol{\theta}_{m}^{t,u+1} \!\!=\! \boldsymbol{\theta}_{m}^{t,u} \!\!-\! \eta \nabla \!F_m(\boldsymbol{\theta}_{m}^{t,u}; z_{m}^{t,u}), \ u \!=\! 0,1,\!\cdots\!,U \!\!-\!\! 1,
    \end{align}
    where $\eta$ denotes the stepsize, and $\nabla F_m(\boldsymbol{\theta}_{m}^{t,u}; z_{m}^{t,u})$ is the stochastic gradient of worker $m$ with a mini-batch $z_{m}^{t,u}$ sampled from the local dataset $\mathcal{D}_m$ with size $|z_{m}^{t,u}|=B$.
    % $B$ is the mini-batch size; $\frac{1}{B} \sum_{b=1}^B \nabla F(\boldsymbol{\theta}_{m}^{t,u}; z_{m,b}^{t,u})$ is the mini-batch gradient to be computed by worker $m$ at iteration $u$ and $z_{m,b}^{t,u}$ are drawn independently from dataset $\mathcal{D}_m$ across all workers, batches, local iterations and training rounds. 
    

    \item Model uploading and aggregation: After local computation, each worker $m \in \mathcal{S}^t$ sends $\mathbf{g}_m^t = \boldsymbol{\theta}_m^{t, U} - \boldsymbol{\theta}^{t}$ to the PS, which represents the discrepancy between the latest local model after $U$ local updates and the current global model received at the beginning of the $t$-th training round. 
    % Specifically, for each worker $m$ in training round $t$, we define $\mathbf{g}_m^t$ as 
    % \begin{align}
    %  \mathbf{g}_m^t = \boldsymbol{\theta}_m^{t, U} - \boldsymbol{\theta}^{t}. 
    % \end{align}
    The PS updates the global model $\boldsymbol{\theta}^{t+1}$ as 
    \begin{align} \label{aggregation}
  \boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^{t}+\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t. 
    \end{align}    
\end{itemize}
Then, the next training round, i.e., the $(t+1)$-th round, starts. This training process repeats until convergence. 
% The setup described above is the basic procedure of FL using local SGD. The modifications to the framework of the proposed algorithm are detailed in the next section.

% The above framework is without a byzantine attack. While in the presence of byzantine attacks, In the randomly selected subset $\mathcal{S}^t$, we assume that there are $A$ malicious clients in a subset $\mathcal{A}^t$. At each training round $t$, after finishing $U$ local updates, each malicious client $m\in\mathcal{A}^t$ multiplies the local model update $\mathbf{g}_m^t$ by a scalar $p_m^t$ which can be either positive or negative and sends it to the server, i.e., reversing the update direction or scaling the module of the update. 

\subsection{Byzantine Attack}

% {\color{red} As shown in Fig.~\ref{attack}, during each training round $t$, after completing $U$ local updates, each malicious client $m\in\mathcal{A}^t$ manipulates their local model update $\boldsymbol{\theta}_m^{t, U} - \boldsymbol{\theta}^{t}$ by multiplying it with a scalar $p_m^t$, denote as $\hat{\mathbf{g}}_m^t=p_m^t(\boldsymbol{\theta}_m^{t, U} - \boldsymbol{\theta}^{t})$. Note that $p_m^t$ can be either positive or negative. } 

% Based on the FL framework, we proceed to elaborate on the scenario with Byzantine attacks that we aim to defend against in this work. 

Due to the decentralized implementation of FL, workers are vulnerable to Byzantine attacks arising from, e.g., poisoned data.
These workers, also known as Byzantine nodes, upload malicious updates to the PS in an attempt to degrade the convergence and learning performance.

Consider $A$ Byzantine nodes out of the $M$ workers. Let $\mathcal{A}$ collect all Byzantine nodes. $\mathcal{A}\subseteq \mathcal{M}$. In the $t$-th training round, $A^t$ malicious workers, collected by $\mathcal{A}^t$, are selected, with $\mathcal{A}^t \subseteq \mathcal{S}^t$ and $\mathcal{A}^t \subseteq \mathcal{A}$. Let $w^t \in [0,1]$ be the intensity level of the Byzantine attack: $A^t = w^t S$. 

To launch Byzantine attacks, any malicious worker $m \in \mathcal{A}^t$ sends 
% an arbitrary vector 
% % $\star \in \mathbb{R}^d$ 
% $\hat{\mathbf{g}}_m^t \in \mathbb{R}^d$ 
% to the PS. 
% % Let $\hat{\mathbf{g}}_m^t$ be the corresponding local update, that is, 
% % \begin{align}
% %     \hat{\mathbf{g}}_m^t = \star, \ \forall t, m \in \mathcal{A}^t.
% % \end{align}
% % % {\color{red} where $A = wS$, $w \in (0,1)$}. 
its malicious update $\hat{\mathbf{g}}_m^t \in \mathbb{R}^d$ 
% are sent
to the PS. 
% Together with the benign local updates $\mathbf{g}_m^t$, 
% reversing the update direction or scaling the magnitude of the global update: 
% The PS then aggregates the received local models as follows:
% Lastly, 
The PS aggregates the local models to update the global model:
\begin{align}
    \boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^{t}+\frac{1}{S}\big(\sum_{m\in\mathcal{A}^t}\hat{\mathbf{g}}_m^t+\sum_{m\in \mathcal{B}^t }\mathbf{g}_m^t\big), \label{byzantine aggregation}
\end{align}
where $\mathcal{B}^t = \mathcal{S}^t \setminus \mathcal{A}^t$ collects the benign workers in round~$t$.
% In this paper, we consider three typical attacks in our Byzantine FL model. ...





% {\color{red} where $\mathcal{B}^t = \mathcal{S}^t \setminus \mathcal{A}^t $.}

% Now we continue to elaborate on the byzantine attack that we aim to defend against in this work. In the randomly selected subset $\mathcal{S}^t$, we assume that there are $A$ malicious clients in a subset $\mathcal{A}^t$. At each training round $t$, after finishing $U$ local updates, each malicious client $m\in\mathcal{A}^t$ multiplies the local model update $\mathbf{g}_m^t$ by a scalar $p_m^t$ which can be either positive or negative and sends it to the server, i.e., reversing the update direction or scaling the module of the update. 

% \begin{figure}[t!]
% \centering
% \includegraphics[width=3in]{ATTACK.pdf}
% \caption{Illustration of byzantine attacks dealt with in this work, where the black color represents the normal clients and the red color represents the malicious ones.}
% \label{attack}
% \end{figure}

\section{DiveRgence-based Adaptive aGgregation}

In this section, we propose the DRAG algorithm, which overcomes client drifts (i.e., the deviation of local model updates from the true global gradient direction, as described in Section I-A) and accelerates convergence. We start with the update flow of DRAG and the overall algorithm.
% In contrast to most state-of-the-art algorithms designed to tackle client drift by employing control variates for local and global model alignment, the proposed DRAG method adopts a heuristic and intuitive manner that ``drags" each local model toward the reference direction through vector manipulation, with the extent of dragging determined by a metric we define as the ``degree of divergence". To establish the theoretical feasibility of DRAG, we conduct a rigorous convergence analysis that provides compelling evidence of its convergence properties. 

\subsection{Algorithm Overview}


% \textbf{Step 3:} The PS first calculates the reference direction $\mathbf{r}^t$ via (\ref{reference_cf}) and the degree of divergence $\lambda_m^t$ via (\ref{dod}). With these two variables, the PS then drags each $\mathbf{g}_m^t$ towards the reference direction $\mathbf{r}^t$ via \eqref{vm}, yielding the modified gradient $\mathbf{v}_m^t$. In the end, the PS aggregates the modified gradients with $\Delta^t=:\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{v}_m^t$ and updates the global model with $\boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^t-\Delta^t$.

\textbf{Algorithm~\ref{Alg_DRAG}} presents the implementation details of DRAG based on a new divergence-based aggregation strategy proposed in Section III-C.
At the beginning of each training round~$t$, after the worker subset $\mathcal{S}^t$ is specified, the PS determines the \textit{reference direction} $\mathbf{r}^t$ (as proposed in Section III-B),  with \eqref{reference}, to measure the degree of local divergence and broadcasts the global model $\boldsymbol{\theta}^t$ and the reference direction $\mathbf{r}^t$ to the selected workers; see Steps 4--12. 

On the local side, the workers perform local updates to generate $\mathbf{g}_m^t$, and compute the DoD $\lambda_m^t$, i.e., using \eqref{dod}, to measure the deviation from the reference direction $\mathbf{r}^t$; see Steps 14 and 15. To cope with data heterogeneity, the workers generate the modified gradient $\mathbf{v}_m^t$ by adaptively ``dragging" their local updates toward $\mathbf{r}^t$; see Step 16. After local computation, the selected workers upload $\mathbf{v}_m^t,\,\forall m \in \mathcal{S}^t$ to the PS for global update. $\mathbf{r}^t$ and $\Delta^t$ are retained to construct the reference direction for the next round; see Step 18. 

Unlike existing FL systems tackling client drifts by employing local control variates or gradient norms~\cite{liang2019variance,li2019feddane,pathak2020fedsplit,mitra2021linear,karimireddy2020scaffold,acar2021federated,karimireddy2020mime,li2020federated,gao2022feddc,pmlr-v162-qu22a}, we design a more effective strategy that guides each local model toward the reference direction through vector manipulation, which enforces directional consistency across clients.
% Detailed analysis of the algorithmic components and their underlying insights are provided in the subsequent subsections.


\begin{algorithm}[!t]
    \caption{DRAG}
    \label{Alg_DRAG}
    \begin{algorithmic}[1]
        \STATE \textbf{Input:} $\boldsymbol{\theta}^0$, $\mathcal{M}$, $S$, $T$, $U$, $\eta$, $\alpha$, $c$;
        \FOR{$t=0,1,\dots,T-1$}
            \STATE The PS selects the worker subset $\mathcal{S}^t \subseteq \mathcal{M}$ UAR;
            \IF{$t=0$}
                \FOR{each worker $m \in \mathcal{S}^t$ in parallel}
                    \STATE Perform $U$ SGD updates by \eqref{update}, and upload $\mathbf{g}_m^t$;
                \ENDFOR
                \STATE The PS obtains the reference direction $\mathbf{r}^t$ with \eqref{reference_a};
            \ELSE
                \STATE The PS obtains $\mathbf{r}^t$ with \eqref{reference_b};
            \ENDIF
            \STATE The PS sends $\boldsymbol{\theta}^t$ and $\mathbf{r}^t$ to workers $m \in \mathcal{S}^t$
            \FOR{each worker $m \in \mathcal{S}^t$ in parallel}
                \STATE Perform $U$ local SGD updates with \eqref{update};
                \STATE Compute the DoD $\lambda_m^t$ with \eqref{dod};
                \STATE Compute the modified gradient $\mathbf{v}_m^t$ with \eqref{vm}, and upload  $\mathbf{v}_m^t$ to the PS;
            \ENDFOR
            \STATE The PS aggregates $\Delta^t$ with \eqref{drag aggregation}, retains $\mathbf{r}^t$ and $\Delta^t$, and updates $\boldsymbol{\theta}^t$ to $\boldsymbol{\theta}^{t+1}$ with \eqref{drag global update};
        \ENDFOR
    \end{algorithmic}
\end{algorithm}

% \textbf{Step 1:} At each training round $t$, the PS first broadcasts the current global parameter $\boldsymbol{\theta}^{t}$ to the workers in a random selected set $\mathcal{S}^t$. 

% \textbf{Step 2:} Each worker in $\mathcal{S}^t$ then performs $U$ local updates via (\ref{update}) and sends its the difference $\mathbf{g}_m^t= \boldsymbol{\theta}_m^{t, U} - \boldsymbol{\theta}^{t}$ back to the PS.

% \textbf{Step 3:} The PS first calculates the reference direction $\mathbf{r}^t$ via (\ref{reference_cf}) and the degree of divergence $\lambda_m^t$ via (\ref{dod}). With these two variables, the PS then drags each $\mathbf{g}_m^t$ towards the reference direction $\mathbf{r}^t$ via \eqref{vm}, yielding the modified gradient $\mathbf{v}_m^t$. In the end, the PS aggregates the modified gradients with $\Delta^t=:\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{v}_m^t$ and updates the global model with $\boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^t-\Delta^t$.
\subsection{Global Reference Direction}
As mentioned in Section I, one fundamental challenge in FL is data heterogeneity, as 
% clients possess data $\mathcal{D}_m$ drawn from non-identically distributed sources. In this sense, 
the local gradients $\mathbf{g}_m^t,\, m\in \mathcal{S}^t$ tend to diverge in direction and magnitude due to discrepancies in the local objective $F_m$, leading to client drifts and global gradient deviation. 
Utilizing $\tilde{\Delta}^t=\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t$ as the global update direction (i.e., using FedAvg, as in conventional FL algorithms) could cause unstable or biased convergence~\cite{li2019convergence}.

To address client drifts, we design a global reference direction $\mathbf{r}^t$ produced by the PS at the beginning of each training round $t$, which 
% serves as a surrogate global descent signal across rounds and adaptively aligns 
helps align the local gradients. At the beginning of training round~$t$, $\mathbf{r}^t$ is constructed as 
% that serves as an implicit surrogate for the global descent direction. Rather than directly averaging the raw local updates, our approach aligns them toward a shared historical direction constructed at the server side. This reference vector acts as a unifying anchor, guiding each client update while preserving adaptive flexibility. Its design is rooted in momentum-style aggregation and is updated progressively across training rounds.
% In contrast to most state-of-the-art algorithms designed to tackle client drift by employing control variates for local and global model alignment, the proposed DRAG method adopts a heuristic and intuitive manner that ``drags" each local model toward the reference direction through vector manipulation, with the extent of dragging determined by a metric we define as the ``degree of divergence".
% A key component of DRAG is the construction of a global reference direction $\mathbf{r}^t$, which provides a consistent update signal across rounds and helps align local gradients with the global descent trajectory. To this end, we define $\mathbf{r}^t$ at the server side using an adaptive update rule that integrates historical global gradient information.
\begin{subequations}\label{reference}
\begin{empheq}[left={\mathbf{r}^t = \left\{
\begin{aligned}}, right={\end{aligned}\right.}]{align}
& \frac{1}{S}\sum\limits_{m \in \mathcal{S}^t} \mathbf{g}_m^t, & \text{for } t = 0 ;  \label{reference_a} \\
&  (1 - \alpha)\mathbf{r}^{t-1} + \alpha\Delta^{t-1}, & \text{for } t \geq 1 , \label{reference_b}
\end{empheq}
\end{subequations}
where $\Delta^{t-1} = \frac{1}{S}\sum_{m \in \mathcal{S}^{t-1}} \mathbf{v}_m^{t-1} $ is the aggregated modified gradient in round $t-1$, and 
$\alpha \in (0,1)$ is a hyperparameter controlling the weights of historical reference directions.
% In training round $t=0$, the reference direction is defined as the average of the local SGD updates, i.e., $\mathbf{r}^0 = \frac{1}{S} \sum_{m \in \mathcal{S}^0} \mathbf{g}_m^0$, since no prior global update is available.
% In subsequent training rounds $t \geq 1$, 
$\mathbf{r}^t,\,t\geq 1$, evolves recursively as a weighted combination of the previous reference direction $\mathbf{r}^{t-1}$ and the latest global update~$\Delta^{t-1}$.

During round $t$, the selected devices generate the modified gradients $\mathbf{v}_m^t$ based on the reference direction $\mathbf{r}^t$ (see Section~III-C). $\mathbf{v}_m^t,\! \forall m \in \mathcal{S}^t$, are aggregated at the PS, given as
\begin{align} \label{drag aggregation}
    \Delta^t = \frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{v}_m^t.
\end{align}
The global model is updated with $ \Delta^t$, as given by
\begin{align} \label{drag global update}
    \boldsymbol{\theta}^{t+1} = \boldsymbol{\theta}^{t} + \Delta^t.
\end{align}

We further derive the closed-form expression of $\mathbf{r}^t$ by deduction, as follows: $ \forall t \geq 1$,
\begin{align} \label{reference_cf}
    \mathbf{r}^t=\frac{(1-\alpha)^t}{S}\sum_{m\in\mathcal{S}^0}\mathbf{g}_m^0+\sum_{i=0}^{t-1}\alpha(1-\alpha)^{t-i-1}\Delta^i. 
\end{align}
Here, $\mathbf{r}^t$ is an exponential moving average over all historical modified global updates~$\Delta^i,\,i=1,\cdots,t-1$, with exponentially decaying weights $(1-\alpha)^{t-i-1}$. This momentum-style construction enables $\mathbf{r}^t$ to preserve long-term memory of update history while prioritizing those more recent. As $\alpha \to 1$, $\mathbf{r}^t$ reduces to the most recent update $\Delta^{t-1}$. As $\alpha \to 0$, a smoother reference direction integrates past variations more conservatively. By tuning~$\alpha$, $\mathbf{r}^t,\,\forall t$, can flexibly balance adaptivity and stability under different data heterogeneity levels. 
% Its design is rooted in momentum-style aggregation and is updated progressively across training rounds.
% This update rule defines $\mathbf{r}^t$ as an exponential moving average (EMA) of all past global updates. 
% This reveals that $\mathbf{r}^t$ maintains a memory of the entire update history, with more recent updates weighted more heavily. Notably, when $\alpha=1$, the update reduces to the last global direction $\Delta^{t-1}$; when $\alpha \rightarrow 0$, it converges to a long-term average. Such a momentum-style reference construction enhances the temporal stability of model updates and provides a meaningful anchor for divergence correction.
% The expression clearly demonstrates that the reference direction has a momentum form which is a weighted sum of all the historical global gradients $\Delta^i$, for $i=0,1,...,t$. Significantly, the weights assigned to the most recent global updates progressively increase as $\alpha$ grows larger. When $\alpha$ is set to 1, the reference direction becomes equivalent to the previous global gradient $\Delta^{t-1}$. Notably, the hyper-parameter $\alpha$ can be adjusted to fit different practical scenarios.

% It is obvious from the expression that the reference direction is a weighted sum of all the historical global update directions. Also, the weights of the most recent global updates increase as $\alpha$ becomes larger. And when $\alpha=1$, the reference direction is tantamount to the previous global gradient $\Delta^{t-1}$. Note that $\alpha$ is a tunable hyper-parameter to adaptively fit different practical scenarios.

% \subsection{Rationale behind DRAG}
% \subsection{Global Reference Direction}
% In this subsection, we elaborate on the rationale behind the ideas of DRAG, including the reference direction, the degree of divergence and the vector manipulation process.

% To commence, we first design the reference direction generated at the PS's side, which plays a vital role in dragging the local gradients as a prerequisite.

% The aim of the reference direction is to provide a feasible and reasonable direction when modifying the local gradients and to better forge a global update direction. Recall that the updating formula of the reference direction is:


% Particularly, the updating formula for the reference direction is given as:
% % \begin{align}
% %     &\mathbf{r}^t=(1-\alpha)\mathbf{r}^{t-1}+\alpha\Delta^{t-1}, for\:t\geq 1\notag\\
% %     &\mathbf{r}^t=\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t, for\:t=0.\label{reference}
% % \end{align}
% \begin{subequations}\label{reference}
% \begin{empheq}[left={\mathbf{r}^t = \left\{
% \begin{aligned}}, right={\end{aligned}\right.}]{align}
% & \frac{1}{S}\sum\limits_{m \in \mathcal{S}^t} \mathbf{g}_m^t, & \text{for } t = 0;  \tag{6a} \label{reference_a} \\
% &  (1 - \alpha)\mathbf{r}^{t-1} + \alpha\Delta^{t-1}, & \text{for } t \geq 1 , \tag{6b} \label{reference_b}
% \end{empheq}
% \end{subequations}

% \begin{align}
% \mathbf{r}^t=
%     \begin{cases}
%       (1-\alpha)\mathbf{r}^{t-1}+\alpha\Delta^{t-1}, ~~\text{for  $t \geq 1$},  \\
%       \frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t,~~~~~~~~~~~~ \text{for  $t=0$}.
%     \end{cases}\label{reference}
% \end{align}

% -------------------------------- To be further discussed --------------------------------%
\iffalse
We further analyze the directional consistency between the reference direction $\mathbf{r}^t$ and the ideal global gradient direction.
% , in comparison with FedAvg, i.e., ${\tilde{\Delta}}^t=\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t$.
We establish the following \textbf{theorem}.
\begin{theorem}   \label{thm:alignment}
Let \textbf{Assumptions 1–2} hold. For $\forall t\ge 1$, $\nabla f(\boldsymbol{\theta}^{t})$ and ${\tilde{\Delta}}^t=\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t$ denote the true global gradient and the aggregated local updates, respectively.
With $\eta \leq \frac{\alpha}{{4LU}} $, the reference direction $\mathbf{r}^t$ attains a higher directional alignment with $\nabla f(\boldsymbol{\theta}^{t})$ in expectation, i.e.,
% \begin{align} \label{eq:angle}
%   \arccos  \mathbb{E} \! \left[
%       \frac{\left\langle \mathbf r^{t},
%                         \nabla f\!\left(\boldsymbol{\theta}^{t}\right)\right\rangle}
%            {\left\| \mathbf r^{t}\right\|
%             \left\|\nabla f\!\left(\boldsymbol{\theta}^{t}\right)\right\|} \right]
%   \leq
%   \arccos \mathbb{E} \! \left[
%       \frac{ \langle \tilde{\Delta}^{t},
%                         \nabla f\!\left(\boldsymbol{\theta}^{t}\right) \rangle}
%            { \| \tilde{\Delta}^{t} \|
%             \left\|\nabla f\!\left(\boldsymbol{\theta}^{t}\right)\right\|} \right] .
% \end{align}
\begin{align} \label{eq:angle}
   \mathbb{E} \! \left[
       {\left\langle \mathbf r^{t},
                        \nabla f\!\left(\boldsymbol{\theta}^{t}\right)\right\rangle} \right]
  \geq
   \mathbb{E} \! \left[
       { \langle \tilde{\Delta}^{t},
                        \nabla f\!\left(\boldsymbol{\theta}^{t}\right) \rangle}
            \right] .
\end{align}
\end{theorem}

\begin{proof}
See Appendix~A.
\end{proof}
{\color{red}
From \textbf{Theorem~\ref{thm:alignment}}, the designed $\mathbf{r}^t$ provides a more precise global update direction compared with the raw aggregation $\tilde{\Delta}^t$, which enables a temporally coherent anchor to guide the DRAG update process.
} 
\fi
% This design not only mitigates erratic update behavior caused by diverging local gradients but also provides a temporally coherent anchor to guide the DRAG update process.
% The objective of the reference direction $\mathbf{r}^t$ is to offer a practical and sensible direction for modifying the local gradients, thereby facilitating the formation of an enhanced global update direction. 


\subsection{Divergence-based Local Gradient Modification}
With the reference direction $\mathbf{r}^t$, we design the adaptive modified gradients $\mathbf{v}_m^t, \forall m \in \mathcal{S}^t$ per training round $t$ to combat client drifts.
We first specify the new metric, DoD,
% To commence, we define two crucial new variables: the reference direction and the degree of divergence. These variables play a key role in the process of dragging the local gradient.
to measure the extent to which the local update $\mathbf{g}_m^t$ of each worker $m$ in round $t$ diverges from $\mathbf{r}^t$. 
The angle $\angle_{m}^t$ is introduced to capture the directional deviation of $\mathbf{g}_m^t$ from $\mathbf{r}^t$ based on the normalized cosine similarity, i.e.,  
\begin{align}
    \angle_{m}^t=\arccos{\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\cdot\|\mathbf{r}^t\|}},
\end{align}
where $\angle_{m}^t = 0$ indicates perfect alignment, and $\angle_{m}^t = \pi$ denotes complete opposition. 
By dividing $\angle_{m}^t$ over $\frac{\pi}{2}$ and approximating the $\arccos$ function with a linear function (e.g., $y=-\frac{\pi}{2}x+\frac{\pi}{2}$), the DoD, $\lambda_m^t$, is defined as 
% (this is one of many choices, one may also use $y=-x+\pi/2$)
\begin{align}
    \lambda_m^t=:c\big(1-\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\cdot \|\mathbf{r}^t\|}\big)\in[0,2c],\label{dod}
\end{align}
% where $c\in[0,1]$ is a hyper-parameter to provide the flexibility to manually adjust the metric to suit various settings. 
where $c \in [0, 1]$ is a tunable hyperparameter modulating the contribution of the divergence signal. 
By adjusting $c$, one can control the extent to which local updates are aligned with $\mathbf{r}^t$.
% , enabling adaptive calibration of the correction strength in response to varying degrees of data heterogeneity or system settings.
Moreover, $\lambda_m^t$ quantifies the geometric inconsistency between $\mathbf{g}_m^t$ and $\mathbf{r}^t$ on the unit hypersphere, offering a continuous and differentiable measure of gradient divergence, as opposed to discrete thresholds or binary decisions. 

% \noindent\textbf{Remark 1.} {\color{red}compare with contrastive learning or domain adaptation}



% It is worth noting that $\lambda_m^t$ is adaptable for each worker $m$ at each round $t$, with a larger value of $\lambda_m^t$ indicating a more significant divergence between the local gradient $\mathbf{g}_m^t$ and the reference direction $\mathbf{r}^t$.
% Notably, $\lambda_m^t = 0$ when $\mathbf{g}_m^t$ is fully aligned with $\mathbf{r}^t$, and $\lambda_m^t = 2c$ when they are exactly opposite. By translating angular misalignment into a scalar divergence measure, this design enables an interpretable and flexible mechanism to modulate the extent of local update modification in the DRAG framework.


% {\color{red}Vector Manipulation}
% \subsubsection{Vector Manipulation}
With the reference direction $\mathbf{r}^t$ and the DoD~$\lambda_m^t$, we establish the modified gradient $\mathbf{v}_m^t$, which ``drags" the local gradients $\mathbf{g}_m^t$ towards $\mathbf{r}^t$ based on $\lambda_m^t$,
% through vector manipulation, 
as given by
\begin{align}\label{vm}
    \mathbf{v}_m^t=(1-\lambda_m^t)\mathbf{g}_m^t+\lambda_m^t\frac{\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t, 
\end{align}
% where $\mathbf{v}_m^t$ is a weighted sum of the local update $\mathbf{g}_m^t$ and the normalized reference direction $\frac{\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t$, with the respective weights of $1-\lambda_m^t$ and $\lambda_m^t$.
% $0\leq \lambda_m^t \leq 2c$.
If there is significant misalignment between $\mathbf{r}^t$ and $\mathbf{g}_m^t$, the local update is dragged toward the normalized reference direction $\frac{\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t$, reducing client drifts while preserving the diversity of local information.

% By adjusting $c$, the client drift can be effectively controlled while preserving the diversity of the local gradients. 
% Additionally, the reference direction $\mathbf{r}^t$ is normalized to match the norm of $\mathbf{g}_m^t$, ensuring that the modified gradient $\mathbf{v}_m^t$ consistently has a greater component aligned with $\mathbf{r}^t$ compared to the original gradient $\mathbf{g}_m^t$. 
% This construction performs a norm-preserving interpolation between the original gradient $\mathbf{g}_m^t$ and the reference direction $\mathbf{r}^t$, with $\lambda_m^t$ adaptively controlling the strength of directional correction. 
% When the local update aligns well with $\mathbf{r}^t$, $\mathbf{v}_m^t$ remains close to $\mathbf{g}_m^t$; when misalignment is significant, the update is dragged toward the global descent path. This design enables DRAG to reduce client drift while preserving meaningful local information, achieving a balance between coherence and adaptability under data heterogeneity.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% \begin{theorem}\label{thm:angle_full}
% Let the initial and post-modification angles be
% $$
% {\angle}_{m}^t
% =
% \arccos\frac{\langle\mathbf{g}_m^t,\mathbf{r}^t\rangle}
%                 {\|\mathbf{g}_m^t\|\,\|\mathbf{r}^t\|}
% \in(0,\pi),\qquad
% \tilde{\angle}_{m}^t
% =
% \arccos\frac{\langle\mathbf{v}_m^t,\mathbf{r}^t\rangle}
%                 {\|\mathbf{v}_m^t\|\,\|\mathbf{r}^t\|},
% $$
% where $\mathbf{v}_m^t$ and $\lambda_m^t$ follow \eqref{vm}–\eqref{dod} with
% $\lambda_m^t=c\bigl(1-\cos{\angle}_{m}^t\bigr)$ and $c\in(0,1]$.

% \begin{enumerate}[label=(\alph*)]
% \item \emph{Strict reduction for any angle.}
%       For every ${\angle}_{m}^t\in(0,\pi)$ and $c>0$,
%       \begin{align}
%         \tilde{\angle}_{m}^t<{\angle}_{m}^t .
%         \label{eq:strict_all}
%       \end{align}

% \item \emph{Explicit contraction factors.}  Define
%       $\kappa(\lambda)=1-\tfrac{\lambda}{2}$ if $\lambda\le1$
%       and $\kappa(\lambda)=\lambda-1$ if $\lambda>1$
%       (note $0\le\kappa(\lambda)\le1$ for $\lambda\in[0,2]$):
%       \begin{align}
%         \tilde{\angle}_{m}^t
%         \le
%         \begin{cases}
%           \kappa(\lambda_m^t)\;{\angle}_{m}^t,
%           & 0<{\angle}_{m}^t\le\displaystyle\frac{\pi}{2},\\[8pt]
%           \pi-\kappa(\lambda_m^t)\bigl(\pi-{\angle}_{m}^t\bigr),
%           & \displaystyle\frac{\pi}{2}<{\angle}_{m}^t<\pi .
%         \end{cases}
%         \label{eq:factor_all}
%       \end{align}

% \item \emph{Guaranteed acuteness for $\,c\ge\frac14$.}  
%       If $c\ge\frac14$ and ${\angle}_{m}^t\in\bigl(\tfrac{\pi}{2},\pi\bigr)$,
%       then
%       \begin{align}
%         \tilde{\angle}_{m}^t<\frac{\pi}{2}.
%         \label{eq:acute}
%       \end{align}
% \end{enumerate}
% \end{theorem}

% \begin{proof}[Sketch\footnote{Full derivation in Appendix~\ref{app:angle_full}.}]
% Let $c_\theta=\cos{\angle}_{m}^t$, $s_\theta=\sin{\angle}_{m}^t>0$ and
% $\lambda=\lambda_m^t=c(1-c_\theta)\in(0,2c]\subseteq(0,2]$.
% With
% \[
% N=(1-\lambda)c_\theta+\lambda,\qquad
% S=(1-\lambda)s_\theta ,
% \]
% we have
% $$
% \cos\tilde{\angle}_{m}^t
% =
% \frac{N}{\sqrt{N^{2}+S^{2}}}.
% $$
% Writing $\Delta=\cos\tilde{\angle}_{m}^t-\cos{\angle}_{m}^t$ and
% after algebra (see Appendix) one obtains
% \[
% \Delta
% =
% \frac{\lambda\!\bigl[\lambda+2(1-\lambda)c_\theta\bigr]}
%      {\sqrt{N^{2}+S^{2}}\;+\;|c_\theta|}\;>\;0,
% \]
% because $\lambda>0$ and the bracket is non-negative for every
% $c_\theta\in[-1,1]$.  Hence $\cos\tilde{\angle}_{m}^t>\cos{\angle}_{m}^t$
% and \eqref{eq:strict_all} follows.

% The factor bounds in \eqref{eq:factor_all} are derived by bounding
% $\tan\tilde{\angle}_{m}^t$ separately on
% $(0,\tfrac{\pi}{2}]$ and $(\tfrac{\pi}{2},\pi)$ and using
% $\arctan x\le x$; details are identical to the acute-angle derivation
% and included in Appendix~\ref{app:angle_full}.

% For \eqref{eq:acute}, note that $\tilde{\angle}_{m}^t<\frac{\pi}{2}$
% is equivalent to $N\ge0$.  
% Since $c_\theta=-u$ with $u\in(0,1)$,
% \[
% N
% =\lambda-(1-\lambda)u
% =c(1+u)-(1-c(1+u))u
% \;=\;
% (1+u)^2c-\;u .
% \]
% One has $N\ge0$ for all $u\in(0,1)$ provided
% $c\ge\max_{u\in(0,1)}\frac{u}{(1+u)^{2}}=\frac14$, giving \eqref{eq:acute}.
% \end{proof}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


% \smallskip
% \smallskip

% \noindent\textbf{Remark 2.} 
Fig.~\ref{vectorm} illustrates the gradient modification in DRAG based on the proposed reference direction $\mathbf{r}^t$ and DoD $\lambda_m^t$. For moderate deviation ($0<\lambda_m^t\le1$) in Fig.~\ref{vectorm}(a), the modification projects $\mathbf{g}_m^t$ toward $\mathbf{r}^t$ and enlarges the aligned component to at least $\frac{\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t$, reducing client drifts.
% while keeping the orthogonal part intact. 
For severe deviation ($1<\lambda_m^t\le2$) in Fig.~\ref{vectorm}(b), the local gradient diverges in the opposite direction to $\mathbf{r}^t$. We reverse $\mathbf{g}_m^t$ according to (\ref{vm}) to ensure adherence to the correct update direction.
% Consequently, each local gradient $\mathbf{g}_m^t$ contributes a stronger consensus signal without sacrificing gradient diversity.


% With the vector manipulation operation, if $0<\lambda_m^t\leq 1$, the resulted $\mathbf{v}_m^t$ both mitigates the drift and at the same time preserves the diversity of each local gradient; if $1<\lambda_m^t\leq 2$, indicating that the local gradient diverges too much from the reference direction, the $\mathbf{g}_m^t$ portion will be reversed according to (\ref{vm}) so as to abide by the correct update direction.
% An illustration demonstrating how the degree of divergence $\lambda_m^t$ and the reference direction $\mathbf{r}^t$ can be utilized to guide the local gradient is provided in Fig.~\ref{vectorm}. Specifically, using the vector manipulation operation, when $0<\lambda_m^t\leq 1$, the resulting $\mathbf{v}_m^t$ effectively mitigates the drift while preserving the diversity of each local gradient. Conversely, when $1<\lambda_m^t\leq 2$, indicating that the local gradient diverges in the opposite direction to the reference direction, the $\mathbf{g}_m^t$ component is reversed according to (\ref{vm}) to ensure adherence to the correct update direction.

% As illustrated in Fig.~\ref{vectorm}, it can be seen that the modified gradient $\mathbf{v}_m^t$ substantially reduces client drift while concurrently preserving the diversity of gradients. This is achieved by increasing the component aligned with the reference direction while maintaining the component in the original direction.

% \smallskip
% \smallskip

% \noindent\textbf{Remark 3.} 
% \textit{
% Compared to conventional drift-mitigation methods~\cite{karimireddy2020scaffold,acar2021federated} with local control variates, DRAG only aligns the local gradients toward the reference direction $\mathbf{r}^t$ with \eqref{vm}, which requires no extra control variables and incurs negligible memory and communication overhead. 
Unlike conventional drift-mitigation methods~\cite{karimireddy2020scaffold,acar2021federated}, which rely on local control variates, DRAG aligns local gradients with $\mathbf{r}^t$ via \eqref{vm}, eliminating the need for extra control variables and incurring negligible memory and communication overhead.
% Moreover, the methods with local gradient-norm regularizers, e.g., the one developed in~\cite{li2020federated}, are unsuitable for handling heterogeneous data distributions, as the local gradients often diverge from the global gradient. By contrast, DRAG achieves tighter alignment with the global objective by dragging the local updates toward a trusted direction based on DoD~\eqref{dod} and accelerates convergence under severe heterogeneity.
Moreover, methods employing local gradient-norm regularizers, e.g.,\cite{li2020federated}, are ineffective under heterogeneous data distributions, since local gradients may deviate substantially from the global gradient. In contrast, DRAG attains closer alignment with the global gradient by steering local updates toward $\mathbf{r}^t$ using the DoD metric $\lambda_m^t$ in~\eqref{dod}, accelerating convergence even under severe heterogeneity.

% Conventional drift-mitigation methods~\cite{karimireddy2020scaffold,acar2021federated} utilize local control variates to keep each client’s trajectory close to the global model, essentially regulating the \emph{magnitude} of local updates. DRAG takes a fundamentally different route: by computing the degree of divergence $\lambda_m^t$ in \eqref{dod} and applying the closed-form transformation in \eqref{vm}, it \emph{reorients} every local gradient toward the reference direction $\mathbf{r}^t$ while preserving its orthogonal component. This directional correction (i) requires no extra control variables, (ii) incurs negligible memory and communication overhead, and (iii) remains effective even when the true global descent direction is strongly obscured by non-IID data. Consequently, DRAG achieves tighter alignment with the global objective than norm-based regularizers, yet retains sufficient gradient diversity to accelerate convergence under severe heterogeneity.

% {\color{red}describe comparison with scaffold feddyn using control variables or fedprox??}
% While the existing methods \cite{li2020federated} generally utilize the norms of the local gradients to quantify the degree of similarity of the local functions, they may not be suitable for handling heterogeneous data distributions, as the local gradients often diverge from the global gradient. Thus, we propose a more appropriate and reasonable metric to quantify the dissimilarity in such scenarios.
% {\color{red}
% In contrast to most state-of-the-art algorithms designed to tackle client drift by employing control variates for local and global model alignment, the proposed DRAG method adopts a heuristic and intuitive manner that ``drags" each local model toward the reference direction through vector manipulation, with the extent of dragging determined by a metric we define as the ``degree of divergence".

% To summarize the DRAG algorithm, it stands apart from control-variates-based methods by offering an intuitive geometric explanation. By adaptively dragging the local gradient towards the reference direction based on each worker's degree of divergence, the algorithm incorporates momentum and provides valuable insights for the desired update direction. The vector manipulation effectively tackles the client drift issue while preventing excessive correction of local gradients, thereby preserving their diversity. Both the reference direction and the degree of divergence include hyper-parameters that enable adaptation to different practical scenarios.

% }



% \noindent\textbf{Usefulness of the degree of divergence.}
% From Figure 1 we can observe that the modified gradient $\mathbf{v}_m^t$ has reduced client drift and at the same time preserved its gradient diversity since it has more component on the reference direction and preserves the the component on its own direction.
% To briefly sum up the DRAG algorithm, unlike most control-variates-based algorithms that compensate for the drift, the proposed DRAG method has an intuitive geometric explanation. According to each worker's degree of divergence, the local gradient is adaptively dragged towards the reference direction, which has a momentum form and provides an insight for the desired update direction. The vector manipulation both deals with the client drift issue and prevents the local gradients from being over-corrected by preserving its gradient diversity. The reference direction and the degree of divergence both incorporates hyper-parameters to adapt themselves to different practical scenarios.


\begin{figure}[t!]
\centering
\subfloat[$0<\lambda_m^t\leq 1$]{\includegraphics[width=1.485in]{illustration_1.pdf}}
\subfloat[$1<\lambda_m^t$]{\includegraphics[width=1.945in]{illustration_2.pdf}}
\caption{Illustration of vector modification of DRAG. For workers with different DoD values, the modified gradient $\mathbf{v}_m^t$ (solid black line) has a larger component on the reference direction $\mathbf{r}^{t}$ (solid yellow line) than the raw local update $\mathbf{g}_{m}^t$ (solid blue line).}
\label{vectorm}
\end{figure}




 


% Then the PS drags each gradient $\mathbf{g}_m^t$ towards the reference direction depending upon the degree of divergence $\lambda_m^t$ via vector manipulation, i,e, the modified gradient $\mathbf{v}_m^t$ is written as $\mathbf{v}_m^t=:(1-\lambda_m^t)\mathbf{g}_m^t+\frac{\lambda_m^t\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t$. In the end, the PS aggregates the modified gradients with $\Delta^t=:\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{v}_m^t$ and updates the global model with $\boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^t-\Delta^t$.


% First, the gradient $\mathbf{g}_m^t$ of worker $m$ at training round $t$, is defined as the difference between the latest local model after $U$ local updates and the original global model received at the beginning of the training round, i.e., $\mathbf{g}_m^t=:\boldsymbol{\theta}_m^{t,U}-\boldsymbol{\theta}^{t}$. Second, the PS will modify the received gradients of each worker in $\mathcal{S}$ and we define the modified gradient of worker $m$ at round $t$ as $\mathbf{v}_m^t$. 


% With these definitions, we are ready to introduce the key part of DRAG. Firstly, we need to determine the reference direction $\mathbf{r}^t$ which acts as a reference vector for dragging each local gradient. The reference direction has a momentum form which is a weighted sum of all the historical global gradients $\Delta^i$, for $i=0,1,...,t$. The updating formula of $\mathbf{r}^t$ is $\mathbf{r}^t=(1-\alpha)\mathbf{r}^{t-1}+\alpha\Delta^{t-1}$ where $\alpha\in(0,1)$ is some constant to control the weights of the historical directions. Particularly, $\mathbf{r}^t=\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t$ for $t=0$. With the reference direction, the PS can readily compute the ``degree of divergence'' $\lambda_m^t$ of each worker $m$ at round $t$, i.e., $\lambda_m^t=:c\left(1-\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\right)\in[0,2c]$, where $c$ is some constant used to adaptively tune the degree of divergence according to different scenarios. 







\section{DiveRgence-based Adaptive aGgregation}
% In this subsection, we demonstrate the necessary modifications made to the DRAG algorithm to render it robust to byzantine attacks, i.e., the malicious clients can reverse the gradient direction or scale the module of it

This section enhances DRAG with resistance to Byzantine attacks, and develops BR-DRAG;
% to tackle data heterogeneity while maintaining secure model training under Byzantine attacks. 
% % We first introduce the update flow of BR-DRAG to illustrate its algorithmic structure.
see \textbf{Algorithm~\ref{Alg_BRDRAG}}.

\subsection{Algorithm Overview}

% Based on the insights of the reference direction and DoD in DRAG, t

Unlike DRAG, BR-DRAG maintains a root dataset $\mathcal{D}_{\mathrm{root}}$  for the PS to generate a trusted reference direction before model training. 
At the beginning of each training round~$t$, the worker subset $\mathcal{S}^t$ is specified. The PS computes the current reference direction $\mathbf{r}^t$ based on $\mathcal{D}_{\mathrm{root}}$ (as opposed to aggregating the local updates $\mathbf{v}_m^t,\,\forall m\in \mathcal{S}^t$, from the selected workers, as in DRAG), and sends $\mathbf{r}^t$ and the current global model $\boldsymbol{\theta}^t$ to the selected workers; see Step 4.

At the selected workers, either malicious or benign, SGD is performed to generate the local updates $\mathbf{g}_m^t,\, \forall m\in \mathcal{S}^t$; see Step~6. 
The selected workers upload $\mathbf{v}_m^t$ to the PS for global update. 
Based on the DoD, $\lambda_m^t$, in \eqref{dod new}, the PS compute the modified gradients $\mathbf{v}_m^t$ with \eqref{vm1} to counteract malicious attacks; see Step 8.
The updated global model $\boldsymbol{\theta}^{t+1}$ is used to start the next round; see Step~9.
Compared to DRAG, BR-DRAG updates both the reference direction $\mathbf{r}^t$ and the module normalization operation to tolerate malicious attacks and ensure system convergence.
% % {\color{red}Based on efficient changes to DRAG, we implement local gradient modifications to ensure a safe global model update.} 
% In the rest of this section, we provide a detailed analysis of reference direction and gradient modification in BR-DRAG under Byzantine attacks.




\begin{algorithm}[!t]
    \caption{BR-DRAG}
    \label{Alg_BRDRAG}
    \begin{algorithmic}[1]
        \STATE \textbf{Input:} $\boldsymbol{\theta}^0$, $\mathcal{M}$, $S$, $T$, $U$, $\eta$, $\alpha$, $c^t$, $\mathcal{D}_{\mathrm{root}}$;
        % \STATE{The PS obtains $\mathcal{D}_{\mathrm{root}}$ with \eqref{}}
        \FOR{$t=0,\dots,T-1$}
            \STATE The PS selects $\mathcal{S}^t \subseteq \mathcal{M}$ UAR;
            \STATE The PS obtains the reference direction $\mathbf{r}^t$ with \eqref{trust}, and sends $\boldsymbol{\theta}^t$ and $\mathbf{r}^t$ to workers $m \in \mathcal{S}^t$;
            \FOR{each worker $m \in \mathcal{S}^t$ in parallel}
                \STATE Perform $U$ local SGD updates with \eqref{update}, and upload $\mathbf{g}_m^t$ to the PS;
            \ENDFOR
            \STATE $\forall m \in \mathcal{S}^t$, the PS computes the DoD $\lambda_m^t$ and modified gradient $\mathbf{v}_m^t$ with \eqref{dod new} and \eqref{vm1};
            \STATE The PS updates $\boldsymbol{\theta}^t$ to $\boldsymbol{\theta}^{t+1}$ with \eqref{byzantine drag global update};
            % \STATE The PS aggregates $\Delta^t$ with \eqref{byzantine drag aggregation}, and updates $\boldsymbol{\theta}^t$ with \eqref{byzantine drag global update}
        \ENDFOR
    \end{algorithmic}
\end{algorithm}

\subsection{Byzantine-Resilient Global Reference Direction}

% \begin{algorithm}[!t]
%     \caption{BR-DRAG: Byzantine-Resilient DRAG}
%     \label{Alg_BRDRAG}
%     \begin{algorithmic}[1]
%         \STATE \textbf{Input:} Initial global model $\boldsymbol{\theta}^0$, total devices $\mathcal{M}$, local steps $U$, rounds $T$, stepsize $\eta$, root dataset $\mathcal{D}_{\text{root}}$, hyperparameter $c$
%         \FOR{$t=0,\dots,T-1$}
%             \STATE Randomly select subset $\mathcal{S}^t \subseteq \mathcal{M}$ of size $S$; broadcast $\boldsymbol{\theta}^t$ to all $m \in \mathcal{S}^t$
%             \FOR{each device $m \in \mathcal{S}^t$ in parallel}
%                 \STATE Perform $U$ local SGD updates from $\boldsymbol{\theta}^{t}$ to obtain $\boldsymbol{\theta}_m^{t,U}$
%                 \STATE Compute local gradient $\mathbf{g}_m^t = \boldsymbol{\theta}_m^{t,U} - \boldsymbol{\theta}^t$
%             \ENDFOR
%             \STATE PS computes trusted reference update on $\mathcal{D}_{\text{root}}$: $\mathbf{r}^t = \boldsymbol{\theta}^{t,U} - \boldsymbol{\theta}^t$ with \eqref{trust}
%             \FOR{each device $m \in \mathcal{S}^t$ in parallel}
%                 \STATE Compute divergence degree $\lambda_m^t$ with \eqref{dod}
%                 \STATE Normalize and modify local gradient: $\mathbf{v}_m^t = (1 - \lambda_m^t) \frac{\|\mathbf{r}^t\|}{\|\mathbf{g}_m^t\|} \mathbf{g}_m^t + \lambda_m^t \mathbf{r}^t$
%             \ENDFOR
%             \STATE Aggregate: $\Delta^t = \frac{1}{S} \left( \sum_{m \in \mathcal{S}^t \setminus \mathcal{A}^t} \mathbf{v}_m^t + \sum_{m \in \mathcal{A}^t} \hat{\mathbf{v}}_m^t \right)$
%             \STATE Update global model: $\boldsymbol{\theta}^{t+1} = \boldsymbol{\theta}^t - \Delta^t$
%         \ENDFOR
%     \end{algorithmic}
% \end{algorithm}


% By tampering with training samples or modifying the local gradients, the malicious effects of the attackers can undermine the reference direction $\mathbf{r}^t$ formed by the weighted sum of all historical global update directions (as done in DRAG), making it less effective in guiding the training.
To defend against the Byzantine attacks, the key improvements of BR-DRAG over DRAG are the selection of $\mathbf{r}^t$ and the modified gradient $\mathbf{v}_m^t$.

% These attacks involve malicious clients attempting to reverse the gradient direction or scale its module. By incorporating specific modifications, DRAG is fortified to handle such adversarial behavior while maintaining the accuracy and security of the federated learning process.

In the presence of malicious local updates (from the malicious workers), $\mathbf{r}^t$ needs to be constructed in a trustworthy way.
% The primary improvement demanded here is the selection of the reference direction. Due to the malicious effects by the attackers, the reference direction formed by the weighted sum of all the historical global update directions will not suffice to serve as a guidance of the training process. 
Before the training, the PS maintains the small, root dataset $\mathcal{D}_{\text{root}} \subset \mathcal{D}$ to compute $\mathbf{r}^t$. 
As assumed in~\cite{cao2020fltrust}, $\mathcal{D}_{\text{root}}$ is trustworthy 
% e.g., from a trusted source or collected directly by the PS. 
% % In practice, $\mathcal{D}_{\text{root}}$ is instantiated from a vetted pool~\cite{northcutt2021confident} (or from attested endpoints under admission control with techniques like de-duplication and outlier filtering~\cite{10.14778/3603581.3603583}). 
, e.g., derived from a vetted pool~\cite{northcutt2021confident} or attested endpoints operating under admission control, employing techniques such as de-duplication and outlier filtering~\cite{10.14778/3603581.3603583}.
To hedge against residual interference like label noise, $\mathbf{r}^t$ can be computed via a robust reducer (e.g., trimmed mean~\cite{fang2020local} or geometric median~\cite{minsker2015geometric}) over per-sample gradients on $\mathcal{D}_{\text{root}}$.
% Before the training starts, the PS maintains a small root dataset $\mathcal{D}_{\text{root}}\in \mathcal{D}$, where $\mathcal{D}=\bigcup_{m\in\mathcal{M}} \mathcal{D}_m$, as done in~\cite{cao2020fltrust}. 

During training round $t$, based on $\mathcal{D}_{\text{root}}$, the global model $\boldsymbol{\theta}^t$ is updated for $U$ SGD iterations to obtain $\boldsymbol{\theta}^{t,U}$: 
% \begin{equation}\label{update_root}
% \boldsymbol{\theta}^{t,u \!+\! 1} \!=\! \boldsymbol{\theta}^{t,u} \!-\! \frac{\eta}{B} \sum_{b=1}^B \nabla F(\boldsymbol{\theta}^{t,u}; z_{b}^{t,u}), \ \ u \!=\! 0,...,U \!-\! 1,
% \end{equation}
\begin{equation}\label{update_root_new}
\boldsymbol{\theta}^{t,u + 1} = \boldsymbol{\theta}^{t,u} - \eta \nabla f(\boldsymbol{\theta}^{t,u}; z^{t,u}), \ \ u = 0,1,\ldots,U - 1,
\end{equation}
where $\boldsymbol{\theta}^{t,0}=\boldsymbol{\theta}^{t}$, and $z^{t,u}$ is drawn independently from $\mathcal{D}_{\text{root}}$ across all batches, local iterations and training rounds.

With $\boldsymbol{\theta}^{t,U}$ obtained from the trusted root dataset $\mathcal{D}_{\text{root}}$, 
% {\color{red}With the...} $\boldsymbol{\theta}^{t,U}$, 
the reference direction $\mathbf{r}^t$ in BR-DRAG  is designed as 
\begin{align} \label{trust}
\mathbf{r}^t = \boldsymbol{\theta}^{t,U}-\boldsymbol{\theta}^{t} = - \eta \sum_{u=0}^{U-1} \nabla f \left(\boldsymbol{\theta}^{t, u} ; z^{t, u}\right) ,
\end{align}
which, unlike \eqref{reference} in DRAG, stems directly from $\mathcal{D}_{\text{root}}$, providing a trustworthy update direction of the global objective and thus mitigating the impact of Byzantine behaviors. 

The PS aggregates the local updates of the workers as in \eqref{byzantine aggregation}. For each uploaded local update, the PS generates the modified gradient $\mathbf{v}_m^t,\,\forall m\in \mathcal{S}^t$, to resist Byzantine attacks (as will be designed in Section IV-C). The global model is updated as 
\begin{align} \label{byzantine drag global update}
    \boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^t+\Delta^t=\boldsymbol{\theta}^t+ \frac{1}{S}\big(\sum_{m\in\mathcal{A}^t}\hat{\mathbf{v}}_m^t+\sum_{m\in \mathcal{B}^t }\mathbf{v}_m^t\big),
    % = \boldsymbol{\theta}^t+\frac{1}{S} \left( \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} \mathbf{v}_m^t + \sum_{m \in \mathcal{A}^t} \mathbf{v}_m^t  \right).
\end{align}
% Unlike exponential moving average (EMA)-based constructions that recursively incorporate past aggregated updates—which may be contaminated by poisoned or manipulated signals—this root-data-driven approach avoids the accumulation of adversarial noise over rounds and offers stronger robustness and consistency. By leveraging freshly computed gradients from a reliable and stationary data source at each round, the server-side reference direction ensures stable convergence guidance without being distorted by untrusted client information.
where ${\mathbf{v}}_m^t$ stands for the benign modified gradient (as in DRAG), while $\hat{\mathbf{v}}_m^t$ denotes the malicious modified gradients.

% \begin{align} \label{byzantine drag aggregation}
%     \Delta^t = \frac{1}{S} \big( \sum_{m \in \mathcal{A}^t} \hat{\mathbf{g}}_m^t + \sum_{m \in \mathcal{B}^t} {\mathbf{g}}_m^t \big),    
% \end{align}





\subsection{Byzantine-Resilient Local Gradient Modification}
% \textbf{Vector manipulation:} 
% For FL systems under Byzantine attacks, the malicious workers could arbitrarily change the value of the local gradients. Therefore, the original vector modification scheme \eqref{vm} becomes inapplicable. This is because $\mathbf{v}_m^t$ is severely affected by the unbounded $\left(1-\lambda_m^t\right) \mathbf{g}_m^t$ and thus becomes unreliable.
In FL systems subjected to Byzantine attacks, malicious workers upload distorted gradients resulting from poisoned data. The vector modification in \eqref{vm} is inapplicable since $\mathbf{v}_m^t$ can be distorted by the unbounded term $\left(1-\lambda_m^t\right) \mathbf{g}_m^t$.
To alleviate the impact of malicious local updates on the global model, we design the modified gradient as \begin{align}\label{vm1}
    \mathbf{v}_m^t:=&(1-\lambda_m^t)\frac{\|\mathbf{r}^t\|}{\|\mathbf{g}_m^t\|}\mathbf{g}_m^t+\lambda_m^t\mathbf{r}^t ,\\
% \end{align}
% \begin{align} 
\label{dod new}
\text{with }  \quad  \lambda_m^t:=&c^t\big(1-\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\cdot \|\mathbf{r}^t\|}\big)\in[0,2c^t],
\end{align}
where $c^t \in [0,1], \forall t$ is set to be adjustable across training rounds, as described later in Section V-B.
% where $\lambda_m^t=:c^t\big(1-\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\cdot \|\mathbf{r}^t\|}\big)\in[0,2c^t]$, with $c^t \in [0,1], \forall t$ set to be adjustable across training rounds.


Unlike DRAG, which modifies the local updates by scaling the reference direction $\mathbf{r}^t$ to match the norm of $\mathbf{g}_m^t$, i.e.,  \eqref{vm}, 
BR-DRAG adopts a different scaling strategy that normalizes $\mathbf{g}_m^t$ to match $\| \mathbf{r}^t \|$; see \eqref{vm1}.
This modification is essential under Byzantine attacks, where the attackers may inflate the norm of their local updates in an attempt to dominate the aggregation process. By scaling $\mathbf{g}_m^t$ to $\|\mathbf{r}^t\|$, BR-DRAG mitigates the impact of malicious norm.
Moreover, with the DoD $\lambda_m^t$, BR-DRAG retains the advantages of divergence-based alignment from DRAG, providing directional correction under data heterogeneity.
% 
% \noindent\textbf{Remark 4.} 
% % With this reference direction coming from a trusted source that provides a general update direction of the training process and the module normalization operation, the modified DRAG can easily tackle the byzantine attack mentioned in Section III. For the attack that scales the module of the local gradient, we tackle it by normalizing each gradient to avoid the module anomaly; for the attack that reverses the direction, (\ref{vm1}) can automatically reverse it back since $1-\lambda_m^t<0$ if the direction diverges too much from the reference direction (with the value of $c$ properly selected), as illustrated in Fig. 2.
% \textit{
With the reference direction $\mathbf{r}^t$ obtained from a trusted dataset and indicating a general update direction for the training process, as well as the module normalization operation in \eqref{vm1}, BR-DRAG can effectively combat the malicious local updates $\hat{\mathbf{g}}_m^t$. 
% {\color{blue}To counter the gradients with arbitrary directions and magnitudes, we normalize each gradient to mitigate abnormal scaling, and leverage DRAG's DoD in \eqref{dod} to reverse excessively divergent gradients when $1-\lambda_m^t<0$.}


% To counter the attack that scales the magnitude of the local gradient, we employ normalization on each gradient to prevent any anomaly in the magnitude. As for the attack that reverses the direction, with (\ref{vm1}), DRAG can automatically correct it by reversing the gradient back if $1-\lambda_m^t<0$, signifying excessive divergence from the reference direction (with the value of $c$ appropriately chosen), which is also illustrated in Fig. 2.


\section{Convergence Analysis of DRAG and BR-DRAG}
In this section, we establish the convergence rate of DRAG and BR-DRAG for  prevalent non-convex objective functions.
% , which are prevalent in extensive machine learning tasks. 
% The convergence proof of DRAG for non-convex objective functions, as in most machine learning tasks, is rigorously established in this subsection, before which we need two assumptions, as stated below.
% \subsection{Assumption}
% This starts with the following assumptions.
We start with the following standard assumptions.

\smallskip
\smallskip

\noindent\textbf{Assumption 1. (Smoothness and lower boundedness)}  \textit{Each local objective 
function $ F_m(\boldsymbol{\theta})$ is $L$-Lipschitz smooth ($L>0$):
\begin{align}
    &\left\|\nabla F_m(\boldsymbol{\theta}_1)-\nabla F_m(\boldsymbol{\theta}_2)\right\| \leq L \left\|\boldsymbol{\theta}_1-\boldsymbol{\theta}_2\right\|, \, \forall m \in \mathcal{M}.  \label{lipschitz_gradient}
\end{align}
Also, assume the objective function $ f$ is lower-bounded by $ f^*$.}
% $\forall \boldsymbol{\theta}_1, \boldsymbol{\theta}_2\in \mathbb{R}^d$. 

\smallskip
\smallskip

\noindent \textbf{Assumption 2. (Unbiasedness and bounded variance)} \textit{For each client $m$, the local gradient estimator $\nabla F_m(\boldsymbol{\theta};z)$ is unbiased, i.e.,
\begin{align}
    \mathbb{E}\left[\nabla F_m(\boldsymbol{\theta};z)\right]=\nabla F_m(\boldsymbol{\theta}), \quad \forall m \in \mathcal{M}.
\end{align}
For both the variance of the local gradient estimator and that of the local gradient from the global one, there also exist two constants $\sigma_L^2$ and $\sigma_G^2$, such that
\begin{align}
    \mathbb{E}\left[\|\nabla F_m(\boldsymbol{\theta};z)-\nabla F_m(\boldsymbol{\theta})\|^2\right]&\leq \sigma_L^2, \quad \forall m \in \mathcal{M};\\
% \end{align}
% \begin{align}
    \mathbb{E}\left[\|\nabla F_m(\boldsymbol{\theta})-\nabla f(\boldsymbol{\theta})\|^2\right]&\leq \sigma_G^2, \quad \forall m \in \mathcal{M}.
\end{align}}

\smallskip
\smallskip

\textbf{Assumptions 1} and \textbf{2} are standard assumptions considered in non-convex optimization and FL convergence analysis~\cite{liang2019variance,li2019feddane,pathak2020fedsplit,mitra2021linear,karimireddy2020scaffold,acar2021federated,karimireddy2020mime,li2020federated,gao2022feddc,pmlr-v162-qu22a}.
Here, $\sigma_G^2$ quantifies the degree of data heterogeneity, and $\sigma_G^2=0$ indicates IID local data among the workers.

\subsection{Convergence Analysis of DRAG}

Under \textbf{Assumptions 1} and \textbf{2}, we establish the upper bound for the expectation of the average squared gradient norm:
% , $\frac{1}{T}\mathbb{E}\left[\sum_{t=0}^{T-1}\left\|\nabla f(\boldsymbol{\theta}^t)\right\|^2\right]$.
\begin{theorem} \label{drag theorem}
Suppose that $ F_m, \forall m \in \mathcal{M}$ is non-convex, and the stepsize is $\eta\leq \frac{1}{8LU}$. Given an initial model $\boldsymbol{\theta}^0$, after $T$ training rounds, DRAG guarantees
\begin{align} \label{drag convergence eq}
    \frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right] \leq \frac{f(\boldsymbol{\theta}^0)-f^*}{\gamma \eta U T} + V,
\end{align}
where $V=\frac{1}{\gamma\eta U} \big(  c \eta (2\sigma_L^2+9 \sigma_G^2)+\frac{\eta^2 U^2 L}{2}\left(\sigma_L^2+3 U \sigma_G^2\right) + (\frac{\eta^3 U^2 L^2(85 c+5)}{2}+\frac{15 \eta^4 U^3 L^3}{2})(\sigma_L^2+6 U \sigma_G^2) \big)  $, and $\gamma > 0 $.
% where $V=\frac{1}{\gamma}\Bigg[\frac{2c\sigma_L^2}{U}+\frac{\eta\sigma_L^2L}{2}+\left(\frac{3c}{U}+ \frac{\eta L}{2} \right)\big(15U^2L^2\eta^2V_1+3U\sigma_G^2)
% +\frac{5(1-c)L^2\eta^2}{2}V_1\big)\Bigg]$ with $V_1=\sigma_L^2+6U\sigma_G^2$, and $\gamma$ is a positive constant defined in \eqref{gamma}.
\end{theorem}

\begin{proof}
    % The proof is presented in Appendix C.
    See Appendix A.
\end{proof}

As revealed in \textbf{Theorem}~\textbf{\ref{drag theorem}}, DRAG achieves fast convergence for non-convex models under data heterogeneity and partial worker participation.
The non-vanishing term $V$ on the RHS of \eqref{drag convergence eq} characterizes the effects of both stochastic gradient variance ($\sigma_L^2$) and data heterogeneity ($\sigma_G^2$) across workers.
Meanwhile, the coefficients of $\sigma_L^2$ and $\sigma_G^2$  increase with the number of local updates, $U$.
To prevent these terms from growing unboundedly as $U$ increases, the stepsize can be set to $\eta = \mathcal{O}(\frac{1}{U})$ to ensure bounded variance accumulation.
% {\color{red}which is effectively controlled by letting the stepsize $\eta$ proportional to its inverse, i.e., $\mathcal{O}\left( \frac{1}{U} \right)$.}

% For DRAG, there exists a trade-off in selecting the hyperparameter $c$ in the DoD metric $\lambda_m^t$: values that are too small or too large can both adversely affect the performance.
In DRAG, the choice of the hyperparameter $c$ in the DoD metric $\lambda_m^t$ exhibits a trade-off, as setting $c$ too small or too large can adversely impact system performance.
% {\color{red}It is also noted that the hyperparameter $c$ in the metric of the DoD, $\lambda_m^t$, plays a non-monotonic role in $V$.}
A larger $c$ increases $\lambda_m^t$ in $\mathbf{v}_m^t$, leading to stronger alignment with the reference direction $\mathbf{r}^t$, reducing client drifts, as discussed in Section III-C.
% , and thus resulting in {\color{red}a smaller variance scaling in $V$, i.e.,} $\frac{5(1-c) \eta^3 L^2 U^2}{2}\left(\sigma_L^2+6 U \sigma_G^2\right)$.
Several terms in $V$, e.g., $c\eta(2\sigma_L^2+9\sigma_G^2)$, increase linearly with $c$, reflecting the susceptibility of $V$ to gradient variance under overly aggressive correction. 


% Furthermore, the final term involving $(1 - c)$ shows that overly small $c$ also leads to excess error by weakening the alignment effect and enlarging the local update deviation.


% {\color{red}Theorem 1 clearly states that provided the objective function is $L$-smooth
% and the global and local variances of the gradients are bounded, the proposed DRAG can then finally converge at a sublinear speed. Further, by selecting a stepsize $\eta$ with the form $\mathcal{O}(1/\sqrt{T})$, we can readily achieve a convergence rate of the level $\mathcal{O}(1/\sqrt{T})$. 
% }



% Note that the impact of the modifications of DRAG to the system setup can be detected in the expression of Theorem 1. As $c$ grows larger, the term $\frac{4c\sigma_L^2}{BU}+\frac{\eta\sigma_L^2L}{2B}+\left(\frac{5c}{U}+\frac{\eta L}{2}\right)\big(15U^2L^2\eta^2V_1+3U\sigma_G^2)$ in $V$ also increases, introducing a larger $V$; Meanwhile, as $c$ increases, the term $\frac{5(1-c)L^2\eta^2}{2}V_1$ decreases, which implies that the hyper-parameter $c$ has to be selected carefully to strike a balance.

\subsection{Convergence Analysis of BR-DRAG}
Next, we prove the convergence of BR-DRAG. The following assumptions,  characterizing the local update properties of the benign workers, are established.

\smallskip
\smallskip

\noindent\textbf{Assumption 3. (Gradient boundedness)} \textit{For the benign workers, the ratio between the norms of the global reference direction $\mathbf{r}^t$ and the local gradient $\mathbf{g}_m^t$ is bounded, i.e.,
% (Bounded Gradient-norm Similarity)
\begin{align}
    \rho_m^t = \frac{\left\|\mathbf{r}^t\right\|}{\left\|\mathbf{g}_m^t\right\|} \in [p ,q], \quad \forall m \in  \mathcal{B}^t.
\end{align}
where $0 < p < q$ are constants.}

\smallskip
\smallskip

\noindent\textbf{Assumption 4. (Bounded divergence)} \textit{For the benign workers, define the cosine similarities $ y_m^t = \frac{\langle\mathbf{g}_m^t, \mathbf{r}^t\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|} $.
The weighted aggregations of $ y_m^t $ is bounded, $\forall t$, i.e.,
% The cosine similarity between the reference direction $\mathbf{r}^t$ and the local gradients $\mathbf{g}_m^t$ is bounded, i.e.,
% In each round $t$, the degree of divergence between the averaged updated gradients from the attacked/unattacked workers and the reference direction is bounded, i.e.,
% \begin{align} \label{assumption 4-1}
%     x^t& = \frac{1}{S}\sum_{m \in \mathcal{A}^t} x_m^t \in [-w^t,0] ;\\
% % \end{align}
% % \begin{align} 
% \label{assumption 4-2}
%     \rho^t &= \frac{1}{S}\sum_{m \in \mathcal{B}^t} y_m^t \rho_m^t \in [0, (1-w^t)q] .
%     % \beta^t = \frac{1}{S}\sum_{m \in \mathcal{B}^t} \beta_m^t \in (0, 1-w^t], \ \forall m \in  \mathcal{B}^t.
% \end{align}}
\begin{align} 
    \rho^t &= \frac{1}{S}\sum_{m \in \mathcal{B}^t} y_m^t \rho_m^t \in [0, (1-w^t)q]; \label{assumption 4-2} 
    % \\ x^t& = \frac{1}{S}\sum_{m \in \mathcal{A}^t} x_m^t \in [-w^t,0] . \label{assumption 4-1}
\end{align}
where $w^t \in [0,1]$ denotes the intensity level of the Byzantine attack defined in Section II-B.} 


\smallskip
\smallskip

Unlike the typical bounded gradient assumptions, e.g.,~\cite[Asm. 4]{li2019convergence} and~\cite[Asm. 5]{zuo2024byzantine}, \textbf{Assumption~3} only requires a bounded ratio between the trusted reference direction $\mathbf{r}^t$ and the benign local gradients $\mathbf{g}_m^t, \forall m$, without restricting their individual magnitudes, relaxing constraints on modeling data heterogeneity, and accommodating diverse data distributions, compared to the typical requirement of uniform gradient norms (i.e., $\|\mathbf{g}_m^t\|\leq G$~\cite{li2019convergence,zuo2024byzantine}). This relaxation better reflects practical federated settings with non-IID data.

% {\color{red}Specifically, .... (Please explain how to benefit from the flexibility of data heterogeneity modeling).}
% In \textbf{Assumption 4}, \eqref{assumption 4-2} constrains the aggregated alignment of the benign workers’ gradients with the reference direction, ensuring that a sufficient proportion of participating workers provide updates consistent with the trusted signal.
% By contrast, \eqref{assumption 4-1} captures the adversarial nature of malicious updates by enforcing a negative or non-positive average cosine similarity with respect to the reference direction, where corrupted gradients tend to reverse or mislead the optimization direction.
In \textbf{Assumption~4}, \eqref{assumption 4-2} enforces that the weighted average of the cosine similarities between the benign workers’ gradients and trusted reference direction $\mathbf{r}^t$, scaled by their gradient-norm ratios $\rho_m^t$, remains within~$[0, (1-w^t)q]$. 
In other words, the benign updates are assumed to be positively aligned with $\mathbf{r}^t$ and have bounded relative magnitudes. The benign contribution to the aggregation maintains a consistent descent direction without being dominated by excessively large updates. 

Interestingly, we do not assume explicit bounds or range for the cosine similarities of the malicious workers $m \in \mathcal{A}^t$. The only assumption about the Byzantine attacks is the intensity level of attacks, i.e., $w^t$ in \eqref{assumption 4-2}. This substantially relaxes the reliance of our analysis and findings on assumptions of a-priori knowledge about the system.
% Conversely, \eqref{assumption 4-1} stipulates that the average cosine similarity of the malicious workers to the reference direction $\mathbf{r}^t$ lies within~$[-w^t, 0]$, explicitly capturing their adversarial tendency to produce updates that are orthogonal or oppositely oriented to $\mathbf{r}^t$. 
% Together, these bounds \eqref{assumption 4-1} and \eqref{assumption 4-2} provide a clear separation between benign and malicious gradient behaviors.


Under \textbf{Assumptions 1}–\textbf{4}, the following theorem establishes.
% \begin{theorem} \label{brdrag theorem}
% Suppose that $ F_m, \forall m \in \mathcal{M}$ is non-convex, and the stepsize satisfies $\eta \leq \frac{(1-\tilde{w})(1-\tilde{c})p}{32U+2U L^2(p+2q^2)}$. Given an initial model $\boldsymbol{\theta}^0$, after $T$ training rounds, BR-DRAG guarantees
% \begin{align} \label{brdrag theorem convergence}
%     \frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right] \leq \frac{f(\boldsymbol{\theta}^0)-f^*}{\chi \eta U T} + W,
% \end{align}
% where $\chi>0$, $\tilde{c} = \max _t \frac{w^t}{w^t - x^t}$, $\tilde{w}=\max _t w^t$, and $W =   \frac{2\eta \sigma_L^2}{\chi}\left(\frac{\left(1+\tilde{w}\right) \tilde{c} \eta U L^2}{2}+\frac{\eta U L^2 q^2}{2\left(1-\tilde{c}\right) p}+1\right) $.
% \end{theorem}

\begin{theorem} \label{brdrag theorem}
Suppose that $ F_m, \forall m \in \mathcal{M}$ is non-convex.
Given an initial model $\boldsymbol{\theta}^0$, after $T$ training rounds,
\begin{itemize}
    \item if $c^t = \frac{w^t}{w^t - x^t} \in [\frac{1}{2},1]$ and the stepsize satisfies $\eta \leq \frac{(1-\tilde{w})(1-\tilde{c})p}{32U+2U L^2(p+2q^2)}$, BR-DRAG guarantees
    \begin{align} \label{brdrag theorem convergence}
    \frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right] \leq \frac{f(\boldsymbol{\theta}^0)-f^*}{\chi \eta U T} + W,
    \end{align}
    where $\chi>0$, $\tilde{c} = \max _t \frac{w^t}{w^t - x^t}$, $\tilde{w}=\max _t w^t$, and $W =   \frac{2\eta \sigma_L^2}{\chi}\left(\frac{\left(1+\tilde{w}\right) \tilde{c} \eta U L^2}{2}+\frac{\eta U L^2 q^2}{2\left(1-\tilde{c}\right) p}+1\right) $.

    \item if $\eta$ and $c^t$ satisfy $\eta \leq \frac{2c^t(1-c^t)(1-w^t)(1-w^t-y^t)p}{32 p U+ c^t U L^2(p^2+q^2)}$ and $\left( (\rho^t)^2 \!-\! 7w^t \!+\! 8x^t \!-\! 1 \right)(c^t)^2 \!+\! (15w^t \!+\! 1)c^t  \geq \frac{8(w^t)^2}{w^t-x^t}, \forall t$, respectively, BR-DRAG guarantees
    \begin{align} \label{brdrag theorem convergence}
    \frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right] \leq \frac{f(\boldsymbol{\theta}^0)-f^*}{\chi \eta U T} + W,
    \end{align}
    where $\chi>0$, $\tilde{c}=\max_t c^t$, $\tilde{w}=\max _t w^t$, $x^t = \frac{1}{S}\sum_{m \in \mathcal{A}^t} x_m^t$ with $x_m^t=\frac{\langle\mathbf{g}_m^t, \mathbf{r}^t\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}$, $y^t=\frac{1}{S} \sum_{m \in \mathcal{B}^t } y_m^t$, and $W =   \frac{2\eta \sigma_L^2}{\chi}\left(\frac{\left(1+\tilde{w}\right) \tilde{c} \eta U L^2}{2}+\frac{\eta U L^2 q^2}{2\left(1-\tilde{c}\right) p}+H+1\right) $ with $H = \max_{t} \frac{\| w^t(1-c^t)+c^t x^t \|}{2c^t\eta U (w^t-x^t) } $.
\end{itemize}
\end{theorem}
\begin{proof}
    See Appendix B.
\end{proof}

According to \textbf{Theorem}~\textbf{\ref{brdrag theorem}}, BR-DRAG converges under Byzantine attacks in the presence of heterogeneity and partial aggregation, unlike conventional methods, e.g,~\cite{li2019rsa,fang2020local,blanchard2017machine}.
The constant $W$ specifies the non-vanishing error caused by gradient variance, data heterogeneity, and adversarial attacks.
In each training round $t$, the malicious updates induce the increased errors $ \delta^t =  { c^t  (w^t-x^t) \eta U}$ and $\kappa^t = c^t (1-w^t-y^t)\eta U$ in \eqref{one round convergence new}, impeding convergence. 
Based on \eqref{vm1}, BR-DRAG projects the malicious vectors on a bounded space $\mathbf{r}^t$, i.e., the trusted reference direction, tolerating the malicious updates.

The resistance of BR-DRAG to malicious worker participation is achieved by reducing the stepsize $\eta$ and enhancing the gradient alignment degree. As the proportion of malicious workers increases (i.e., $w^t$ grows), the hyperparameter $c^t$ needs to increase, which in turn reduces the effective stepsize. 
% This adjustment sacrifices the convergence speed to ensure that the conditions for stability and robustness are still satisfied.
This adjustment slows convergence to maintain the conditions necessary for stability and robustness.
An excessively large $c^t$ amplifies the variance-related components, e.g., $\frac{\tilde{c}(1+\tilde{w})L^2}{\chi}\eta^2 U \sigma_L^2$ in $W$ in BR-DRAG, while an excessively small $c^t$ may degrade gradient alignment in \eqref{vm1}, heightening the vulnerability to adversarial gradients. To this end, $c^t$ should be calibrated meticulously to suppress adversarial deviation while controlling gradient variance in BR-DRAG.


% This leads to a trade-off: increasing $c$ improves robustness to malicious directions but risks inflating sensitivity to gradient randomness.
% A larger $c$ enhances alignment by giving more weight to the reference direction, thereby suppressing the influence of adversarial updates and reducing the corresponding error terms $\delta^t$ and $\kappa^t$. 
% However, as in DRAG, an overly large $c$ amplifies variance-related components in $W$, such as those arising from stochastic gradient noise. 
% This leads to a trade-off: increasing $c$ improves robustness to malicious directions but risks inflating sensitivity to gradient randomness.
% Thus, BR-DRAG achieves Byzantine resilience by balancing directional correction strength and variance accumulation through careful tuning of $c$.


% To suppress the adversarial errors $\delta^t$ and $\kappa^t$ introduced by malicious updates, BR-DRAG adopts a carefully designed update rule that combines cosine similarity-based weighting and norm-bounded interpolation. As shown in the proof (Appendix C), the modified gradient $\mathbf{v}_m^t$ for each worker is expressed as a convex combination of the normalized local gradient and the trusted reference direction. This design ensures that even 

% Specifically, the adversarial terms $\delta^t$ and $\kappa^t$—which appear in the convergence decomposition as upper bounds on inner product misalignment and curvature deviation—are both proportional to the similarity-weighted correction coefficient $c$. 

% The first effect of increasing $c$ is positive: by enlarging the interpolation weight $h_m^t = c(1 - \alpha_m^t)$ for malicious clients and $l_m^t = c(1 - \beta_m^t)$ for benign ones, BR-DRAG promotes tighter alignment with $\mathbf{r}^t$, reducing the misalignment between aggregated gradients and the true descent direction. This tightens the negative drift term in $D_1^t$ and reduces its adverse effect on convergence.

% However, the second-order effect of $c$ is more subtle: larger $c$ also magnifies several variance-related components in $D_2^t$, such as $\frac{c \eta U L^2}{2}$ and $\frac{1}{2 c \eta U}$, forming a non-monotonic trade-off. When $c$ is too small, malicious updates are insufficiently corrected, allowing $\delta^t$ and $\kappa^t$ to grow and weaken the descent term. When $c$ is too large, these correction terms can dominate the residual noise in $W$, resulting in higher steady-state error. Thus, BR-DRAG achieves Byzantine robustness through a tunable balance: $c$ must be calibrated to jointly suppress adversarial deviation while controlling noise amplification.

% To mitigate such Byzantine-induced errors, BR-DRAG introduces two core mechanisms: a server-generated reference direction that remains uncontaminated, and a divergence-based update rule that adaptively modulates each client's contribution according to its directional agreement with the reference. Specifically, the modification formula in \eqref{vm1} leverages the dissimilarity metric $\lambda_m^t$ to interpolate between the potentially corrupted local gradient and the trusted reference, where the parameter $c$ governs the interpolation strength. 



Unlike a majority of the existing Byzantine-resilient FL algorithms~\cite{li2019rsa,cao2020fltrust,9721118,9153949,10100920,zuo2024byzantine}, which typically require the condition of $A^t < \frac{S}{2}, \forall t$,
% the number of malicious workers be less than half of the participating set (i.e., $A^t < S/2$),
BR-DRAG allows an arbitrary proportion of malicious workers in each round. 
This is enabled by the design of the reference direction in \eqref{trust} and gradient modification in \eqref{vm1}, suppressing the effects of malicious gradient aggregation.
Moreover, by plugging $\eta = \mathcal{O}(\frac{1}{\sqrt{T}})$ and $U = \mathcal{O}(1)$ into \eqref{brdrag theorem convergence}, BR-DRAG maintains the same convergence rate of $\mathcal{O}(\frac{1}{\sqrt{T}})$ as DRAG, which is consistent with FL algorithms without Byzantine attacks~\cite{li2020federated}.



% The algorithm ensures that even in rounds dominated by malicious clients ($w^t \rightarrow 1$), the trusted reference direction remains intact and benign updates are adaptively amplified to preserve global convergence. 



\section{Experiments and Results}

% In this section, we illustrate the advantages of DRAG by comparing its performance against state-of-the-art algorithms including SCAFFOLD \cite{karimireddy2020scaffold}, AdaBest \cite{varno2022adabest}, FedProx \cite{li2020federated}, and the vanilla FedAvg. These algorithms are tested on the EMNIST dataset and the more challenging CIFAR-10 dataset with partial and full worker participation to provide a more comprehensive perspective into the performances of these algorithms. We also on byzantine attack on these two datasets. 




% Our evaluation is carried out on both the EMNIST dataset and the CIFAR-10 dataset, with scenarios of both partial and full worker participation. Additionally, we assess the robustness of the algorithms against byzantine attacks on these two datasets. 

This section experimentally evaluates DRAG and BR-DRAG compared to state-of-the-art FL and Byzantine-resilient FL algorithms, with the following datasets and models:
\begin{itemize}
    \item EMNIST dataset~\cite{cohen2017emnist}: For the
     digit and letter classification task, we utilize the ``balanced" data split containing 131,600 samples with 47 classes. The designed Convolutional Neural Network (CNN) comprises two $5 \times 5$ convolutional layers and two fully connected layers, with a final output layer for 47-class classification.
     % There are in total six different splits in the dataset and we used the ``balanced'' data split in this work with 47 balanced classes of data. Each piece of data is a $28\times 28$ grey image and there are 131,600 characters in the 47 balanced classes.
     % The EMNIST dataset is an extended version of the famous MNIST dataset. Besides the handwritten digits from the MNIST dataset, the EMNIST dataset also includes handwritten letters. 
    
    \item CIFAR-10 dataset~\cite{krizhevsky2009learning}: 
    For the image classification task, the CIFAR-10 dataset is composed of 60,000 color images in a total of 10 classes. We employ a CNN with two $5 \times 5$ padded convolutional layers to extract hierarchical features and predict across 10 classes.
    % with 6000 images in each class. There are 50000 images for training and 10000 for testing. Since it is a color image dataset, it is supposed to be more challenging than the EMNIST dataset.

    \item CIFAR-100 dataset~\cite{krizhevsky2009learning}:
    For the fine-grained image classification task, the CIFAR-100 dataset contains 60,000 color images categorized into 100 classes.
    We adopt a deeper CNN comprising three $3 \times 3$ padded convolutional layers followed by max pooling operations, and two fully connected layers, culminating in a 100-class output layer.
\end{itemize}
Consider an FL system with $40$ workers, and each worker performs $U=5$ local updates in each training round.
To characterize data heterogeneity, we utilize the Dirichlet distribution with parameter $\beta$~\cite{pmlr-v162-qu22a,jhunjhunwala2023fedexp,Kim_2024_CVPR}. Specifically, we sample $p_k \sim \operatorname{Dir}(\beta)$ and allocate a proportion $p_{k,j}$ of image class $k$ instances to device $j$. A smaller $\beta$ represents a more imbalanced data distribution. Here, we choose $\beta=0.1$ and $0.5$ to represent strong and moderate data heterogeneity, respectively.
By default, we set the stepsize $\eta = 0.01$ and the batch size is 10.
To simulate partial worker participation, we set the worker subset size to $S=10$ per round. The FL environment is set up in PyTorch 2.5.1 on a Windows 10 operating system with one 4090 GPU.

% As in \cite{cao2020fltrust}, we define the heterogeneity of data as follows. Any training data with label $\ell$ is assigned to client $\ell \bmod M$ with probability $q$ and to any other client with probability $\frac{1-q}{M-1}$. In dealing with client-drift, the high heterogeneous data distribution is implemented by setting $q=1$. In countering the byzantine attack, we define low data heterogeneity as $q=\frac{1}{M}$ and high data heterogeneity as $q=1$.


\subsection{Evaluation of DRAG}

We first compare DRAG with the following benchmarks, which have been successfully deployed in FL systems.
% by conducting a comprehensive performance comparison against state-of-the-art algorithms, including SCAFFOLD \cite{karimireddy2020scaffold}, AdaBest \cite{varno2022adabest}, FedProx \cite{li2020federated}, and the vanilla FedAvg \cite{pmlr-v54-mcmahan17a}. 
\begin{itemize}
    \item FedAvg~\cite{pmlr-v54-mcmahan17a}: The local and global models are updated as \eqref{update} and \eqref{aggregation}, respectively.

    % \item AdaBest~\cite{varno2022adabest}

    \item FedProx~\cite{li2020federated}: In any training round $t$, the workers perform $U$ local updates with a regularization term to approximate a $\gamma$-inexact solution: $\left\|\nabla F_m(\boldsymbol{\theta}_{m}^{t,u})+\mu(\boldsymbol{\theta}_{m}^{t,u}-\boldsymbol{\theta}^t)\right\| 
        \leq \gamma\left\|\nabla F_m(\boldsymbol{\theta}^t)\right\|$.
        
    \item SCAFFOLD~\cite{karimireddy2020scaffold}:
    In any training round $t$, the local update is employed with control variates: $\boldsymbol{\theta}_{m}^{t,u}=\boldsymbol{\theta}_{m}^{t,u}-\eta ( \nabla F_m(\boldsymbol{\theta}_{m}^{t,u}; z_{m}^{t,u})  - h_{m}^t + h^t )$, where $h_{m}^t$ and $h^t$ are updated as $h_{m}^{t+1}=\nabla F_m (\boldsymbol{\theta}^{t}; z_{m}^{t,0} ) $ and $h^{t+1}=h^{t}+\frac{1}{M}\sum_{i \in \mathcal{S}^t} (\nabla F_m (\boldsymbol{\theta}^{t}; z_{m}^{t,0} )-h_{m}^t  )$.

    \item FedExP~\cite{jhunjhunwala2023fedexp}: With $\tilde{\Delta}^t = \frac{1}{S}\sum_{i \in \mathcal{S}^t}\mathbf{g}_m^t$ and a small positive constant $\epsilon$, the global model is updated as $\boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^t+\frac{\eta_g^t}{S} \sum_{m \in \mathcal{S}^t} \mathbf{g}_m^t$, where $\eta_g^t=\max \Big\{1, \frac{\sum_{i \in \mathcal{S}^t}\left\|\mathbf{g}_m^t\right\|^2}{2 S\left(\left\|\tilde{\Delta}^t\right\|^2+\epsilon\right)}\Big\}$.   
    
    % \item FedSAM~\cite{pmlr-v162-qu22a}: SAM optimizers are utilized at the workers' side: ${\boldsymbol{\theta}}_m^{t,u+1} = {\boldsymbol{\theta}}_m^{t,u} - \eta \nabla F_m(\tilde{\boldsymbol{\theta}}_{m}^{t,u}; z_{m}^{t,u}) $, where $\tilde{\boldsymbol{\theta}}_m^{t,u}$ is obtained by $\tilde{\boldsymbol{\theta}}_m^{t,u} = \boldsymbol{\theta}_m^{t,u} + \rho \frac{\nabla F_m(\boldsymbol{\theta}_{m}^{t,u}; z_{m}^{t,u})}{\|\nabla F_m(\boldsymbol{\theta}_{m}^{t,u}; z_{m}^{t,u})\|}$.

    \item FedACG~\cite{Kim_2024_CVPR}: The local updates are performed as 
    $\boldsymbol{\theta}_m^{t,u+1} = \boldsymbol{\theta}_m^{t,u} - \eta ( \nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m}^{t,u}) + \beta \| \boldsymbol{\theta}_m^{t,u} - (\boldsymbol{\theta}^{t-1} + \lambda m^{t-1} ) \| )$ with the lookahead gradient updated as $m^t = \lambda m^{t-1} + \frac{1}{S} \sum_{m \in \mathcal{S}^t} \mathbf{g}_m^t $.


\end{itemize} 
After fine-tuning these benchmarks, we set $\mu=0.2$ for FedProx; $\epsilon=0.001$ for FedExP; $\beta=0.2$ and $\lambda=0.85$ for FedACG. For DRAG, $\alpha=0.25$, $c=0.1$ for moderate data heterogeneity, and $c=0.25$ for strong data heterogeneity.

Figs.~\ref{fig_EMNIST_DRAG} to \ref{fig_CIFAR100_DRAG} show the convergence of the algorithms. Across different datasets and data heterogeneity levels, DRAG significantly outperforms the benchmarks, validating \textbf{Theorem~1}. 
The accuracy gap between DRAG and FedAvg is larger when data heterogeneity is more severe (from $\beta=0.5$ to $\beta=0.1$), confirming the effective client-drift mitigation of DRAG.
% For example, as shown in Fig.~\ref{fig_CIFAR10_DRAG}(a) and \ref{fig_CIFAR100_DRAG}(a), DRAG achieves 3.8\% and 4.2\% accuracy improvement over FedAvg, respectively. 
Meanwhile, DRAG converges considerably faster than the benchmarks, and achieves 70\% test accuracy within 600 rounds in Fig.~\ref{fig_CIFAR10_DRAG}(b). Over 1,400 rounds are needed for other algorithms, e.g., SCAFFOLD and FedExP, verifying that the adaptive aggregation of DRAG in \eqref{vm} accelerates the decrease in the global loss function.
% When the data distributions tend to be IID, DRAG approaches FedAvg in accuracy, as the local and global updates follow nearly the same direction, making the modified gradient in \eqref{vm} degenerate into $\mathbf{g}_m^t$.

FedProx demonstrates subpar performance due to the lack of enough local computations to achieve the required $\gamma$-inexact solution. With the additional lookahead gradient, FedACG achieves a limited accuracy improvement over FedProx. FedExP is limited by the later stage of training and degrades to FedAvg.
Since the control variate $\nabla F_m(\boldsymbol{\theta}^t) - \mathbb{E}[h_m^t] + \mathbb{E}[h^t] \neq 0$ when $\boldsymbol{\theta}^t=\boldsymbol{\theta}^*$, the local models of SCAFFOLD deviate from $\boldsymbol{\theta}^*$ and necessitate additional training rounds to mitigate this imbalance of control variate.
By balancing the local gradient direction and the reference direction, DRAG avoids using control variates or regularization terms, realizing adaptive model training without additional communication overhead.

% Its success can be attributed to its heuristic and intuitive design, which avoids using control variates or regularization terms. Instead, DRAG leverages tunable hyperparameters $c$ and $\alpha$ to balance the weights of the local gradient direction and the reference direction, respectively. This adaptability enables DRAG to develop a more accurate reference direction, further contributing to its superior performance.

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{ new_figs/emnist_01.pdf}
		\label{fig_EMNIST_1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{ new_figs/emnist_05.pdf }
		\label{fig_EMNIST_2}}
	\caption{Convergence performance of different algorithms on EMNIST.}
	\label{fig_EMNIST_DRAG}
\end{figure}
\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{ new_figs/cifar10_01.pdf}
		\label{fig_CIFAR10_1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{  new_figs/cifar10_05.pdf}
		\label{fig_CIFAR10_2}}
	\caption{Convergence performance of different algorithms on CIFAR-10.}
	\label{fig_CIFAR10_DRAG}
\end{figure}
\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{ new_figs/cifar100_01.pdf}
		\label{fig_CIFAR100_1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{ new_figs/cifar100_05.pdf}
		\label{fig_CIFAR100_2}}
	\caption{Convergence performance of different algorithms on CIFAR-100.}
	\label{fig_CIFAR100_DRAG}
\end{figure}

% \begin{figure}[hbt]
% \centering
% \subfloat[full participation]{\includegraphics[width=1.5in]{ACAG_full_EMNIST.pdf}}
% \subfloat[partial participation]{\includegraphics[width=1.5in]{ACAG_partial_EMNIST.pdf}}
% \caption{Performance comparison of DRAG with state-of-the-art algorithms under the EMNIST dataset with full and partial worker participation.}
% \end{figure}

% Fig. 3(a) and (b) plot the performance comparison of different algorithms under the EMNIST dataset ($q=1$), with full and partial worker participation, respectively. To be specific, there are in total $M=40$ workers and $S=10$ for partial worker participation. Each worker performs $U=5$ local updates. All the algorithms stop updating once the test accuracy reaches 80\%. The stpdfize $\eta$ is 0.1. We used a two-layer fully connected network with 500 hidden units. The hyperparameters of the algorithms are listed below. For AdaBest, $\mu=0.02, \beta=0.8$; 

% EMNIST for DRAG, $c=0.25$, $\alpha=0.6$ for partial participation and $\alpha=1$ for full participation; $\mu=0.2$ for FedProx.




% from Fig. 3 that AdaBest has the fastest convergence speed in the first stage with full participation but the final convergence accuracy is lower than the other algorithms. This is mainly because Adabest reduces the weight of historical information when utilizing control variates compared to SCAFFOLD. While SCAFFOLD outperforms AdaBest, FedAvg, and FedProx both under full and partial participation. However, DRAG is consistently the best for the two cases, thanks to the momentum-based reference direction that provides a relatively accurate training direction and the vector manipulation operation.



% In Fig. 4, we test the performances of these algorithms on the CIFAR-10 dataset ($q=1$) with $M=20$ workers for full participation and $S=5$ workers for partial participation. Each worker performs $U=5$ local updates. All the algorithms stop updating once the test accuracy reaches 70\%. The stpdfize $\eta$ is 0.1. We used a CNN with two convolution layers and three fully-connected layers. For AdaBest, $\mu=0.02, \beta=0.2$; for DRAG, $c=0.1$, $\alpha=0.2$ for partial participation and $\alpha=1$ for full participation; $\mu=0.2$ for FedProx.

% In this setting, AdaBest still performs relatively better with partial participation; SCAFFOLD and FedProx do not have much improvement compared to FedAvg, as observed in \cite{li2021federated}. However, DRAG achieves distinctly better performance than other algorithms. To be specific, it only needs half of the training rounds to reach 70\% test accuracy compared to FedAvg with partial participation and the proportion becomes 1/4 in full participation.

Fig.~\ref{fig_diff_S_DRAG} evaluates DRAG under different numbers of participating workers, where a UAR selection of $S =5$, $15$, $25$, and $35$ out of $M=40$ workers is considered.
Clearly, more worker participation leads to faster convergence under both data distributions.
A larger number of selected workers $S$ contributes to more stable reference directions $\mathbf{r}^t, \forall t$, in \eqref{reference_cf} and helps better guide local gradient modifications.
\begin{figure}[t!]
\centering
\subfloat[\centering Accuracy on CIFAR-10 with $\beta=0.1$]{\includegraphics[width=1.7in]{new_figs/diff_S_cifar10_01.pdf}}
\subfloat[\centering Accuracy on CIFAR-10 with $\beta=0.5$]{\includegraphics[width=1.7in]{new_figs/diff_S_cifar10_05.pdf}}
\caption{Convergence performance of DRAG with different numbers of participating workers on CIFAR-10.}
\label{fig_diff_S_DRAG}
\end{figure}

Figs. \ref{diff_alpha_DRAG} and \ref{diff_c_DRAG} examine the impact of the hyperparameters $\alpha$ and $c$ in \eqref{reference_cf} and \eqref{dod} on the performance of DRAG.
Fig.~\ref{diff_alpha_DRAG}(a) demonstrates an overly small $\alpha=0.01$ in \eqref{reference_cf} overuses historical gradients, where unstable gradients in the early training stage hinder convergence. When $\alpha$ is large, e.g., $\alpha > 0.25 $, the reference direction $\mathbf{r}^t$ depends excessively on $\Delta^{t-1}$, making it difficult to find a fast descent path.
In Fig.~\ref{diff_alpha_DRAG}(b), although the data heterogeneity becomes moderate, an overly large or small value of $\alpha$ slows down the convergence.
Fig.~\ref{diff_c_DRAG} shows that too large or too small $c$ significantly impairs training, since it would penalize gradient variance control or gradient alignment.
% , which leads to a moderate value setting.
This is consistent with the analysis in Section V-B, confirming the importance of careful configuration of $c$. 


\begin{figure}[t!]
\centering
\subfloat[\centering $\beta=0.1$]{\includegraphics[width=1.7in]{new_figs/diff_alpha_cifar10_01.pdf}}
\subfloat[\centering $\beta=0.5$]{\includegraphics[width=1.7in]{new_figs/diff_alpha_cifar10_05.pdf}}
\caption{Convergence performance of DRAG with different values of hyperparameter $\alpha$ on CIFAR-10.}
\label{diff_alpha_DRAG}
\end{figure}

\begin{figure}[t!]
\centering
\subfloat[\centering $\beta=0.1$]{\includegraphics[width=1.7in]{new_figs/diff_c_cifar10_01.pdf}}
\subfloat[\centering $\beta=0.5$]{\includegraphics[width=1.7in]{new_figs/diff_c_cifar10_05.pdf}}
\caption{Convergence performance of DRAG with different values of hyperparameter $c$ on CIFAR-10.}
\label{diff_c_DRAG}
\end{figure}

% {\color{red}DRAG exhibits significant superiority, especially in the presence of full participation. This advantage is attributed to the momentum-based reference direction, which provides a more accurate training direction, and the vector manipulation operation, which contributes to effective alignment with the reference direction.}

% In this setting, AdaBest still performs relatively better with partial participation; however, SCAFFOLD and FedProx show limited improvement compared to FedAvg, as previously observed in \cite{li2021federated}. On the other hand, DRAG outperforms all other algorithms significantly. Specifically, it achieves 70\% test accuracy with only half of the training rounds compared to FedAvg with partial participation, and this proportion further reduces to 1/4 under full participation.



% To sum up, with the proposed heuristic and intuitive design of algorithm instead of using control variates or regularization terms, DRAG is consistently the best algorithm to deal with client drift under various settings. It has strong adaptability thanks to the tunable hyper-parameters $c$ and $\alpha$, which are designed to balance the weights of the direction of the local gradient and the reference direction, and to develop a better reference direction, respectively.



\subsection{Evaluation of BR-DRAG}

We proceed to evaluate the BR-DRAG framework under Byzantine attacks by considering the following benchmarks.
\begin{itemize}
    \item FedAvg~\cite{pmlr-v54-mcmahan17a}: The local and global models are updated following \eqref{update} and \eqref{byzantine aggregation}, respectively.

    \item FLTrust~\cite{cao2020fltrust}: In any training round $t$, the local gradients are modified with ReLU-clipped cosine similarity: $\tilde{\mathbf{g}}_m^t \!=\! \max \! \left(0, \cos \left(\mathbf{g}_m^t, \mathbf{r}^t\right)\right) \! \left\|\mathbf{g}_m^t\right\| \! \frac{\mathbf{r}^t}{\left\|\mathbf{r}^t\right\|}$, where $\mathbf{r}^t$ is the reference direction designed in a similar way to \eqref{trust}.
    
    \item RFA~\cite{9721118}: The global model is updated via geometric median aggregation: $\boldsymbol{\theta}^{t+1} = \operatorname{GeoMed}\left( \{ \boldsymbol{\theta}_m^{t,U} \}_{m \in \mathcal{S}^t} \right)$, where $\operatorname{GeoMed}(\cdot)$ minimizes the sum of distances to all local models and is solved with the Weiszfeld algorithm~\cite{beck2015weiszfeld}.

    \item RAGA~\cite{zuo2024byzantine}: The global gradient $\mathbf{g}^t$ is acquired with the geometric median metric: $\mathbf{g}^t \!=\! \operatorname{GeoMed}\left( \{ \mathbf{g}_m^{t} \}_{m \in \mathcal{S}^t} \right)$, and the global model is updated as $\boldsymbol{\theta}^{t+1}=\boldsymbol{\theta}^{t} + \mathbf{g}^t$.
\end{itemize}
For malicious workers, we deploy three different attack methods, including noise injection~\cite{9614992}, sign flipping~\cite{li2019rsa}, and label flipping~\cite{10054157}; see Section I-B. For the noise injection attack, we set malicious local gradients to $p_m^t\mathbf{g}_m^t, \forall m \in \mathcal{A}^t$, where $p_m^t$ conforms to the Gaussian distribution $\mathcal{N}(0,3)$. For the sign flipping attack, the malicious workers generate $-\mathbf{g}_m^t, \forall m \in \mathcal{A}^t$ as the local gradient. For the label flipping attack, half of local data labels of the malicious workers are reversed. Under BR-DRAG, $c^t=0.5, \forall t$, and $\mathcal{D}_{\text{root}}$ is constructed with 3,000 data samples drawn randomly and uniformly in equal proportions from trusted workers; see Section IV-B.
% {\color{blue}construct $\mathcal{D}_{\text{root}}$ with 3,000 data samples drawn randomly and uniformly in equal quantities from each worker.}

We set the ratio of malicious workers to $ \frac{A}{M} = 30\%$ in the system.
Figs.~\ref{fig_CIFAR10_BRDRAG_NI} to \ref{fig_CIFAR100_BRDRAG_LF} evaluate the training performance of BR-DRAG under the different Byzantine attack methods, datasets, and data heterogeneity levels.
% It is observed that BR-DRAG outperforms the benchmarks significantly, validating its effectiveness.
Compared to the benchmarks, BR-DRAG demonstrates superior adaptability and robustness against severe data heterogeneity and attacks, validating \textbf{Theorem 2}. 
For example, BR-DRAG surpasses FLTrust in accuracy by 13.4\% and 9.8\% in Figs.~\ref{fig_CIFAR10_BRDRAG_SF}(a) and~\ref{fig_CIFAR10_BRDRAG_SF}(b), respectively.
Based on \eqref{vm1} and using a larger $c$ than DRAG, BR-DRAG effectively suppresses malicious gradients while amplifying the impact of benign updates. 

On the other hand, no benchmarks maintain stability under the diverse Byzantine attacks. RFA and RAGA can approach BR-DRAG in performance since the applied Weiszfeld algorithm can resist the interference of extreme gradient values. The ReLU-clipping design in FLTrust lacks fine-grained control, as it would entirely preclude the attackers or workers with excessively biased local data. 
In contrast, a key advantage of BR-DRAG is that it retains only the directionally consistent component of each received update with respect to $\mathbf{r}^t$, and suppresses misalignment.
% As shown in Figs.~\ref{fig_CIFAR10_BRDRAG_NI_more_attacker} to \ref{fig_CIFAR10_BRDRAG_LF_more_attacker}, we increase the proportion of Byzantine workers within the system to $\frac{A}{M}=60\%$. Unlike other algorithms, which need the assumption of $\frac{A}{M} < 50\%$, BR-DRAG still maintains convergence, aligning with the statement in Section V-C. 
As shown in Figs.~\ref{fig_CIFAR10_BRDRAG_NI_more_attacker} to \ref{fig_CIFAR10_BRDRAG_LF_more_attacker}, even when the proportion of Byzantine workers increases to $\frac{A}{M}=60\%$, BR-DRAG maintains convergence. This surpasses all benchmarks, which typically require $\frac{A}{M} < 50\%$, validating the analysis in Section V-C.
% The benchmarks based on geometric median are particularly impaired, as the excessive number of malicious nodes drastically affects the determination of centroids. 
The benchmarks based on the geometric median, i.e., RFA and RAGA, suffer severe degradation, as the high proportion of malicious workers significantly distorts the centroid estimation.
% In contrast, BR-DRAG mitigates the risk of complete non-convergence by regulating both the magnitude and direction of malicious updates. This further confirms the unique advantages of the proposed algorithm under high data heterogeneity and severe Byzantine attacks.
In contrast, BR-DRAG converges by regulating the magnitude and direction of malicious updates, demonstrating its distinctive robustness under high data heterogeneity and severe Byzantine attacks.

% -------------------------------------------------------
% \begin{figure}[!t]
% 	\centering
% 	\subfloat[{\centering Accuracy on EMNIST with $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_emnist_01.pdf}
% 		\label{fig_EMNIST_BRDRAG_NI1}}
% 	\subfloat[{\centering Accuracy on EMNIST with $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_emnist_05.pdf}
% 		\label{fig_EMNIST_BRDRAG_NI2}}
% 	\caption{Convergence performance of different algorithms on EMNIST under the noise injection attack.}
% 	\label{fig_EMNIST_BRDRAG_NI}
% \end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_cifar10_01.pdf}
		\label{fig_CIFAR10_BRDRAG_NI1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_cifar10_05.pdf}
		\label{fig_CIFAR10_BRDRAG_NI2}}
	\caption{Convergence performance of different algorithms on CIFAR-10 under the noise injection attack.}
	\label{fig_CIFAR10_BRDRAG_NI}
\end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_cifar100_01.pdf}
		\label{fig_CIFAR100_BRDRAG_NI1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_cifar100_05.pdf}
		\label{fig_CIFAR100_BRDRAG_NI2}}
	\caption{Convergence performance of different algorithms on CIFAR-100 under the noise injection attack.}
	\label{fig_CIFAR100_BRDRAG_NI}
\end{figure}
% -------------------------------------------------------

% \begin{figure}[!t]
% 	\centering
% 	\subfloat[{\centering Accuracy on EMNIST with $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_emnist_01.pdf}
% 		\label{fig_EMNIST_BRDRAG_SF1}}
% 	\subfloat[{\centering Accuracy on EMNIST with $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_emnist_05.pdf}
% 		\label{fig_EMNIST_BRDRAG_SF2}}
% 	\caption{Convergence performance of different algorithms on EMNIST under the sign flipping attack.}
% 	\label{fig_EMNIST_BRDRAG_SF}
% \end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_cifar10_01.pdf}
		\label{fig_CIFAR10_BRDRAG_SF1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_cifar10_05.pdf}
		\label{fig_CIFAR10_BRDRAG_SF2}}
	\caption{Convergence performance of different algorithms on CIFAR-10 under the sign flipping attack.}
	\label{fig_CIFAR10_BRDRAG_SF}
\end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_cifar100_01.pdf}
		\label{fig_CIFAR100_BRDRAG_SF1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_cifar100_05.pdf}
		\label{fig_CIFAR100_BRDRAG_SF2}}
	\caption{Convergence performance of different algorithms on CIFAR-100 under the sign flipping attack.}
	\label{fig_CIFAR100_BRDRAG_SF}
\end{figure}

% -------------------------------------------------------

% \begin{figure}[!t]
% 	\centering
% 	\subfloat[{\centering Accuracy on EMNIST with $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_label_emnist_01.pdf}
% 		\label{fig_EMNIST_BRDRAG_LF1}}
% 	\subfloat[{\centering Accuracy on EMNIST with $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_label_emnist_05.pdf}
% 		\label{fig_EMNIST_BRDRAG_LF2}}
% 	\caption{Convergence performance of different algorithms on EMNIST under the label flipping attack.}
% 	\label{fig_EMNIST_BRDRAG_LF}
% \end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_label_cifar10_01.pdf}
		\label{fig_CIFAR10_BRDRAG_LF1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_label_cifar10_05.pdf}
		\label{fig_CIFAR10_BRDRAG_LF2}}
	\caption{Convergence performance of different algorithms on CIFAR-10 under the label flipping attack.}
	\label{fig_CIFAR10_BRDRAG_LF}
\end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_label_cifar100_01.pdf}
		\label{fig_CIFAR100_BRDRAG_LF1}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_label_cifar100_05.pdf}
		\label{fig_CIFAR100_BRDRAG_LF2}}
	\caption{Convergence performance of different algorithms on CIFAR-100 under the label flipping attack.}
	\label{fig_CIFAR100_BRDRAG_LF}
\end{figure}

% -------------------------------------------------------

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_cifar10_01_more_attacker.pdf}
		\label{fig_CIFAR10_BRDRAG_NI1_more_attacker}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_noise_cifar10_05_more_attacker.pdf}
		\label{fig_CIFAR10_BRDRAG_NI2_more_attacker}}
	\caption{Convergence performance of different algorithms on CIFAR-10 under the noise injection attack, where 60\% of the workers are attackers.}
	\label{fig_CIFAR10_BRDRAG_NI_more_attacker}
\end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_cifar10_01_more_attacker.pdf}
		\label{fig_CIFAR10_BRDRAG_SF1_more_attacker}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_sign_cifar10_05_more_attacker.pdf}
		\label{fig_CIFAR10_BRDRAG_SF2_more_attacker}}
	\caption{Convergence performance of different algorithms on CIFAR-10 under the sign flipping attack, where 60\% of the workers in the system are attackers.}
	\label{fig_CIFAR10_BRDRAG_SF_more_attacker}
\end{figure}

\begin{figure}[!t]
	\centering
	\subfloat[{\centering $\beta=0.1$.}]{\includegraphics[width=1.65in]{new_figs/br_label_cifar10_01_more_attacker.pdf}
		\label{fig_CIFAR10_BRDRAG_LF1_more_attacker}}
	\subfloat[{\centering $\beta=0.5$.}]{\includegraphics[width=1.65in]{new_figs/br_label_cifar10_05_more_attacker.pdf}
		\label{fig_CIFAR10_BRDRAG_LF2_more_attacker}}
	\caption{Convergence performance of different algorithms on CIFAR-10 under the label flipping attack, where 60\% of the workers in the system are attackers.}
	\label{fig_CIFAR10_BRDRAG_LF_more_attacker}
\end{figure}

% -------------------------------------------------------

% \begin{figure}[t!]
% \centering
% \includegraphics[width=1.5in]{ACAG_byzantine_iid.pdf}
% \caption{Performance comparison of DRAG with state-of-the-art algorithms using the CIFAR-10 dataset under byzantine attacks with low data heterogeneity.} \label{byzan-iid}
% \end{figure}

% Fig. 5 demonstrates the performances of DRAG, FLTrust \cite{cao2020fltrust} and FedAvg under byzantine attacks with low data heterogeneity. The algorithms stop training once the accuracy approaches 70\%. It is observed that the curve of FedAvg quickly falls down because it has no
% counter-attack techniques. And DRAG clearly achieves a better performance than FLTrust mainly because DRAG does not completely remove the attacker while taking advantage of its useful information through scaling and vector manipulation.


% Fig.~\ref{byzan-iid} presents the performance of DRAG, FLTrust \cite{cao2020fltrust}, and FedAvg under byzantine attacks with low data heterogeneity ($q=\frac{1}{M}$), and there is $A=1$ attacker among the $M=S=10$ clients. In the plot, FedAvg's performance declines rapidly due to its lack of counter-attack techniques. 
% In contrast, FLTrust demonstrates convergence by utilizing a small data sharing (root dataset) and ReLU-clipped cosine similarity. However, DRAG exhibits greater stability than FLTrust with just one data sharing step. The key advantage of DRAG lies in its ability to preserve the attacker's useful information through scaling and vector manipulation, instead of entirely removing the attacker. This property enables DRAG to achieve superior performance in handling Byzantine attacks while maintaining the integrity of the FL process.

 % By leveraging the attacker's information, DRAG adapts dynamically to the presence of malicious clients, effectively mitigating their negative impact on the training process.



% \begin{figure}[t!]
% \centering
% \includegraphics[width=1.5in]{ACAG_byzantine_noniid.pdf}
% \caption{Performance comparison of DRAG with state-of-the-art algorithms using the CIFAR-10 dataset under byzantine attacks with high data heterogeneity.} \label{byzan-noniid}
% \end{figure}

% In Fig.~\ref{byzan-noniid}, we increase the heterogeneity of data distribution among clients by setting $q=1$. Under this setting, the performance of FedAvg is still degraded as expected while FLTrust cannot converge at all. This is primarily due to the ReLU-clipping operation in FLTrust. As the heterogeneity of data increases, it is common that the local model update of a normal client diverges significantly from the trusted root direction. However, in FLTrust, such normal clients are erroneously identified as malicious and removed, causing the training process to crash. As analyzed above, DRAG retains its superiority even with the increased data heterogeneity.



% As previously analyzed, DRAG retains its superiority even with the increased data heterogeneity. Its advantageous properties, such as preserving useful information from attackers through scaling and vector manipulation instead of complete removal, allow DRAG to effectively handle byzantine attacks and maintain robustness in the face of highly heterogeneous data distributions.


% \begin{figure}[t!]
% \centering
% \includegraphics[width=1.5in]{ACAG_byzantine_emnist_noniid.pdf}
% \caption{Performance comparison of DRAG with state-of-the-art algorithms using the EMNIST dataset under byzantine attacks with low data heterogeneity.} \label{byzan-iidemnist}
% \end{figure}

% \begin{figure}[t!]
% \centering
% \includegraphics[width=1.5in]{ACAG_byzantine_emnist_iid.pdf}
% \caption{Performance comparison of DRAG with state-of-the-art algorithms using the EMNIST dataset under byzantine attacks with high data heterogeneity.} \label{byzan-noniidemnist}
% \end{figure}
% To further substantiate the effectiveness of DRAG in handling Byzantine attacks, we conduct performance testing on EMNIST, as shown in Fig.\ref{byzan-iidemnist} and Fig.\ref{byzan-noniidemnist}. In this setting, we increase the number of attackers to $A=4$ among a total of $M=10$ clients, as EMNIST is less challenging compared to CIFAR-10. We observe that FedAvg also achieves decent performance under low data heterogeneity since the malicious gradients are diluted through averaging, leading to a less negative impact on the training process compared to using CIFAR-10. Similarly, FLTrust performs well under low data heterogeneity, but the performance is limited when data is highly heterogeneously distributed, as it tends to mistake good clients for malicious ones. In contrast, DRAG consistently maintained its superiority across different data heterogeneity levels. 





\section{Conclusion and Future Work}
% In this work, we proposed a novel scheme justified by the name divergence-based adaptive aggregation (DRAG) to deal with the client drift issue in heterogeneous data distribution with local SGD. Unlike most methods that use control variates or regularization terms, the proposed DRAG algorithm heuristically drags each local gradient towards the reference direction according to its degree of divergence. The convergence analysis of DRAG is rigorously established to corroborate its feasibility theoretically. The performance of DRAG is also tested against state-of-the-art algorithms on EMNIST and CIFAR-10 datasets, the results of which substantiate the superiority of DRAG. Additionally, DRAG is proved to be robust to byzantine attacks that scale and reverse the direction of the local gradient.

% DRAG is strongly adaptive and has the potential in other scenarios such as its interaction with the wireless medium, its robustness against more complicated attacks, etc. These interesting directions will be explored in future works.

In this paper, we designed DRAG to address client drifts caused by data heterogeneity in FL systems. With the proposed reference direction and DoD, the local updates can be adaptively dragged towards the global update direction.
We further developed BR-DRAG, a Byzantine-resilient variant of DRAG, to resist Byzantine attacks.
Comprehensive convergence analyses of DRAG and BR-DRAG with non-convex loss functions were conducted under data heterogeneity, partial worker participation, and Byzantine attacks.
It was revealed that both DRAG and BR-DRAG achieve fast convergence compared to the existing algorithms.
Numerical results validated the superiority of DRAG in mitigating data heterogeneity and accelerating model convergence, and the robustness of BR-DRAG against malicious attacks.
In the future, we will extend DRAG and BR-DRAG to decentralized FL settings.






% to address the client-drift issue in heterogeneous data distribution with local SGD. Unlike methods that rely on control variates or regularization terms, DRAG employs a heuristic approach, dynamically dragging each local gradient towards the reference direction based on the degree of divergence. 
% We provide rigorous convergence analysis to theoretically support the feasibility of DRAG. Through extensive testing on EMNIST and CIFAR-10 datasets against state-of-the-art algorithms, DRAG demonstrates superior performance. Furthermore, we establish DRAG's resilience to Byzantine attacks that scale and reverse the direction of the local gradient.
% The adaptability of DRAG opens up exciting possibilities for exploring its utilization in other scenarios. To be specific, the wireless medium could be taken into consideration to explore more practical applications; we could also search for improvements of DRAG for attacks such as the label flipping attack, the Krum attack and so forth.



% To be specific, the wireless medium could be taken into consideration to explore more practical applications; we could also search for improvements of DRAG for attacks such as the label flipping attack, the Krum attack and so forth.






% \bibliography{bib}
% \bibliographystyle{iclr2022_conference}



% \section*{Appendix A \\ Proof of Theorem \ref{thm:alignment}}

% % By decomposing $\tilde{\Delta}^t$, we first have
% % \begin{align}
% %     \tilde{\Delta}^t = \nabla f\left(\boldsymbol{\theta}^t\right) + \mathbf{b}^t + \mathbf{e}^t,
% % \end{align}
% % where $\mathbf{b}^t = \frac{-\eta}{S} \sum_{m \in \mathcal{S}^t} \sum_{u=0}^{U-1} \nabla F_m(\theta_{m}^{t,u}) - \nabla f\left(\boldsymbol{\theta}^t\right) $ and $\mathbf{e}^t = \frac{-\eta}{S} \sum_{m \in \mathcal{S}^t} \sum_{u=0}^{U-1} (\nabla F_m(\theta_{m}^{t,u};z_{m}^{t,u})
% %                     -\nabla F_m(\theta_{m}^{t,u})) $.
% We first prove $\| \delta_m^{t,u} \| \!=\! \| \boldsymbol{\theta}_m^{t,u} \!-\! \boldsymbol{\theta}^{t} \| \leq 2\eta u \|\nabla F_m (\boldsymbol{\theta}^t)\|, \forall u =0,\cdots,U-1$ by induction. It holds for $u=0$ by definition. When $u\geq 1$, we have
% \begin{align} \label{delta bound}
%     &\left\|\delta_m^{t,u+1}\right\| \overset{(a)}{\leq} (1+\eta L)\left\|\delta_m^{t,u}\right\|+\eta \|\nabla F_m (\boldsymbol{\theta}^t)\| \nonumber \\
%     &\overset{(b)}{\leq}  \left(1+\frac{1}{4u} \right) 2 \eta u  \|\nabla F_m (\boldsymbol{\theta}^t)\|+\eta  \|\nabla F_m (\boldsymbol{\theta}^t)\| \nonumber \\
%     &\leq 2 \eta(u+1)  \|\nabla F_m (\boldsymbol{\theta}^t)\|,
% \end{align}
% where (a) comes from the triangle inequality and \textbf{Assumption 1}, and (b) is due to $\eta L U \leq \frac{1}{4}$. With \eqref{delta bound}, it is noted that
% \begin{align} \label{temp bound1}
%     &\left\| \mathbb{E}\left[\mathbf{g}_m^t\right]+\eta U \nabla F_m\left(\boldsymbol{\theta}^t\right) \right\| \nonumber \\
%     &= \left\|\eta \sum_{u=0}^{U-1}\left[\nabla F_m\left(\boldsymbol{\theta}^t\right)-\nabla F_m\left(\boldsymbol{\theta}_m^{t, u}\right)\right]\right\| \overset{(a)}{\leq} \eta L \sum_{u=0}^{U-1}\left\|\delta_m^{t,u}\right\| \nonumber \\
%     & \overset{(b)}{\leq} \eta L \sum_{u=0}^{U-1} 2 \eta u  \|\nabla F_m (\boldsymbol{\theta}^t)\| \leq \eta^2 L U^2  \|\nabla F_m (\boldsymbol{\theta}^t)\|,
% \end{align}
% where (a) is due to \textbf{Assumption 1} and the Cauchy-Schwarz inequality, and (b) comes from \eqref{delta bound}. Based on \eqref{temp bound1} and UAR worker sampling, we obtain
% \begin{align} \label{temp bound2}
%     \left\|\mathbb{E}\left[\tilde{\Delta}^t\right]+\eta U \nabla f(\boldsymbol{\theta}^t) \right\| \leq \eta^2 L U^2\left\|\nabla f(\boldsymbol{\theta}^t)\right\|.
% \end{align}

% Taking the inner products between $\tilde{\Delta}^t$ and $\nabla f(\boldsymbol{\theta}^t)$, we proceed to derive the following bound, as given by
% \begin{align}
% &\mathbb{E}\left[\left\langle\tilde{\Delta}^t, \nabla f(\boldsymbol{\theta}^t)\right\rangle  \right] \nonumber \\
% & =-\eta U\left\|\nabla f(\boldsymbol{\theta}^t)\right\|^2+\left\langle\mathbb{E}\left[\tilde{\Delta}^t\right]+\eta U \nabla f(\boldsymbol{\theta}^t), \nabla f(\boldsymbol{\theta}^t)\right\rangle \nonumber \\
% & \overset{(a)}{\leq} -\eta U(1 \!-\! \eta L U)\left\|\nabla f(\boldsymbol{\theta}^t)\right\|^2 \overset{(b)}{\leq} -\frac{3}{4}\eta U\left\|\nabla f(\boldsymbol{\theta}^t)\right\|^2 ,
% \end{align}
% where (a) is based on the Cauchy-Schwarz inequality and \eqref{temp bound2}, and (b) is due to the fact that $\eta L U \leq \frac{1}{4}$.

% Next, according to the definition of $\mathbf{v}_m^t$ in \eqref{vm}, we have
% \begin{align} \label{temp bound3}
%     &\left\langle \mathbf{v}_m^t, \mathbf{r}^t\right\rangle \nonumber \\
%     &= \! \left\langle \mathbf{g}_m^t, \mathbf{r}^t\right\rangle \!+\! \lambda_m^t\left(\left\|\mathbf{g}_m^t\right\| \! \left\|\mathbf{r}^t\right\| \!-\! \left\langle \mathbf{g}_m^t, \mathbf{r}^t\right\rangle\right) \!  \overset{(a)}{\geq} \! \left\langle \mathbf{g}_m^t, \mathbf{r}^t\right\rangle,
% \end{align}
% where (a) is due to the Cauchy-Schwarz inequality $\|\mathbf{x}\|\|\mathbf{y}\|\geq \langle \mathbf{x}, \mathbf{y}\rangle$. Averaging \eqref{temp bound3} over each worker $m \in \mathcal{S}^t$ yields
% \begin{align}\label{temp bound4}
%     \left\langle\Delta^t, \mathbf{r}^t\right\rangle \geq\left\langle\tilde{\Delta}^t, \mathbf{r}^t\right\rangle .
% \end{align}



% \newpage
% With \eqref{temp bound4}, the following bounds can be derived as 
% % the inner product between $\mathbf{r}^t$ and $\nabla f(\boldsymbol{\theta}^{t-1})$ is then lower-bounded by
% \begin{align}
%     &\mathbb{E}\left[\left\langle \mathbf{r}^t, \nabla f(\boldsymbol{\theta}^{t-1})\right\rangle \right] \nonumber \\
%     &\geq(1-\alpha) \mathbb{E}\left\langle\tilde{\Delta}^{t-1}, \nabla f(\boldsymbol{\theta}^{t-1})\right\rangle+\alpha \mathbb{E}\left\langle\tilde{\Delta}^{t-1}, \mathbf{r}^{t-1}\right\rangle , \\
% % \end{align}
% % \begin{align}
%   &\left\|\mathbb{E}\left[\langle \mathbf{r}^{t},\nabla f(\boldsymbol{\theta}^t)-\nabla f(\boldsymbol{\theta}^{t-1})\rangle\right]\right\| \! \overset{(a)}{\leq} \! L\,\mathbb{E}\bigl[\lVert \mathbf{r}^{t}\rVert \lVert\Delta^{t-1}\rVert\bigr] ,
%   % &\overset{(b)}{\leq} 4\,\eta^{2}LU^{2}\,
%   %      \mathbb{E}\lVert \nabla f(\boldsymbol{\theta}^{t-1})\rVert^{2} 
%   % \overset{(c)}{\leq} \alpha\,\eta U\,\mathbb{E}\lVert \nabla f(\boldsymbol{\theta}^{t-1})\rVert^{2},
% \end{align}
% where (a) comes from \textbf{Assumption 1}. Based on \eqref{temp bound1}, it is first noted that 
% \begin{align} \label{bound of g_m^t}
%     \| \mathbf{g}_m^{t-1} \| \leq 2\eta U \| \nabla F_m(\boldsymbol{\theta}^{t-1}) \| \overset{(a)}{\leq} 2\eta U \|\nabla f(\boldsymbol{\theta}^{t-1})\|,
% \end{align}
% where (a) comes from the variant of \textbf{Assumption 2}, i.e., $\mathbb{E}\left[\left\|\nabla F_m\left(\boldsymbol{\theta}^{t-1}\right)\right\|\right] \leq\left\|\nabla f(\boldsymbol{\theta}^{t-1})\right\|+\sigma_G$.
% For the upper bound of $\|\mathbf{r}^t\|$, we have
% \begin{align}
    
% \end{align}



% For the upper bound of $\|\Delta^{t-1}\|$, we have
% \begin{align}
%     & \|\Delta^{t-1}\| 
%     % \overset{(a)}{\leq} \frac{1}{S} \sum_{m \in \mathcal{S}^{t-1}}\left\|\mathbf{v}_m^{t-1}\right\| \nonumber \\
%     \overset{(a)}{\leq} \frac{1}{S} \sum_{m \in \mathcal{S}^{t-1}} \left\|\mathbf{g}_m^{t-1}\right\| 
%     % &\overset{(b)}{\leq} \frac{1}{S} \sum_{m \in \mathcal{S}^{t-1}} \!\! \max \left\{\left\|\mathbf{g}_m^{t-1}\right\|,\left\|\mathbf{r}^{t-1}\right\|\right\} 
%     \overset{(b)}{\leq} 2\eta U \| \nabla f(\boldsymbol{\theta}^{t-1})\| ,
% \end{align}
% % \begin{align}
% %     & \|\Delta^{t-1}\| 
% %     % \overset{(a)}{\leq} \frac{1}{S} \sum_{m \in \mathcal{S}^{t-1}}\left\|\mathbf{v}_m^{t-1}\right\| \nonumber \\
% %     \overset{(a)}{\leq} \frac{1}{S} \sum_{m \in \mathcal{S}^{t-1}} \left(1-\lambda_m^{t-1}\right)\left\|\mathbf{g}_m^{t-1}\right\|+\lambda_m^{t-1}\left\|\mathbf{r}^{t-1}\right\| \nonumber\\ 
% %     &\overset{(b)}{\leq} \frac{1}{S} \sum_{m \in \mathcal{S}^{t-1}} \!\! \max \left\{\left\|\mathbf{g}_m^{t-1}\right\|,\left\|\mathbf{r}^{t-1}\right\|\right\} \overset{(c)}{\leq} 2\eta U \| \nabla f(\boldsymbol{\theta}^{t-1})\| ,
% % \end{align}
% where (a) is due to the triangle inequality, and (b) comes from \eqref{bound of g_m^t}.

%  $&\overset{(b)}{\leq} 4\,\eta^{2}LU^{2}\,
%        \mathbb{E}\lVert \nabla f(\boldsymbol{\theta}^{t-1})\rVert^{2} 
%   \overset{(c)}{\leq} \alpha\,\eta U\,\mathbb{E}\lVert \nabla f(\boldsymbol{\theta}^{t-1})\rVert^{2},$



% % $L$-smoothness gives
% % $\lVert c^{t}-c^{t-1}\rVert\le L\lVert\theta^{t}-\theta^{t-1}\rVert
% %                               = L\lVert\Delta^{t-1}\rVert$.
% % From Lemma A.1 with $u=U$ we have
% % $\lVert g_{m}^{t-1}\rVert\le 2\eta U\,\lVert c^{t-1}\rVert$;
% % because $v_{m}^{t-1}$ is a convex combination of $g_{m}^{t-1}$ and
% % $r^{t-1}$, and $r^{t-1}$ itself is an average of earlier $v$’s,
% % induction yields
% % \(
% %   \lVert\Delta^{t-1}\rVert\le 2\eta U\,\lVert c^{t-1}\rVert
% % \)
% % and
% % \(
% %   \lVert r^{t}\rVert\le 2\eta U\,\lVert c^{t-1}\rVert.
% % \)
% % Consequently

% % where the last step uses \eqref{eq:eta-upper}.

% Adding \eqref{eq:rec-align} and \eqref{eq:drift-bound},
% and substituting \eqref{eq:key-neg} at round $t-1$,
% gives
% \begin{equation}\label{eq:r-better}
%   \mathbb{E}\langle r^{t},c^{t}\rangle
%   \ge \mathbb{E}\langle\tilde{\Delta}^{t},c^{t}\rangle
%        + (\alpha-4\eta LU)\,\eta U\,
%          \mathbb{E}\lVert c^{t}\rVert^{2}.
% \end{equation}

% % Because $\arccos(\cdot)$ is strictly decreasing on $(0,1]$,
% % \eqref{eq:app-goal} is equivalent to the cosine inequality
% % \begin{align}
% %   \mathbb{E}\! \Biggl[
% %      \frac{\langle \mathbf r^{t},\mathbf g^{t}\rangle}
% %           {\|\mathbf r^{t}\|\,\|\mathbf g^{t}\|}
% %   \Biggr]
% %   \;\;\ge\;\;
% %   \mathbb{E}\! \Biggl[
% %      \frac{\langle \tilde{\Delta}^{t},\mathbf g^{t}\rangle}
% %           {\|\tilde{\Delta}^{t}\|\,\|\mathbf g^{t}\|}
% %   \Biggr],
% %   \qquad
% %   \mathbf g^{t}:=\nabla f(\boldsymbol{\theta}^{t}). \label{eq:cos-target}
% % \end{align}

% % We prove \eqref{eq:cos-target} in four detailed steps.

% % % ----------------------------------------------------------------
% % \subsection*{A.1 \;  Bias–Noise Decomposition of $\tilde{\Delta}^{t}$}

% % Each client performs $U$ local SGD updates; the \emph{raw} model
% % difference is  
% % \[
% %   \mathbf g_{m}^{t}
% %   = \boldsymbol{\theta}_{m}^{t,U}-\boldsymbol{\theta}^{t}
% %   = -\eta\sum_{u=0}^{U-1}
% %       \nabla F\!\bigl(\boldsymbol{\theta}_{m}^{t,u};z_{m}^{t,u}\bigr).
% % \]
% % The server averages the $S$ selected clients (\emph{UAR without
% % replacement}) to get  
% % \[
% %   \tilde{\Delta}^{t}:=\frac{1}{S}\sum_{m\in\mathcal S^{t}}\mathbf g_{m}^{t}.
% % \]

% % \vspace{4pt}\noindent
% % \textbf{Unbiased local stochastic gradient.}
% % By Assumption 2,
% % $$
% %   \mathbb{E}\!\bigl[\nabla F_{m}(\boldsymbol{\theta};z)\bigr]
% %   =\nabla F_{m}(\boldsymbol{\theta}),\quad\forall m.
% % $$

% % \vspace{4pt}\noindent
% % \textbf{Step-wise bias.}
% % Taking conditional expectation w.r.t.\ the local data randomness $\xi$ and
% % client sampling $\mathcal S^{t}$,
% % \begin{align}
% %   \mathbb{E}\!\left[\tilde{\Delta}^{t}\mid\boldsymbol{\theta}^{t}\right]
% %   &= -\eta\sum_{u=0}^{U-1}
% %      \frac1M\sum_{m=1}^{M}
% %             \nabla F_m(\boldsymbol{\theta}_{m}^{t,u}) \notag\\
% %   &= -\eta U\nabla f(\boldsymbol{\theta}^{t})
% %      + \mathbf b_{t}
% %   =  \mathbf g^{t}+\mathbf b_{t},                            \label{eq:def-bias}
% % \end{align}
% % where
% % \begin{align}
% %   \mathbf b_{t}:= -\eta\sum_{u=0}^{U-1}
% %              \Bigl[\nabla f(\boldsymbol{\theta}_{m}^{t,u})
% %                   -\nabla f(\boldsymbol{\theta}^{t})\Bigr].
% % \end{align}
% % Using $L$–smoothness \eqref{lipschitz_gradient} and
% % \(
% %   \|\boldsymbol{\theta}_{m}^{t,u}-\boldsymbol{\theta}^{t}\|
% %    \le \eta u G
% % \)
% % with
% % $G:=\sqrt{\sigma_{L}^{2}+\sigma_{G}^{2}}$, we obtain the \emph{bias bound}
% % \begin{align}
% %   \|\mathbf b_{t}\|
% %   \le \eta\,\frac{U(U-1)}{2}\,LG
% %   =:\beta_{t}.                                                \label{eq:biasbdd}
% % \end{align}
% % Let $\displaystyle\beta_{\max}:=\max_{i<t}\beta_{i}$ for brevity.

% % \vspace{4pt}\noindent
% % \textbf{Zero-mean noise.}
% % Define
% % \[
% %   \boldsymbol{\varepsilon}_{t}
% %        :=\tilde{\Delta}^{t}-\bigl(\mathbf g^{t}+\mathbf b_{t}\bigr).
% % \]
% % Then
% % \[
% %   \mathbb{E}\!\left[\boldsymbol{\varepsilon}_{t}
% %                     \bigm|\boldsymbol{\theta}^{t}\right]=\mathbf 0,
% %   \quad
% %   \mathbb{E}\!\left[\|\boldsymbol{\varepsilon}_{t}\|^{2}
% %                     \bigm|\boldsymbol{\theta}^{t}\right]
% %        \le\frac{\sigma_{G}^{2}}{S}.                          \label{eq:var-noise}
% % \]

% % \vspace{4pt}\noindent
% % \textbf{Summary decomposition.}
% % \begin{align}
% %   \boxed{
% %   \tilde{\Delta}^{t}
% %   \;=\;\mathbf g^{t}\;+\;\mathbf b_{t}\;+\;\boldsymbol{\varepsilon}_{t}.
% %   }                                                         \label{eq:raw-decomp}
% % \end{align}

% % % ----------------------------------------------------------------
% % \subsection*{A.2 \;  Error Dynamics of the Reference Direction}

% % Define the deviation of $\mathbf r^{t}$ from $\mathbf g^{t}$:
% % $$
% %   \boldsymbol{\xi}_{t}:=\mathbf r^{t}-\mathbf g^{t},
% %   \qquad
% %   \boldsymbol{\xi}_{0}:=\mathbf r^{0}-\mathbf g^{0}.
% % $$
% % For $t\ge 1$, substitute
% % \(
% %   \Delta^{t-1}=\tilde{\Delta}^{t-1}
% % \)
% % into \eqref{reference_b}, then use
% % \eqref{eq:raw-decomp} to obtain
% % \begin{align}
% %   \boldsymbol{\xi}_{t}
% %   &=(1-\alpha)\boldsymbol{\xi}_{t-1}
% %      +\alpha\bigl(\mathbf b_{t-1}+\boldsymbol{\varepsilon}_{t-1}\bigr).
% % \end{align}
% % Taking $\|\cdot\|^{2}$ and total expectation yields
% % \begin{align}
% %   \mathbb{E}\!\bigl[\|\boldsymbol{\xi}_{t}\|^{2}\bigr]
% %   &\le (1-\alpha)^{2}\mathbb{E}\!\bigl[\|\boldsymbol{\xi}_{t-1}\|^{2}\bigr]
% %      +2\alpha^{2}\beta_{t-1}^{\,2}
% %      +2\alpha^{2}\tfrac{\sigma_{G}^{2}}{S}.                   \label{eq:xi-rec}
% % \end{align}
% % Iterate \eqref{eq:xi-rec} from $t=1$ and use
% % $\mathbb{E}\|\boldsymbol{\xi}_{0}\|^{2}\le\sigma_{G}^{2}/S$:
% % \begin{align}
% %   \boxed{
% %   \mathbb{E}\!\bigl[\|\boldsymbol{\xi}_{t}\|^{2}\bigr]
% %   \le(1-\alpha)^{2t}\frac{\sigma_{G}^{2}}{S}
% %      +\frac{2\alpha}{1-(1-\alpha)^{2}}
% %        \!\Bigl(\beta_{\max}^{2}+\tfrac{\sigma_{G}^{2}}{S}\Bigr).
% %   }                                                          \label{eq:xi-var}
% % \end{align}

% % % ----------------------------------------------------------------
% % \subsection*{A.3 \;  Expected Cosine Lower Bounds}

% % Let $\mathbf g\neq\mathbf0$, $\mathbf z$ s.t.
% % $\langle\mathbf g,\mathbf z\rangle=0$.
% % Then
% % \begin{align}
% %   \cos\!\angle(\mathbf g+\mathbf z,\mathbf g)
% %    \;=\;\frac{\|\mathbf g\|}
% %               {\sqrt{\|\mathbf g\|^{2}+\|\mathbf z\|^{2}}}
% %    =:h(X),\quad
% %    X:=\frac{\|\mathbf z\|^{2}}{\|\mathbf g\|^{2}},
% % \end{align}
% % where $h(x)=(1+x)^{-1/2}$ is \emph{convex} on $x\ge0$.
% % By Jensen,
% % \begin{align}
% %   \mathbb{E}\!\bigl[\cos\angle(\mathbf g+\mathbf z,\mathbf g)\bigr]
% %   \ge h\!\bigl(\mathbb E[X]\bigr).                            \label{eq:jensen}
% % \end{align}

% % \paragraph*{(i) Raw aggregate.}
% % Here $\mathbf z=\mathbf b_{t}+\boldsymbol{\varepsilon}_{t}$; using
% % \eqref{eq:biasbdd} and \eqref{eq:var-noise},
% % \begin{align}
% %   \mathbb{E}\!\Biggl[
% %     \frac{\langle\tilde{\Delta}^{t},\mathbf g^{t}\rangle}
% %          {\|\tilde{\Delta}^{t}\|\|\mathbf g^{t}\|}\Biggr]
% %   \ge
% %   h\!\Bigl(
% %      \frac{\beta_{t}^{\,2}+\sigma_{G}^{2}/S}
% %           {\mathbb{E}\|\mathbf g^{t}\|^{2}}
% %     \Bigr).                                                  \label{eq:cos-raw}
% % \end{align}

% % \paragraph*{(ii) Reference direction.}
% % Here $\mathbf z=\boldsymbol{\xi}_{t}$; by \eqref{eq:xi-var},
% % \begin{align}
% %   \mathbb{E}\!\Biggl[
% %     \frac{\langle\mathbf r^{t},\mathbf g^{t}\rangle}
% %          {\|\mathbf r^{t}\|\|\mathbf g^{t}\|}\Biggr]
% %   \ge
% %   h\!\Bigl(
% %      \frac{\alpha(\beta_{\max}^{2}+\sigma_{G}^{2}/S)}
% %           {\mathbb{E}\|\mathbf g^{t}\|^{2}}
% %     \Bigr).                                                 \label{eq:cos-ref}
% % \end{align}

% % Because $0<\alpha<1$ and $h(x)$ is \emph{strictly decreasing},
% % \[
% %   h\!\Bigl(
% %         \alpha(\beta_{\max}^{2}+\sigma_{G}^{2}/S)
% % /\mathbb{E}\|\mathbf g^{t}\|^{2}\Bigr)
% %   \;\ge\;
% %   h\!\Bigl(
% %         \beta_{t}^{\,2}+\sigma_{G}^{2}/S
% %         \;\big/\;\mathbb{E}\|\mathbf g^{t}\|^{2}\Bigr).
% % \]
% % Combining \eqref{eq:cos-raw}–\eqref{eq:cos-ref},
% % \begin{align}
% %   \mathbb{E}\!\Biggl[
% %     \frac{\langle\mathbf r^{t},\mathbf g^{t}\rangle}
% %          {\|\mathbf r^{t}\|\|\mathbf g^{t}\|}\Biggr]
% %  \ge
% %   \mathbb{E}\!\Biggl[
% %     \frac{\langle\tilde{\Delta}^{t},\mathbf g^{t}\rangle}
% %          {\|\tilde{\Delta}^{t}\|\|\mathbf g^{t}\|}\Biggr],   \label{eq:cos-final}
% % \end{align}
% % which is exactly inequality \eqref{eq:cos-target}.

% % % ----------------------------------------------------------------
% % \subsection*{A.4 \;  Concluding the Proof}

% % Since $\arccos(\cdot)$ is strictly decreasing on $(0,1]$,
% % applying $\arccos$ to \eqref{eq:cos-final} yields \eqref{eq:app-goal},
% % completing the proof. \hfill $\square$



%     % =================================================================
% %   Appendix A – Full-detail Proof of Theorem \ref{thm:alignment}
% % =================================================================







\section*{Appendix A \\ Proof of Theorem \ref{drag theorem}}
% \subsection{Lemma 1}
% \textit{The global update $\Delta^t$ at round $t$ can be written in the form as}
% \begin{align}
%     \Delta^t=(1-c)\mathbf{g}^t+\frac{c}{M}\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t+\frac{c}{M}\sum_{m\in\mathcal{M}}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t.\label{lemma1}
% \end{align}
% \begin{proof}
%     With the definition of $\Delta^t$, we can write:
%     \begin{align}
%         \Delta^t&=\frac{1}{M}\sum_{m=1}^M\left((1-\lambda_m^t)\mathbf{g}_m^t+\frac{\lambda_m^t\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t\right)\notag\\
%         \overset{()}{=}\frac{1}{M}\sum_{m=1}^M\left((1-\lambda_m^t)\mathbf{g}_m^t+\frac{\lambda_m^t\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t\right)\notag\\
%     \end{align}
% \end{proof}

Based on the $L$-smoothness of $f$, we first have:
\begin{align}
    &\mathbb{E}[f(\boldsymbol{\theta}^{t+1})]  \!\leq \!\!f(\boldsymbol{\theta}^{t}) \!+\! \left\langle\nabla\! f(\boldsymbol{\theta}^{t}), \mathbb{E}\!\left[\boldsymbol{\theta}^{t+1}\!\!\!-\!\boldsymbol{\theta}^t\right]\right\rangle\!\! +\!\! \frac{L}{2} \!\mathbb{E}\!\left[\|\boldsymbol{\theta}^{t+1}\!\!\!\!-\!\!\boldsymbol{\theta}^t\|^2\right]\notag\\
    &=f(\boldsymbol{\theta}^{t}) + \left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[\Delta^t+a_1-a_1\right]\right\rangle + \frac{L}{2} \mathbb{E}\left[\|\Delta^t\|^2\right]\notag\\
    & \overset{(a)}{=} f(\boldsymbol{\theta}^{t}) -(1-c)\eta U\mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right]+ \frac{L}{2}T_1+T_2,\label{one-step} 
    % &=f(\boldsymbol{\theta}^{t}) -(1-c)\eta U\|\nabla f(\boldsymbol{\theta}^t)\|^2+ \underbrace{\left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[\Delta^t+a_1\right]\right\rangle}_{T_2} + \frac{L}{2} \underbrace{\mathbb{E}\left[\|\Delta^t\|^2\right]}_{T_1}.\label{one-step}
\end{align}
where $\Delta^t = \boldsymbol{\theta}^{t+1}-\boldsymbol{\theta}^t$, and $a_1=(1-c)\eta U \nabla f(\boldsymbol{\theta}^t)$; $(a)$ is due to $T_1=\mathbb{E}\left[\|\Delta^t\|^2\right]$ and $T_2=\left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[\Delta^t+a_1\right]\right\rangle$. 

Next, we bound $T_1$ and $T_2$. For $T_1$, we have
\begin{align}  \label{T1}
    & T_1 = \mathbb{E}\Bigg[\Big\|\frac{1}{S}\sum_{m\in\mathcal{S}^t}\big((1-\lambda_m^t)\mathbf{g}_m^t+\frac{\lambda_m^t\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\mathbf{r}^t\big)\Big\|^2\Bigg]\notag\\
    &\overset{(a)}{\leq}\mathbb{E}\Big[\Big(\frac{1}{S}\sum_{m\in\mathcal{S}^t}\big((1-\lambda_m^t)\|\mathbf{g}_m^t\|+\frac{\lambda_m^t\|\mathbf{g}_m^t\|}{\|\mathbf{r}^t\|}\|\mathbf{r}^t\|\big)\Big)^2\Big]\notag\\
    % &\overset{(c)}{=}\frac{1}{S^2}\mathbb{E}\left[\left(\sum_{m\in\mathcal{S}^t}\|\mathbf{g}_m^t\|\right)^2\right] \\ \notag
    % & \overset{(c)}{=}\frac{1}{S^2 }\left(\sum_{m\in\mathcal{S}^t}\mathbb{E}\left[\left\|\sum_{u=0}^{U-1} \eta \nabla F(\boldsymbol{\theta}_m^{t,u};z_{m}^{t,u})\right\|\right]\right)^2\notag\\
    &\overset{(b)}{\leq}\frac{\eta^2}{S}\sum_{m\in\mathcal{S}^t}\mathbb{E}\Big[\big\|\sum_{u=0}^{U-1}\nabla F(\boldsymbol{\theta}_m^{t,u};z_{m}^{t,u})\big\|^2\Big]\notag\\
    &\overset{(c)}{=}\frac{\eta^2}{S}\sum_{m\in\mathcal{S}^t} \Big( \mathbb{E}\Big[\big\|\sum_{u=0}^{U-1} a_2\big\|^2\Big] + \mathbb{E}\Big[\big\|\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\big\|^2\Big] \Big) \notag\\
    &\overset{(d)}{=} {\eta^2U\sigma_L^2} +\frac{\eta^2}{M}\sum_{m\in\mathcal{M}}\mathbb{E}\Big[\big\|\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\big\|^2\Big] ,
    % &\overset{(a8)}{=}\frac{\eta^2U\sigma_L^2}{B}+\frac{\eta^2}{M}\sum_{m\in\mathcal{M}}\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right]\notag\\ 
\end{align}
where $a_2=\nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})-\nabla F_m(\boldsymbol{\theta}_m^{t,u})$; $(a)$ is due to the triangle inequality; $(b)$ is based on the definition of $\mathbf{g}_m^t$ and Cauchy-Schwartz inequality; $(c)$ follows from $\mathbb{E}[\|\boldsymbol{x}\|^2]=\mathbb{E}[\|\boldsymbol{x}-\mathbb{E}[\boldsymbol{x}]\|^2+\|\mathbb{E}[\boldsymbol{x}]\|^2]$ and \textbf{Assumption 2};
% $\mathbb{E}[\nabla F_m(\boldsymbol{\theta}_{m}^{j,u};z_{m,b}^{j,u})]=\nabla F_m(\boldsymbol{\theta}_{m}^{j,u})$ 
$(d)$ comes from $\mathbb{E}[\|\boldsymbol{x}_1\!\!+\!\ldots\!+\!\boldsymbol{x}_n\|^2]\!\!=\!\mathbb{E}[\|\boldsymbol{x}_1\|^2\!\!+\!\!\ldots\!+\!\!\|\boldsymbol{x}_n\|^2]$ if $\boldsymbol{x}_i,\,\forall i$, is zero-mean and independent, \textbf{Assumption 2}, and UAR worker sampling.

We can rewrite $T_2$ as
\begin{align} \label{bound of T2}
    & T_2\overset{(a)}{=}\Bigg\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\Bigg[(1-c)\mathbf{g}^t+\frac{c}{S}\sum_{m\in\mathcal{S}^t}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t\nonumber\\
    &+\frac{c}{S}\sum_{m\in\mathcal{S}^t}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t+(1-c)\eta U \nabla f(\boldsymbol{\theta}^t)\Bigg]\Bigg\rangle\notag\\
    &\overset{(b)}{=} \underbrace{\left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[(1-c)\mathbf{g}^t+a_1\right]\right\rangle}_{T_{2,1}} + T_{2,2} + T_{2,3},
    % \underbrace{\left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[(1-c)\mathbf{g}^t+a_1\right]\right\rangle}_{T_{2,1}}\nonumber\\
    % &+\underbrace{\left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[\frac{c}{M}\sum_{m\in\mathcal{S}^t}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t\right]\right\rangle}_{T_{2,2}}\notag\\
    % &+\underbrace{\left\langle\nabla f(\boldsymbol{\theta}^t), \mathbb{E}\left[\frac{c}{M}\sum_{m\in\mathcal{S}^t}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t\right]\right\rangle}_{T_{2,3}},
\end{align}
where $T_{2,2} = \langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}[\frac{c}{M}\sum_{m\in\mathcal{S}^t}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t]\rangle$, and $T_{2,3}= \langle\nabla f(\boldsymbol{\theta}^t), \mathbb{E} [\frac{c}{M}\sum_{m\in\mathcal{S}^t}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t ] \rangle$;
$(a)$ is based on the definition of $\Delta^t$, and $(b)$ comes from decomposition of $\mathbf{g}^t=\frac{1}{S}\sum_{m\in\mathcal{S}^t}\mathbf{g}_m^t$. 

We proceed to bound $T_{2,1}$, $T_{2,2}$ and $T_{2,3}$ in \eqref{bound of T2}.
The upper bound of $T_{2,1}$ is given by
\begin{align}
    &T_{2,1} 
    % = \left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[(1-c)\mathbf{g}^t+a_1\right]\right\rangle\notag\\
    \overset{(a)}{=} \left\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\left[(1-c)\bar{\mathbf{g}}^t+a_1\right]\right\rangle\notag\\
    &\overset{(b)}{=}\Big\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\big[-(1-c)\frac{1}{M}\sum_{m \in \mathcal{M}}\sum_{u=0}^{U-1}\eta\nabla F_m(\boldsymbol{\theta}_m^{t,u})\nonumber\\
    &+(1-c)\eta U\frac{1}{M}\sum_{m \in \mathcal{M}} \nabla F_m(\boldsymbol{\theta}^t)\big]\Big\rangle\notag\\
    &{=}\Big\langle\sqrt{\eta U(1-c)}\nabla f(\boldsymbol{\theta}^{t}),-\frac{\sqrt{\eta(1-c)}}{M\sqrt{U}}\mathbb{E}\big[\sum_{m \in \mathcal{M}}\sum_{u=0}^{U-1}a_3\big]\Big\rangle\notag\\
    &\overset{(c)}{\leq}a_4+\frac{(1-c)\eta}{2UM^2}\mathbb{E}\Bigg[\Big\|\sum_{m \in \mathcal{M}}\sum_{u=0}^{U-1}a_3\Big\|^2\Bigg]\notag\\
    % &\overset{(d)}{\leq}a_4+\frac{(1-c)\eta}{2M}\sum_{m=1}^M\sum_{u=0}^{U-1}\mathbb{E}\left[\left\|a_3\right\|^2\right]\notag\\
    &\overset{(d)}{\leq}a_4+\frac{(1-c)\eta L^2}{2M}\sum_{m \in \mathcal{M}}\sum_{u=0}^{U-1}\mathbb{E}\left[\left\|\boldsymbol{\theta}_m^{t,u}-\boldsymbol{\theta}^t\right\|^2\right]
    % -\frac{(1-c)\eta}{2UM^2}\mathbb{E}\left[\left\|\sum_{m=1}^M\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right]
    ,\label{T21}
\end{align}
where $\Bar{\mathbf{g}}^t \!\!=\!\! \frac{1}{M} \! \sum_{m\in\mathcal{M}}\mathbf{g}_m^t$, $a_3 \!\!=\!\! \nabla F_m(\boldsymbol{\theta}_m^{t,u}) \!-\! \nabla F_m(\boldsymbol{\theta}^t)$, 
and $a_4 \!=\! \frac{(1-c)\eta U}{2} \mathbb{E}\left[ \|\nabla f(\boldsymbol{\theta}^{t})\|^2 \right] $; $(a)$ is due to identical sampling distribution in each round; $(b)$ comes from the definition of $\Bar{\mathbf{g}}^t$; $(c)$ follows from $\langle\mathbf{x},\mathbf{y}\rangle \! \leq \! \frac{1}{2}[\|\mathbf{x}\|^2+\|\mathbf{y}\|^2]$; and $(d)$ is due to the Cauchy-Schwartz inequality and \textbf{Assumption 1}.

% Similarly, by following the bounding operation of $T_{2,1}$, we can have the following two inequalities
% \begin{align}
%     T_{2,2}\leq a_4+\frac{c\eta}{SU}\sum_{m\in\mathcal{S}^t}\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right],\label{T22}
% \end{align}

% \begin{align}
%     T_{2,3}\leq a_4+\frac{4c\eta\sigma_L^2}{B}+\frac{4c\eta}{MU}\sum_{m\in\mathcal{M}}\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right],\label{T23}
% \end{align}


The upper bound of $T_{2,2}$ is derived as  
\begin{align}
    & 
    T_{2,2}
    % =\Bigg\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\Big[\frac{c}{S}\sum_{m\in\mathcal{S}^t}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t\Big]\Bigg\rangle\notag\\
    % \overset{(a)}{=}\Bigg\langle\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\Big[\frac{c\eta}{S}\sum_{m\in\mathcal{S}^t}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\Big]\Bigg\rangle\notag\\
    \overset{(a)}{=}\Bigg\langle\sqrt{c\eta U}\nabla f(\boldsymbol{\theta}^{t}), \mathbb{E}\Bigg[\frac{\sqrt{c\eta}}{S\sqrt{U}}\sum_{m\in\mathcal{S}^t}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}a_5\Bigg]\Bigg\rangle\notag\\
    &\overset{(b)}{\leq}\frac{c\eta U}{2}\mathbb{E}\left[ \|\nabla f(\boldsymbol{\theta}^{t})\|^2 \right]+\frac{c\eta}{S^2U}\mathbb{E}\Bigg[\Bigg\|\sum_{m\in\mathcal{S}^t}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}a_5\Bigg\|^2\Bigg]\notag\\
    % &\overset{(d)}{\leq}\frac{c\eta U}{2}\mathbb{E}\left[ \|\nabla f(\boldsymbol{\theta}^{t})\|^2 \right]+\frac{c\eta}{S^2U}\left(\sum_{m\in\mathcal{S}^t}\mathbb{E}\left[\left\|a_5\right\|\right]\right)^2\notag\\
    &\overset{(c)}{\leq}\frac{c\eta U}{2}\mathbb{E}\left[ \|\nabla f(\boldsymbol{\theta}^{t})\|^2 \right]+\frac{c\eta}{SU}\sum_{m\in\mathcal{S}^t}\mathbb{E}\left[\left\|a_5\right\|^2\right],\label{T22}
\end{align}
where $a_5 \!\!\!=\!\!\! \sum_{u=0}^{U-1} \nabla F_m \! \left(\boldsymbol{\theta}_m^{t, u}\right)$; $(a)$ is due to the definition of $\mathbf{g}_m^t$; $(b)$ follows from $\langle\mathbf{x},\mathbf{y}\rangle \!\! \leq \!\! \frac{1}{2}[\|\mathbf{x}\|^2  \!+\! \|\mathbf{y}\|^2]$; and $(c)$ is based on the Cauchy-Schwartz inequality and triangle inequality.

The upper bound of $T_{2,3}$ is given by 
\begin{align}
    & 
    T_{2,3} 
    % \notag \\
    % =\left\langle\nabla f(\boldsymbol{\theta}^t), \mathbb{E}\left[\frac{c}{S}\sum_{m\in\mathcal{S}^t}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t\right]\right\rangle\notag\\
     {=} \Bigg\langle \!\!\! \sqrt{c\eta U}\nabla \!f(\boldsymbol{\theta}^t), \mathbb{E} \! \Bigg[\!\frac{\sqrt{c}}{S\sqrt{\eta U}} \!\! \! \sum_{m\in\mathcal{S}^t} \!\!\!
 \! \frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\| \!\!-\!\! \left\langle\mathbf{g}_m^t, \!\mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t\! \Bigg] \! \Bigg\rangle\notag\\
    &\overset{(a)}{\leq} \frac{c \eta U}{2} \mathbb{E}\left[ \|\nabla f(\boldsymbol{\theta}^{t})\|^2 \right] \nonumber \\
    &+\frac{c}{2\eta US^2}\mathbb{E}\Bigg[\Bigg\|\sum_{m\in\mathcal{S}^t}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\Bigg\|^2\|\mathbf{r}^t\|^2\Bigg]\notag\\
    &\overset{(b)}{\leq}\frac{c\eta U}{2}\mathbb{E}\left[ \|\nabla f(\boldsymbol{\theta}^{t})\|^2\right] +\frac{c}{2\eta US^2}\mathbb{E}\Bigg[\Big(2\sum_{m\in\mathcal{S}^t}\|\mathbf{g}_m^t\|\Big)^2\Bigg]\notag\\
    &\overset{(c)}{\leq}\frac{c\eta U}{2}\mathbb{E} \! \left[ \|\nabla f(\boldsymbol{\theta}^{t})\|^2 \right] \!+\! {2c\eta\sigma_L^2} \!+\! \frac{2c\eta}{MU} \!\!\! \sum_{m\in\mathcal{M}} 
    \!\! \mathbb{E} \! \left[\left\|a_5\right\|^2\right],\label{T23}
\end{align}
where $(a)$ comes from $\langle\mathbf{x},\mathbf{y}\rangle\leq\frac{1}{2}[\|\mathbf{x}\|^2+\|\mathbf{y}\|^2]$; $(b)$ follows from the triangle inequality, and $(c)$ is based on \textbf{Assumption~2} and uniform worker sampling without replacement.

We further bound $\mathbb{E}[\| a_5\|^2]$ in \eqref{T22} and \eqref{T23}; i.e.,
% To bound $T_{2,2}$ and $T_{2,3}$ in \eqref{T22} and \eqref{T23}, we proceed to derive the upper bound of $\mathbb{E}[\| \sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_{m}^{t,u})\|^2]$, as given by
% {\color{red}We continue to bound the following term in \eqref{T22} and \eqref{T23}:} 
\begin{align}
    &\mathbb{E} \! \left[\left\| a_5 \right\|^2\right] \!=\! \mathbb{E} \! \Bigg[\Bigg\| \! \sum_{u=0}^{U-1}(a_3 \!+\! \nabla F_m(\boldsymbol{\theta}^t) \!-\! \nabla f(\boldsymbol{\theta}^t) \!+\! \nabla f(\boldsymbol{\theta}^t))\Bigg\|^2\Bigg]\notag\\    
    % &=\mathbb{E}\Bigg[\Bigg\| \sum_{u=0}^{U-1}(\nabla F_m(\boldsymbol{\theta}_{m}^{t,u})-\nabla F_m(\boldsymbol{\theta}^t) +\nabla F_m(\boldsymbol{\theta}^t)\notag\\
    % &-\nabla f(\boldsymbol{\theta}^t)+\nabla f(\boldsymbol{\theta}^t))\Bigg\|^2\Bigg]\notag\\
    &\overset{(a)}{\leq} 3UL^2\sum_{u=0}^{U-1}\mathbb{E}[\|\boldsymbol{\theta}_{m}^{t,u}-\boldsymbol{\theta}^t\|^2]+3U^2\sigma_G^2+3U^2\|\nabla f(\boldsymbol{\theta}^t)\|^2\notag\\
    % &\overset{(f2)}{\leq} 15U^3L^2\eta^2(\sigma_L^2+6U\sigma_G^2)+(90U^4L^2\eta^2+3U^2)\|\nabla f(\boldsymbol{\theta}^t)\|^2+3U^2\sigma_G^2\notag\\
    &\overset{(b)}{=}C_1 \mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right]+C_2,\label{ti}
\end{align}
where $C_1=90U^4L^2\eta^2+3U^2$; $C_2=15U^3L^2\eta^2(\sigma_L^2+6U\sigma_G^2)+3U^2\sigma_G^2$;
$(a)$ is due to the Cauchy-Schwartz inequality; $(b)$ is based on \cite[Lemma 3]{reddi2020adaptive}, which yields: $\forall 1 \leq u \leq U $,
% \begin{align}
% \mathbb{E}\left[\left\|\boldsymbol{\theta}_{m}^{t,u} \!-\! \boldsymbol{\theta}^t\right\|^2\right]\!\leq\! 5U\eta^2(\sigma_L^2 \!+ \! 6U\sigma_G^2) \!+ \! 30U^2\eta^2 \mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right]. \notag 
% \end{align}
\begin{align}
\mathbb{E} \! \left[ \! \left\|\boldsymbol{\theta}_{m}^{t,u} \!-\! \boldsymbol{\theta}^t\right\|^2 \! \right] \!\! \leq \! 5U\eta^2(\sigma_L^2 \!+\!  6U\sigma_G^2) \!+\!  30U^2\eta^2 \mathbb{E} \! \left[ \! \|\nabla \! f(\boldsymbol{\theta}^t)\|^2\right]. \notag 
\end{align}




% Together with the upper bounds of $T_1$, $T_{2,1}$, $T_{2,2}$ and $T_{2,3}$ based on (\ref{T1}), (\ref{T21}), (\ref{T22}) and (\ref{T23}), \eqref{one-step} can be rewritten as
% % Substituting (\ref{T_2final}) and (\ref{T1}) into (\ref{one-step}), we have:
% % \begin{align}
% %     &\mathbb{E}[f(\boldsymbol{\theta}^{t+1})]\leq f(\boldsymbol{\theta}^{t})-(1-c)\eta U\|\nabla f(\boldsymbol{\theta}^{t})\|^2\notag\\
% %     &+ \frac{(1\!+\!c)\eta U}{2}\|\nabla f(\boldsymbol{\theta}^{t})\|^2\!+\!\frac{4c\eta\sigma_L^2}{B}+\frac{5c\eta}{U}(C_1\|\nabla f(\boldsymbol{\theta}^t)\|^2\!+\!C_2)\notag\\
% %     &+\frac{(1-c)\eta L^2U}{2}(5U\eta^2(\sigma_L^2+6U\sigma_G^2)+30U^2\eta^2\|\nabla f(\boldsymbol{\theta}^{t})\|^2) \notag\\
% %     &+\frac{\eta^2U\sigma_L^2L}{2B}+\frac{\eta^2L}{2}(C_1\|\nabla f(\boldsymbol{\theta}^t)\|^2+C_2)\notag\\
% %     &=f(\boldsymbol{\theta}^{t})-\eta U\Bigg(\frac{1-3c}{2}-\left(\frac{5c}{U}+\frac{\eta L}{2}\right)\left(90U^3L^2\eta^2+3U\right)\notag\\
% %     &-\frac{30(1-c)L^2U^2\eta^2}{2}\Bigg)\|\nabla f(\boldsymbol{\theta}^{t})\|^2+C_3,\label{onestepnew}
% % \end{align}
% \begin{align}
%     &\mathbb{E}[f(\boldsymbol{\theta}^{t+1})] \notag\\
%     &\leq f(\boldsymbol{\theta}^{t})-\eta U\Bigg(\frac{1-3c}{2}-\left(\frac{5c}{U}+\frac{\eta L}{2}\right)\left(90U^3L^2\eta^2+3U\right)\notag\\
%     &-\frac{30(1-c)L^2U^2\eta^2}{2}\Bigg)\mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right]+C_3,\label{onestepnew}
% \end{align}
% where $C_3={4c\eta\sigma_L^2}{}+\frac{\eta^2U\sigma_L^2L}{2}+(\frac{5c\eta}{U}+\frac{\eta^2L}{2})C_2+\frac{5(1-c)\eta^3 L^2U^2}{2}(\sigma_L^2+6U\sigma_G^2)$.

%  With $\eta \leq \frac{1}{8LU}$, we can prove that there exists a positive constant $\gamma$ satisfying
% \begin{align}
%     \gamma < \frac{1-3c}{2}-\left(\frac{5c}{U^2}+\frac{\eta L}{2U}\right)C_1-\frac{30(1-c)L^2U^2\eta^2}{2}. \label{gamma}
% \end{align}

Together with the upper bounds of $T_1$, $T_{2,1}$, $T_{2,2}$ and $T_{2,3}$ in (\ref{T1}), (\ref{T21}), (\ref{T22}) and (\ref{T23}), respectively, we rewrite \eqref{one-step} as
\begin{align}
    \mathbb{E}[f(&\boldsymbol{\theta}^{t+1})] 
    % \notag\\
    % &
    \!\leq \!f(\boldsymbol{\theta}^{t})\!-\!\eta U\Big(\frac{1\!\!-\!3c}{2}\!-\!\big(\frac{3c}{U}\!\!+\!\frac{\eta L}{2}\big)(90U^3L^2\eta^2\!\!+\!3U)\notag\\
    &-\frac{30(1-c)L^2U^2\eta^2}{2}\Big)\mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right]+C_3,\label{onestepnew}
\end{align}
where $C_3={2c\eta\sigma_L^2}{}+\frac{\eta^2U\sigma_L^2L}{2}+(\frac{3c\eta}{U}+\frac{\eta^2 L}{2})C_2+\frac{5(1-c)\eta^3 L^2U^2}{2}(\sigma_L^2+6U\sigma_G^2)$.
With $\eta \leq \frac{1}{8LU}$, we can prove that there exists a positive constant $\gamma$ satisfying
\begin{align}
    \gamma < \frac{1-3c}{2}-\left(\frac{3c}{U^2}+\frac{\eta L}{2U}\right)C_1-\frac{30(1-c)L^2U^2\eta^2}{2}. \label{gamma}
\end{align}
Rearranging \eqref{onestepnew} and summing from $t=0,\ldots,T-1$ yields
\begin{align} \label{drag final conv}
    \frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^t)\|^2\right] \leq \frac{f(\boldsymbol{\theta}^0)-f^*}{\gamma \eta U T} + V,
\end{align}
where $V=\frac{C_3}{\gamma U\eta}$. 
% $V=\frac{1}{\gamma}\left[\frac{4c\sigma_L^2}{BU}+\frac{\eta\sigma_L^2L}{2B}+\left(\frac{5c}{U^2}+\frac{\eta L}{2U}\right)C_2+\frac{5(1-c)L^2U\eta^2}{2}(\sigma_L^2+6U\sigma_G^2)\right]$. 
The proof is complete.

% Recalling the bounded result of $T_{2,1}$ in (\ref{T21}), we can further bound the term $\mathbb{E}\left[\left\|\boldsymbol{\theta}_m^{t,u}-\boldsymbol{\theta}^t\right\|^2\right]$:
% \begin{align}
%     &\mathbb{E}\left[\left\|\boldsymbol{\theta}_m^{t,u}-\boldsymbol{\theta}^t\right\|^2\right]=\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\frac{\eta}{B} \sum_{b=1}^B\nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})\right\|^2\right]\leq \eta^2 U^2G^2
% \end{align}
% \begin{align}
%     T_1&= \mathbb{E}\left[\|\Delta^t\|^2\right]\notag\\
%     &=\mathbb{E}\left[\left\|(1-c)\mathbf{g}^t+\frac{c}{M}\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t+\frac{c}{M}\sum_{m\in\mathcal{M}}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t\right\|^2\right]\notag\\
%     &\leq \underbrace{\mathbb{E}\left[3(1-c)^2\|\mathbf{g}^t\|^2\right]}_{T_{1,1}}+\underbrace{\mathbb{E}\left[3c^2\left\|\frac{1}{M}\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t\right\|^2\right]}_{T_{1,2}} + \underbrace{\mathbb{E}\left[3c^2\left\|\frac{1}{M}\sum_{m\in\mathcal{M}}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t\right\|^2\right]}_{T_{1,3}}, \label{T_1_general}
% \end{align}
% where the equality holds due to the definition of $\Delta^t$ and the inequality is due to Cauchy-Schwartz inequaltiy. The three terms $T_{1,1}, T_{1,2}, T_{1,3}$ are further bounded individually as follows.

% For the term $T_{1,1}$, we have:
% \begin{align}
%     T_{1,1}&=3(1-c)^2\mathbb{E}\left[\|\mathbf{g}^t\|^2\right]\notag\\
%     &\overset{(a1)}{=}3(1-c)^2\mathbb{E}\left[\left\|\frac{1}{M}\sum_{m\in\mathcal{M}}\mathbf{g}_m^t\right\|^2\right]\notag\\
%     &\overset{(a2)}{=}\frac{3(1-c)^2\eta^2}{M^2}\mathbb{E}\left[\left\|\sum_{m=1}^M\sum_{u=0}^{U-1}\frac{1}{B}\sum_{b=1}^B\nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})\right\|^2\right]\notag\\
%     &\overset{(a3)}{=}\frac{3(1-c)^2\eta^2}{M^2}\mathbb{E}\left[\left\|\sum_{m=1}^M\sum_{u=0}^{U-1}\frac{1}{B}\sum_{b=1}^B\left(\nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})-\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right)\right\|^2\right]\notag\\
%     &+\frac{3(1-c)^2\eta^2}{M^2}\mathbb{E}\left[\left\|\sum_{m=1}^M\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right]\notag\\
%     &\overset{(a4)}{=}\frac{3(1-c)^2\eta^2 U \sigma_L^2}{MB}+\frac{3(1-c)^2\eta^2}{M^2}\mathbb{E}\left[\left\|\sum_{m=1}^M\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right], \label{T11}
% \end{align}
% where $(a1)$ is due to the definition $\mathbf{g}^t=\frac{1}{M}\sum_{m=1}^M\mathbf{g}_m^t$; $(a2)$ is due to the definition of $\mathbf{g}_m^t$; $(a3)$ is because of $\mathbb{E}[\|x\|^2]=\mathbb{E}[\|x-\mathbb{E}[x]\|^2+\|\mathbb{E}[x]\|^2]$ and $\mathbb{E}[\nabla F_m(\boldsymbol{\theta}_{m}^{j,u};z_{m,b}^{j,u})]=\nabla F_m(\boldsymbol{\theta}_{m}^{j,u})$; and $(a4)$ is due to the fact that $\mathbb{E}[\|x_1+...+x_n\|^2]=\mathbb{E}[\|x_1\|^2+...+\|x_n\|^2]$ if $x_i'$s are independent with zero mean.

% Likewise, we can bound the term $T_{1,2}$:
% \begin{align}
%     T_{1,2}&=3c^2\mathbb{E}\left[\left\|\frac{1}{M}\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\mathbf{g}_m^t\right\|^2\right]\notag\\
%     &\overset{(b1)}{=}\frac{3c^2\eta^2}{M^2B^2}\mathbb{E}\left[\left\|\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\sum_{u=0}^{U-1}\sum_{b=1}^B\nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})\right\|^2\right]\notag\\
%     &\overset{(b2)}{=}\frac{3c^2\eta^2}{M^2B^2}\mathbb{E}\left[\left\|\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\sum_{u=0}^{U-1}\sum_{b=1}^B\left(\nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})-\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right)\right\|^2\right]\notag\\
%     &+\frac{3c^2\eta^2}{M^2}\mathbb{E}\left[\left\|\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right]\notag\\
%     &\overset{(b3)}{\leq}\frac{3c^2\eta^2 U\sigma_L^2}{MB}+\frac{3c^2\eta^2}{M^2}\mathbb{E}\left[\left\|\sum_{m\in\mathcal{M}}\frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right],\label{T12}
% \end{align}
% where $(b1)$ is due to the definition of $\mathbf{g}_m^t$; $(b2)$ and $(b3)$ come in the same way as $(a3)$ and $(a4)$.

% For the term $T_{1,3}$, we have:
% \begin{align}
%     T_{1,3}&=3c^2\mathbb{E}\left[\left\|\frac{1}{M}\sum_{m\in\mathcal{M}}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\mathbf{r}^t\right\|^2\right]\\
%     &\overset{(c1)}{=}3c^2\mathbb{E}\left[\left(\frac{1}{M}\sum_{m\in\mathcal{M}}\frac{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|-\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\|\mathbf{r}^t\|^2}\right)^2\left\|\mathbf{r}^t\right\|^2\right]\notag\\
%     &\overset{(c2)}{\leq}3c^2\mathbb{E}\left[\left(\frac{1}{M}\sum_{m\in\mathcal{M}}\frac{2\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}{\|\mathbf{r}^t\|^2}\right)^2\left\|\mathbf{r}^t\right\|^2\right]\notag\\
%     &\overset{(c3)}{=}\frac{12c^2}{M^2}\mathbb{E}\left[\left(\sum_{m=1}^M\|\mathbf{g}_m^t\|\right)^2\right]\notag\\
%     &\overset{(c4)}{=}\frac{12c^2\eta^2}{M^2B^2}\left(\sum_{m=1}^M\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\sum_{b=1}^B\nabla F(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})\right\|\right]\right)^2\notag\\
%     &\overset{(c5)}{\leq}\frac{12c^2\eta^2}{MB^2}\sum_{m=1}^M\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\sum_{b=1}^B\nabla F(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})\right\|^2\right]\notag\\
%     &\overset{(c6)}{=}\frac{12c^2\eta^2}{MB^2}\sum_{m=1}^M\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\sum_{b=1}^B\left(\nabla F_m(\boldsymbol{\theta}_m^{t,u};z_{m,b}^{t,u})-\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right)\right\|^2\right]\notag\\
%     &+\frac{12c^2\eta^2}{M}\sum_{m=1}^M\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right]\notag\\
%     &=\frac{12c^2\eta^2U\sigma_L^2}{B}+\frac{12c^2\eta^2}{M}\sum_{m=1}^M\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\nabla F_m(\boldsymbol{\theta}_m^{t,u})\right\|^2\right],\label{T13}
% \end{align}

\section*{Appendix B \\ Proof of Theorem \ref{brdrag theorem}}
% \begin{align}
%     & \|\mathbf{v}_m^t\| \overset{(a)}{\leq} \| c \mathbf{r}^t \| + \|  (1-c) \frac{\left\|\mathbf{r}^t\right\|}{\left\|\mathbf{g}_m^t\right\|} \mathbf{g}_m^t \| \nonumber \\
%     & + \| \frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\left\|\mathbf{g}_m^t\right\|^2} \mathbf{g}_m^t-c \frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\left\|\mathbf{g}_m^t\right\|\left\|\mathbf{r}^t\right\|} \mathbf{r}^t \| \nonumber \\
%     & \overset{(b)}{\leq} \left\|\mathbf{r}^t\right\| + c\| \frac{\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\left\|\mathbf{g}_m^t\right\|^2} \left\|\mathbf{r}^t\right\|(\left\|\mathbf{r}^t\right\|\mathbf{g}_m^t-\left\|\mathbf{g}_m^t\right\| \mathbf{r}^t) \| \nonumber \\
%     & \overset{(c)}{\leq} \left\|\mathbf{r}^t\right\|+\frac{2c}{\left\|\mathbf{g}_m^t\right\|} \| \left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle \| \overset{(d)}{\leq} (1+2c)\left\|\mathbf{r}^t\right\|. 
% \end{align}

As with Appendix A, with the $L$-smoothness of $f$, we have
\begin{align}\label{one-step new}
    \mathbb{E}[f(\boldsymbol{\theta}^{t+1})] 
    % \notag\\
    % &
    & \! \leq \! f(\boldsymbol{\theta}^{t})\! +\! \left\langle\nabla\! f(\boldsymbol{\theta}^{t}), \mathbb{E}[\boldsymbol{\theta}^{t+1}\!\!\!-\!\boldsymbol{\theta}^t]\right\rangle \!+ \!\frac{L}{2} \mathbb{E}[\| \Delta^t \|^2]\notag\\
    & = f(\boldsymbol{\theta}^{t}) + T_1 +T_2 + \frac{L}{2} T_3 ,
    %    &=f(\boldsymbol{\theta}^{t}) -(1-c)\eta U\|\nabla f(\boldsymbol{\theta}^t)\|^2+ T_1+T_2. 
\end{align}
where $T_1 = 
    \mathbb{E}\left[ \left\langle \nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S} \sum_{m \in \mathcal{A}^t} \mathbf{v}_m^t  \right\rangle \right] $, $T_2 = \mathbb{E}\left[ \left\langle \nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S} \sum_{m \in \mathcal{B}^t } \mathbf{v}_m^t  \right\rangle \right] $, and $T_3 =  \mathbb{E}\left[\|\Delta^t\|^2\right] $.
% For the upper bound of $T_1$, we first have
% \begin{align}
%     T_1 = 
%     \underbrace{\mathbb{E}\left[ \left\langle \nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S} \sum_{m \in \mathcal{A}^t} \mathbf{v}_m^t  \right\rangle \right] }_{T_{1,1}} + \underbrace{\mathbb{E}\left[ \left\langle \nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} \mathbf{v}_m^t  \right\rangle \right] }_{T_{1,2}}.
% \end{align}

With the definition of $\mathbf{v}_m^t$ in (\ref{vm1}) and \textbf{Assumption 4},  the modified gradient $\mathbf{v}_m^t$ can be rewritten as 
\begin{align} \label{expansion of v}
    \mathbf{v}_m^t = h_m^t {{\bf{r}}^t} + (1 - h_m^t)\left\| {{{\bf{r}}^t}} \right\|\frac{{{\bf{g}}_m^t}}{{\left\| {{\bf{g}}_m^t} \right\|}}, \ \forall m \in \mathcal{A}^t,
\end{align}
where $h_m^t = c^t(1 - x_m^t)$. 
% Based on \eqref{expansion of v}, the upper bound of
Then, 
$T_{1}$ can be upper bounded by 
% \begin{align}
%     & T_{1,1} \leq \mathbb{E}\left[ \frac{1}{S} \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\| \cdot \left\| \sum_{m \in \mathcal{A}^t} \mathbf{v}_m^t\right\| \right] \nonumber \\
%     & \overset{(a)}{\leq} w \mathbb{E}\left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\| \right] \cdot \mathbb{E}\left[ \left\| \mathbf{r}^t \right\| \right] ,
% \end{align}
\begin{align} \label{T11 new}
    & T_1 \!=\! \mathbb{E} \! \Bigg[ \! \Bigg\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S}\sum_{m \in \mathcal{A}^t} h_m^t {{\bf{r}}^t} + (1 - h_m^t)\left\| {{{\bf{r}}^t}} \right\|\frac{{{\bf{g}}_m^t}}{{\left\| {{\bf{g}}_m^t} \right\|}} \Bigg\rangle \! \Bigg ] \nonumber \\
    & \overset{(a)}{=} -c^t (w^t \!-\! x^t) \eta U\left\|\nabla \! f\left(\boldsymbol{\theta}^t\right)\right\|^2 \nonumber \\
    & + \mathbb{E}\left[\left\langle\nabla \! f\left(\boldsymbol{\theta}^t\right), c^t (w^t \!-\! x^t)\eta U \nabla f\left(\boldsymbol{\theta}^t\right) \!+\! c^t (w^t \!-\! x^t) \mathbf{r}^t\right\rangle\right] \nonumber \\
    &+ \mathbb{E}\Bigg[\Bigg\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S}\sum_{m \in \mathcal{A}^t}  (1 - h_m^t)\left\| {{{\bf{r}}^t}} \right\|\frac{{{\bf{g}}_m^t}}{{\left\| {{\bf{g}}_m^t} \right\|}}\Bigg\rangle\Bigg] , 
\end{align}
where $(a)$ is due to $A^t = w^t S$.

For the malicious workers $m \in \mathcal{A}^t$, define $x_m^t = \frac{\langle\mathbf{g}_m^t, \mathbf{r}^t\rangle}{\|\mathbf{g}_m^t\|\|\mathbf{r}^t\|}$ and $x^t = \frac{1}{S}\sum_{m \in \mathcal{A}^t} x_m^t$.
The second term on the RHS of \eqref{T11 new} is upper bounded, as given by
\begin{align} \label{second of T11 new}
    & \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), c^t (w^t \!-\! x^t)\eta U \nabla f\left(\boldsymbol{\theta}^t\right) + c^t (w^t \!-\! x^t) \mathbf{r}^t\right\rangle\right] \nonumber \\
    & \overset{(a)}{=} \left\langle \sqrt{\delta^t} \nabla f\left(\boldsymbol{\theta}^t\right),\frac{\sqrt{\delta^t}}{ {U} } \mathbb{E}\left[ \sum_{u=0}^{U-1} \nabla f\left(\boldsymbol{\theta}^t\right)-\nabla f\left(\boldsymbol{\theta}^{t, u}\right) \right]\right\rangle \nonumber \\
    & \overset{(b)}{\leq} \frac{\delta^t}{2} \mathbb{E}\left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \right] + \frac{\delta^t L^2}{2U} \sum_{u=0}^{U-1} \mathbb{E}\left[\left\|\boldsymbol{\theta}^{t, u}-\boldsymbol{\theta}^t\right\|^2\right].
\end{align}
Here, $ \delta^t\! = \! c^t (w^t\!-\! x^t) \eta U$. $(a)$ follows from~\eqref{trust}; $(b)$ is due to $\langle\mathbf{x}, \mathbf{y}\rangle \leq \frac{1}{2}\left[\|\mathbf{x}\|^2+\|\mathbf{y}\|^2\right]$ and \textbf{Assumption~1}. 

The last term on the RHS of \eqref{T11 new} is bounded as 
\begin{align} \label{third of T11 new}
    & \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S}\sum_{m \in \mathcal{A}^t}  (1 - h_m^t)\left\| {{{\bf{r}}^t}} \right\|\frac{{{\bf{g}}_m^t}}{{\left\| {{\bf{g}}_m^t} \right\|}}\right\rangle\right]  \nonumber \\
    & = \mathbb{E}\left[\left\langle \sqrt{\delta^t} \nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S\sqrt{\delta^t}} \sum_{m \in \mathcal{A}^t}(1 - h_m^t)\left\| {{{\bf{r}}^t}} \right\|\frac{{{\bf{g}}_m^t}}{{\left\| {{\bf{g}}_m^t} \right\|}}\right\rangle \right] \nonumber \\
    & \overset{(a)}{\leq} \! \frac{\delta^t}{2}\mathbb{E} \! \left[ \! \left\|\nabla \! f \! \left( \boldsymbol{\theta}^t\right)\right\|^2 \right] \!+\! \frac{\left\| \sum_{m \in \mathcal{A}^t}(1 \!-\! h_m^t)\right\|^2}{2S^2 \delta^t}
    \mathbb{E} \! \left[ \! \left\| \mathbf{r}^t \right\|^2 \right],
    % \frac{\frac{w^t}{S} \! \sum_{m \in \mathcal{A}^t} \! \left(1 \!-\! h_m^t\right)^2}{2 \delta^t}  \mathbb{E} \! \left[ \! \left\| \mathbf{r}^t \right\|^2 \right],
\end{align}
where $(a)$ comes from $\langle\mathbf{x}, \mathbf{y}\rangle \leq \frac{1}{2}\left[\|\mathbf{x}\|^2+\|\mathbf{y}\|^2\right]$.

Similarly, for the upper bound of $T_{2}$, we have
\begin{align} \label{T12 new}
    & T_{2} 
    % = \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S}\sum_{m \in \mathcal{B}^t} l_m^t {{\bf{r}}^t} + (1 - l_m^t)\rho_m^t {\bf{g}}_m^t\right\rangle\right] \nonumber \\
    = -c^t (1-w^t-y^t) \eta U \mathbb{E}\left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \right] \nonumber \\ 
    &+ \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), c^t (1 \!-\! w^t \!-\! y^t) \left( \mathbf{r}^t + \eta U \nabla f\left(\boldsymbol{\theta}^t\right) \right) \right\rangle\right] \nonumber \\
    &+ \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S}\sum_{m \in \mathcal{B}^t} \left(1- l_m^t\right) \rho_m^t \mathbf{g}_m^t\right\rangle\right] ,
    % &\overset{(a)}{=} -c(1-w) \eta U \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \nonumber \\
    % & + \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), a_2 + c(1-w) \mathbf{r}^t \right\rangle\right] \nonumber \\
    % & +\mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} \left(\mathbf{v}_m^t-c \mathbf{r}^t - c \frac{\mathbf{r}^t\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\left\|\mathbf{g}_m^t\right\|\left\|\mathbf{r}^t\right\|} \right) \right\rangle\right] \nonumber \\ 
    % & - \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} c \frac{\mathbf{r}^t \left\langle \mathbf{g}_m^t, \mathbf{r}^t \right\rangle}{\left\|\mathbf{g}_m^t\right\| \left\|\mathbf{r}^t\right\|}\right\rangle\right] 
\end{align}
where $l_m^t = c^t(1 - y_m^t)$ and $y^t=\frac{1}{S} \sum_{m \in \mathcal{B}^t } y_m^t$.
% $a_2 = c(1-w)(1-w-\beta^t) \eta U \nabla f\left(\boldsymbol{\theta}^t\right)$ and $\beta^t=\frac{1}{S} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} \beta_m^t$.

As with \eqref{second of T11 new}, with $\kappa^t = c^t (1-w^t-y^t)\eta U$, the second term on the RHS of \eqref{T12 new} can be bounded as 
\begin{align} \label{second of T12 new}
    &\mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), c^t (1-w^t-y^t) \left( \mathbf{r}^t + \eta U \nabla f\left(\boldsymbol{\theta}^t\right) \right) \right\rangle\right] \nonumber \\
    & \leq \frac{\kappa^t}{2}\mathbb{E}\!\left[\left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2\right] \!\!+\!\frac{\kappa^t L^2}{2U} \!\!\sum_{u=0}^{U-1} \mathbb{E}\!\left[\left\|\boldsymbol{\theta}^{t, u}\!\!-\!\boldsymbol{\theta}^t\right\|^2\right],
\end{align}
based on the Cauchy-Schwarz inequality and \textbf{Assumption 1}.  

The third term on the RHS of \eqref{T12 new} is upper bounded as
\begin{align} \label{third of T12 new}
    &  \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S}\sum_{m \in \mathcal{B}^t} \left(1- l_m^t\right) \rho_m^t \mathbf{g}_m^t\right\rangle\right]  \nonumber \\
    % & \overset{(d1)}{\leq} \mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{-1}{S}\sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} b_m^t \sum_{u=0}^{U-1} \eta \nabla F_m\left(\boldsymbol{\theta}_m^{t, u}; z_{m,b}^{t,u}\right)\right\rangle\right] \nonumber \\
    % & \overset{(d1)}{=} \frac{-1}{b^t} \mathbb{E}\left[\left\langle  b^t \sqrt{\eta U} \nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{\sqrt{\eta U}S}\sum_{m \in \mathcal{B}^t} b_m^t \sum_{u=0}^{U-1} \eta \nabla F_m\left(\boldsymbol{\theta}_m^{t, u}\right)\right\rangle\right] \nonumber \\
    & \overset{(a)}{=} \frac{-1}{b^t} \mathbb{E}\left[\left\langle  b^t \nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{ S}\sum_{m \in \mathcal{B}^t} b_m^t \sum_{u=0}^{U-1} \eta \nabla F_m\left(\boldsymbol{\theta}_m^{t, u}\right)\right\rangle\right] \nonumber \\
    & \overset{(b)}{\leq} \frac{-b^t \eta U}{2} \mathbb{E} \left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \right] \nonumber\\
    &  + \frac{1}{2 b^t} \mathbb{E}\Bigg[ \Bigg\| \frac{1}{S} \sum_{m \in \mathcal{B}^t} b_m^t \frac{\sqrt{\eta}}{\sqrt{U}} \sum_{u=0}^{U-1} \nabla F_m(\boldsymbol{\theta}^t) - \nabla F_m\left(\boldsymbol{\theta}_m^{t, u}\right) \Bigg\|^2  \Bigg]  \nonumber \\
    & \overset{(c)}{\leq} \!\! \frac{\eta L^2}{2b^t S} \!\!\! \sum_{m \in \mathcal{B}^t} \!\! ( b_m^t )^2 \!\! \sum_{u=0}^{U-1} \! \mathbb{E} \!\! \left[ \! \left\| \boldsymbol{\theta}^{t} \!\!-\! \boldsymbol{\theta}_m^{t, u} \right\|^2 \! \right] \!\!-\!\! \frac{b^t \eta U}{2}  \!\! \left\|  \nabla \! f \! \left(  \boldsymbol{\theta}^t  \right) \! \right\|^2 ,
    % & \leq \frac{-(1-w)(1-c)p\eta U}{2}\left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 + \frac{ (2-c)^2 q^2\eta L^2}{2(1-w)(1-c)pS} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} (b_m^t)^2 \sum_{u=0}^{U-1} \left\| \boldsymbol{\theta}^{t} - \boldsymbol{\theta}_m^{t, u} \right\|^2 ,
\end{align}
where $b_m^t = \left(1 - l_m^t\right) \rho_m^t $, and $b^t = \frac{1}{S}\sum_{m \in \mathcal{B}^t} b_m^t$;
$(a)$ is due to $\mathbf{g}_m^t = \boldsymbol{\theta}_m^{t,U}-\boldsymbol{\theta}^t$ and \textbf{Assumption 2}; $(b)$ comes from $-\!2\left\langle {\bf{x}},{\bf{y}} \right\rangle \!\leq\! -\| {\bf{x}} \|^2\! + \!\| {\bf{x}}\!-\!{\bf{y}} \|^2$; and $(c)$ is due to \textbf{Assumption~1} and the Cauchy-Schwarz inequality.

% Under Assumption 4,
% the last term of the RHS of \eqref{second of T11 new} can be bounded as 
% \begin{align}
%     & -\mathbb{E}\left[\left\langle\nabla f\left(\boldsymbol{\theta}^t\right), \frac{1}{S} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t} c \frac{\mathbf{r}^t\left\langle\mathbf{g}_m^t, \mathbf{r}^t\right\rangle}{\left\|\mathbf{g}_m^t\right\|\left\|\mathbf{r}^t\right\|}\right\rangle\right] \nonumber \\
%     & \leq  - \left\langle\sqrt{c \kappa (1-w) \eta U} \nabla f\left(\boldsymbol{\theta}^t\right), \frac{\sqrt{c\kappa (1-w)}}{\sqrt{\eta U}} \mathbb{E}\left[ \mathbf{r}^t \right] \right\rangle \nonumber \\
%     & \leq -\frac{c \kappa(1-w) \eta U}{2} \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \nonumber \\
%     & +\frac{c \kappa(1-w) \eta L^2}{2} \sum_{u=0}^{U-1} \mathbb{E}\left[\left\|\boldsymbol{\theta}^{t, u}-\boldsymbol{\theta}^t\right\|^2\right] ,
% \end{align}

By the triangle inequality, $\| {\bf{v}}_m^t \| \leq \| \left(1-\lambda_m^t\right) \frac{\left\|\mathbf{r}^t\right\|}{\left\|\mathbf{g}_m^t\right\|} \mathbf{g}_m^t \| + \| \lambda_m^t \mathbf{r}^t \| \leq \left\|\mathbf{r}^t\right\|, \forall i \in \mathcal{S}^t$. Then, $T_3$ is upper bounded by
\begin{align}
    T_3 
    % \!=\! \mathbb{E}\Bigg[ \Bigg\| \frac{1}{S} \! \Big(\sum_{m \in \mathcal{B}^t } \! \mathbf{v}_m^t \!+\! \sum_{m \in \mathcal{A}^t} \! {\mathbf{v}}_m^t \!\Big) \! \Bigg\|^2 \Bigg]  
    \overset{(a)}{\leq} \frac{1}{S} \sum_{m \in \mathcal{S}^t} \mathbb{E}\left[ \left\|\mathbf{r}^t\right\|^2 \right] = \mathbb{E}\left[\left\|\mathbf{r}^t\right\|^2\right], 
\end{align}
where $(a)$ is based on the Cauchy-Schwarz inequality.


% From Appendix A, we note that
% % \begin{align}
% %     &\mathbb{E}\left[\left\| \sum_{u=0}^{U-1}\nabla f(\boldsymbol{\theta}^{t,u})\right\|^2\right] \notag\\
% %     &=\mathbb{E}\Bigg[\Bigg\| \sum_{u=0}^{U-1}( \nabla f(\boldsymbol{\theta}^{t,u})-\nabla f(\boldsymbol{\theta}^t)+\nabla f(\boldsymbol{\theta}^t))\Bigg\|^2\Bigg]\notag\\    
% %     &\overset{(f1)}{\leq} 2UL^2\sum_{u=0}^{U-1}\mathbb{E}[\|\boldsymbol{\theta}^{t,u}-\boldsymbol{\theta}^t\|^2]+ 2U^2\|\nabla f(\boldsymbol{\theta}^t)\|^2\notag\\
% %     &\overset{(f2)}{=}(2U^2+60 U^4 \eta^2 L^2)\|\nabla f(\boldsymbol{\theta}^t)\|^2+ 10U^3 L^2 \eta^2 \sigma_L^2,
% % \end{align}
% % where $(f1)$ is due to Cauchy-Schwartz inequality; $(f2)$ is from \cite[Lemma 3]{reddi2020adaptive}, which proves the inequality 
% \begin{align}
% \mathbb{E}\left[\left\|\boldsymbol{\theta}^{t,u}-\boldsymbol{\theta}^t\right\|^2\right]\!\leq\! 5U\eta^2 \sigma_L^2   \!+ \! 30U^2\eta^2 \left\|\nabla f(\boldsymbol{\theta}^t)\right\|^2 \nonumber .
% \end{align}

% {\color{red}with bounded gradient, hard to prove!!}
% {\color{red}Assume $\left\| \nabla f\left(\boldsymbol{\theta}^{t, u}\right)\right\| \leq G, \forall t,u$, making $\left\|\sum_{u=0}^{U-1} \nabla f\left(\boldsymbol{\theta}^{t, u};z^{t,u}\right)\right\|$ bounded}
By the definition of $\mathbf{r}^t$ in \eqref{trust}, it follows that 
% A tighter version:
\begin{align} \label{bound of r}
    & \mathbb{E}\left[\left\|\mathbf{r}^t\right\|^2\right] 
    % = \mathbb{E}\left[\left\|\boldsymbol{\theta}^{t, U}-\boldsymbol{\theta}^t\right\|^2\right] \nonumber \\
    \overset{(a)}{\leq} \mathbb{E}\left[\left\|\boldsymbol{\theta}^{t, U-1} \!-\! \boldsymbol{\theta}^t \!-\! \eta \nabla f\left(\boldsymbol{\theta}^{t, U-1}\right) \right\|^2\right] \!+\! \eta^2\sigma_L^2 \nonumber \\
    & \overset{(b)}{\leq} (1+\frac{1}{2U-1})\mathbb{E}\left[\left\|\boldsymbol{\theta}^{t, U-1}-\boldsymbol{\theta}^t \right\|^2\right] + \eta^2\sigma_L^2 \nonumber \\
    & + 2\eta^2 U \mathbb{E}\left[\left\|\nabla f\left(\boldsymbol{\theta}^{t,U-1}\right)-\nabla f\left(\boldsymbol{\theta}^{t}\right) + \nabla f\left(\boldsymbol{\theta}^{t}\right)\right\|^2\right]  \nonumber \\
     & \overset{(c)}{\leq} (1 + \frac{1}{2U-1}+4\eta^2 U L^2)\mathbb{E}\left[\left\|\boldsymbol{\theta}^{t, U-1}-\boldsymbol{\theta}^t \right\|^2\right] \nonumber \\
     & + 4\eta^2 U \mathbb{E}\left[\left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2\right] + \eta^2 \sigma_L^2 \nonumber \\
     & \overset{(d)}{\leq}  \frac{U}{U \!-\! 1}\mathbb{E} \! \left[\left\|\boldsymbol{\theta}^{t, U\!-\!1} \!-\! \boldsymbol{\theta}^t \right\|^2\right] \!+\! 4\eta^2 U \mathbb{E}\!\left[\left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2\right] \!+\! \eta^2 \sigma_L^2 \nonumber \\
     & \overset{(e)}{\leq} 8 \eta^2 U^2 \mathbb{E}\left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \right] + 2\eta^2 U \sigma_L^2 ,
\end{align}
where $(a)$ is due to $\mathbb{E}\left[\|\boldsymbol{x}\|^2\right]=\mathbb{E}\left[\|\boldsymbol{x}-\mathbb{E}[\boldsymbol{x}]\|^2+\|\mathbb{E}[\boldsymbol{x}]\|^2\right]$ and \textbf{Assumption 2}; $(b)$ comes from Young's inequality; $(c)$ is based on the Cauchy-Schwarz inequality and \textbf{Assumption 1}; $(d)$ holds since $\frac{1}{2 U-1}+4\eta^2 U L^2 \leq \frac{1}{U-1} $ with $\eta \leq \frac{1}{3UL}$; and $(e)$ is due to $(1+ \frac{1}{U-1})^U \leq 2U$. 
It can also be verified that \eqref{bound of r} holds when $U=1$. For $1 \leq u \leq U$, we now have
\begin{align} \label{drift bound new}
    \mathbb{E} \! \left[ \left\| \boldsymbol{\theta}^{t, u} \!-\! \boldsymbol{\theta}^t \right\|^2 \right] \!\leq\! 8 \eta^2 U^2 \mathbb{E} 
 \! \left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \right] \!+\! 2 \eta^2 U \sigma_L^2.
\end{align} 

Together with the upper bounds of $T_1$, $T_2$, and $T_3$ based on \eqref{second of T11 new}, \eqref{third of T11 new}, \eqref{second of T12 new}, \eqref{third of T12 new} and \eqref{drift bound new}, we can rewrite \eqref{one-step new} as
% \begin{align}
%     &\mathbb{E}\left[f\left(\boldsymbol{\theta}^{t+1}\right)\right] \leq f\left(\boldsymbol{\theta}^t\right) \nonumber \\
%     &- \eta U \left( \frac{c(1-w)(1-w-\beta^t)}{2} + \frac{b^t}{2} - 8\eta U -\frac{4 \eta^2 U^2 L^2}{ b^t S} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t}\left(b_m^t\right)^2 - \frac{\frac{4}{S} \sum_{m \in \mathcal{A}^t}\left(1-c+c \alpha_m^t\right)^2}{ c\left(w-\alpha^t\right)} \right. \nonumber \\
%     & \left. + 4cw(w-\alpha^t)L^2 \eta^2 U^2 + 4c(1-w)(1-w-\beta^t)L^2 \eta^2 U^2 \right) \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \nonumber \\
%     & + 2\eta^2 U \sigma_L^2 \left( \frac{c w\left(w-\alpha^t\right) \eta U L^2}{2} + \frac{c(1-w)\left(1-w-\beta^t\right) \eta U L^2}{2} + \frac{\frac{1}{S} \sum_{m \in \mathcal{A}^t}\left(1-c+c \alpha_m^t\right)^2}{2 c\left(w-\alpha^t\right) \eta U} + \frac{\eta U L^2}{2 b^t S} \sum_{m \in \mathcal{S}^t \backslash \mathcal{A}^t}\left(b_m^t\right)^2 + 1 \right) ,
% \end{align}
\begin{align} \label{one round convergence new}
    &\mathbb{E} \! \left[f\left(\boldsymbol{\theta}^{t+1}\right)\right] \leq f\left(\boldsymbol{\theta}^t\right) \!-\! \eta U D_1^t \mathbb{E} \! \left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \right] \!+\! D_2^t ,
\end{align}
where $D_1^t = \frac{\kappa^t}{2\eta U} + \frac{b^t}{2} - 8\eta U - 4 (\delta^t+\kappa^t) \eta U L^2 -\frac{4 \eta^2 U^2 L^2}{ b^t S} \sum_{m \in \mathcal{B}^t }(b_m^t)^2 - 4\eta U \frac{ \| \sum_{m \in \mathcal{A}^t}(1- h_m^t) \|^2 }{ S^2 \delta^t}  $, and $D_2^t =2\eta^2 U \sigma_L^2 \Big( \frac{ (\delta^t + \kappa^t) L^2}{2} + \frac{  \left\| \sum_{m \in \mathcal{A}^t}\left(1- h_m^t\right) \right\|^2 }{2 S^2 \delta^t} + \frac{\eta U L^2}{2 b^t S} \sum_{m \in \mathcal{B}^t }\left(b_m^t\right)^2 + 1 \Big)$.
% \begin{align}
%    D_1^t & = \frac{\kappa^t}{2\eta U} + \frac{b^t}{2} - 8\eta U -\frac{4 \eta^2 U^2 L^2}{ b^t S} \sum_{m \in \mathcal{B}^t }\left(b_m^t\right)^2 \nonumber \\
%   & \frac{- 4\eta U \left\| \sum_{m \in \mathcal{A}^t}\left(1- h_m^t\right) \right\|^2 }{ S^2 \delta^t} - 4 (\delta^t+\kappa^t) \eta U L^2 , \nonumber   
% \end{align}
% and
% \begin{align}
%      D_2^t =&2\eta^2 U \sigma_L^2 \Big( \frac{ (\delta^t + \kappa^t) L^2}{2} + \frac{  \left\| \sum_{m \in \mathcal{A}^t}\left(1- h_m^t\right) \right\|^2 }{2 S^2 \delta^t}  \nonumber \\
%     &  + \frac{\eta U L^2}{2 b^t S} \sum_{m \in \mathcal{B}^t }\left(b_m^t\right)^2 + 1 \Big). \nonumber    
% \end{align}

Since the malicious workers can potentially upload any local updates, we have $x^t  \in [-w^t,w^t)$ and $\delta^t   \in (0, 2c^t w^t \eta U]$. Hence,
$y^t  \in (0, 1-w^t ]$,
$(b_m^t)^2  \in [ 0, q^2 ]$,
$b^t  \in [ (1-w^t)(1-c^t)p,(1-w^t)q ]$,
and
$\kappa^t  \in [0,c^t(1-w^t) \eta U)$ under \textbf{Assumptions 3} and \textbf{4}.
Note that $ \left\| \sum_{m \in \mathcal{A}^t}\left(1- h_m^t\right) \right\|^2=0, \forall t$, when $c^t = \frac{w^t}{w^t - x^t} \in [\frac{1}{2},1]$. With the stepsize $\eta \leq \frac{(1-w^t)(1-c^t)p}{32U+2U L^2(p+2q^2)}, \forall t$, and the above ranges of $\delta^t$, $\kappa^t$, $b^t$ and $b_m^t$, we can prove that there is a positive constant $\chi^t$ satisfying
\begin{align} \label{chi}
    \chi^t
    & <  \frac{(1-w^t)(1-c^t)p}{2} - 8\eta U - \frac{4\eta^2 U^2 L^2}{(1-c^t)p} q^2 \nonumber \\
    & 
    % - \frac{4(1-c+c\alpha)^2}{c} 
    - 4c^t (1+ w^t) L^2 \eta^2 U^2 < D_1^t , \ \forall t,
\end{align}
% where $\eta \leq \frac{8(1-w)(1-c)p}{96U+U L^2(16q^2+p)}$. 
% \begin{align} \label{chi}
%     \chi
%     & <  \frac{(1-w)(1-c)p}{2} - 8\eta U - \frac{4\eta^2 U^2 L^2}{(1-c)p}(2-c)^2 q^2 \nonumber \\
%     & 
%     % - \frac{4(1-c+c\alpha)^2}{c} 
%     - 4c w^2 L^2 \eta^2 U^2 - 4c (1-w)^2 L^2 \eta^2 U^2 < D_1^t , \ \forall t,
% \end{align}
% \begin{align}
%     \gamma 
%     & <  \frac{(1-w)(1-c)p}{2} - 8\eta U - \frac{4\eta^2 U^2 L^2}{(1-c)p}(2-c)^2 q^2 \nonumber \\
%     & - {\color{red}\frac{4(1-c+c\alpha)^2}{c}} - 4c w^2 L^2 \eta^2 U^2 - 4c (1-w)^2 L^2 \eta^2 U^2 < C_1^t
% \end{align}
% since $\exists c \in [0,1]$, for $ \forall \alpha \in [-1,0] $, $\frac{4(1-c+c\alpha)^2}{c} =0$.
Similarly, $D_2^t$ can be upper bounded as
% \begin{align}
%     D_2 
%     &:= 2\eta^2 U \sigma_L^2 \left( \frac{c (w^2+(1-w)^2)  \eta U L^2}{2}  +\frac{ 1 }{2 c  \eta U} \right. \nonumber \\
%     & \left. + \frac{ \eta U L^2}{2(1-c) p}(2-c)^2 q^2 + 1 \right) \geq D_2^t  , \ \forall t.
% \end{align}
\begin{align} \label{D_2^t}
    D_2^t \!\leq\! 2\eta^2 U \sigma_L^2 \! \left( \! \frac{(1 \!+\! w^t) c^t  \eta U L^2}{2} \!+\! \frac{ \eta U L^2 q^2}{2(1 \!-\! c^t) p}  \!+\! 1 \! \right) .
\end{align}

Define $\chi = \max_{t} \chi^t$, $D_2 = \max_{t} D_2^t$, $\tilde{c}=\max _t c^t$ and $\tilde{w}=\max _t w^t$. By reorganizing \eqref{one round convergence new} and summing it over $t=0,\cdots,T-1$, we finally obtain
% Let $\tilde{c} = \max_t c^t $ and $\tilde{w} = \max_t w^t $, and substitute $c^t$ and $w^t$ in \eqref{chi} and \eqref{D_2^t} with $\tilde{c}$ and $\tilde{w}$ to acquire $\chi$ and $D_2$,
% % $\chi = \min_{t} \chi^t$, $D_2 = \max_t D_2^t$, $\tilde{c} = \max_t c^t $, $\tilde{w} = \max_t w^t $ 
% rearranging \eqref{one round convergence new} and summing from $t=0, \ldots, T-1$, we finally obtain
\begin{align} \label{brdrag final conv}
    \frac{1}{T} \sum_{t=0}^{T-1} \mathbb{E} \left[ \left\|\nabla f\left(\boldsymbol{\theta}^t\right)\right\|^2 \right] \leq \frac{f\left(\boldsymbol{\theta}^0\right)-f^*}{\chi \eta U T} + W .
\end{align}
Here, $\eta \!\leq \!\frac{(1-\tilde{w})(1-\tilde{c})p}{32U\!+\!2U L^2(p\!+\!2q^2)}$; $W\!=\!\frac{D_2}{\chi  \eta U}$. 

% lr: $  \frac{\kappa^t}{2\eta U}  - 8\eta U - 4 (\delta^t+\kappa^t) \eta U L^2 -\frac{4 \eta^2 U^2 L^2}{ b^t S} \sum_{m \in \mathcal{B}^t }(b_m^t)^2  \geq 0  $
Since the PS may be unaware of the malicious workers, $c^t = \frac{w^t}{w^t-x^t}$ may not hold. Nevertheless, with the stepsize $\eta \leq \frac{2c^t(1-c^t)(1-w^t)(1-w^t-y^t)p}{32 p U+ c^t U L^2(p^2+q^2)} , \forall t$, if $c^t$ satisfies
\begin{align}
    \left( (\rho^t)^2 \!-\! 7w^t \!+\! 8x^t \!-\! 1 \right)(c^t)^2 \!+\! (15w^t \!+\! 1)c^t  \geq \frac{8(w^t)^2}{w^t-x^t} ,
\end{align}
we can prove that there also exists a positive constant $\chi^t \leq D_1^t$, and $D_2^t$ can be upper bounded as 
\begin{align} \label{D_2^t new}
    D_2^t \!\leq\! 2 \eta^2 U \sigma_L^2 \! \left( \! \frac{(1 \!+\! w^t) c^t  \eta U L^2}{2} \!\!+\!\! \frac{ \eta U L^2 q^2}{2(1 \!-\! c^t) p}  \!\!+\!\! H^t \!\!+\!\! 1 \! \right) ,
\end{align}
where $H^t = \frac{  \left\| \sum_{m \in \mathcal{A}^t}\left(1- h_m^t\right) \right\|^2 }{2 S^2 \delta^t}$ is bounded. Then, \eqref{brdrag final conv} can also be attained. 
This proof concludes.


\bibliographystyle{IEEEtran}
\bibliography{bib}

\end{document}