\section{Introduction}
Federated Learning (FL) has gained popularity as a framework for use-cases involving training machine learning models on edge devices without transmitting the collected data to a central server \citep{mcmahan2017communication}\blfootnote{$^*$Denotes equal contribution.\\This research was supported by Army Cooperative Agreement W911NF2120076.
\\ A. S. Bedi, F. Huang, and D. Manocha are with the University of Maryland, College Park, MD, USA. Email: \{amritbd,furongh,dmanocha\}@umd.edu. C. Fan is with the 	University of Massachusetts, Amherst, MA, USA. Email: \{chenfan\}@umass.edu
. A. K. Sahu and A. Koppel are with Amazon. Email: \{anit.sahu@gmail.com, {aekoppel}@amazon.com\}. B. M. Sadler is with DEVCOM Army Research Laboratory, Adelphi, MD, USA. Email: \{brian.m.sadler6.civ\}@army.mil.\\}. A typical cross-device setting in a federated learning setup involves a central server which seeks to train a global model based on the data collected and resident at edge devices, with the constraint that raw data is not to be transmitted back to the server. Mathematically, we can write it as:
\begin{align}\label{eq_main_problem}
	\min_{{\bbx\in\mathcal{X}}} F(\bbx): = \frac{1}{N}\sum_{i=1}^{N} f_i(\bbx),
\end{align}
where $\mathcal{X}\subset \mathbb{R}^d$ is a compact convex subset and $F(\bbx)$ is the sum of $N$ possibly non-convex local objectives $f_i(\bbx)$ \citep{li2020federated}. All the devices\footnote{We use nodes and devices interchangeably throughout the paper.} in the network are connected in a star topology, as is standard in FL literature \citep{mcmahan2017communication,karimireddy2020scaffold}). Local training data sets can have different sample sizes and are not necessarily drawn from a common distribution, meaning that $\mathbb{P}(\zeta_i)$ is allowed to be distinct (heterogeneous) for each device $i$.
The goal of standard FL is to train a global model $\bbx^*$ by solving \eqref{eq_main_problem}, which performs well across all the clients \citep{mcmahan2017communication,li2018federated}.
In the presence of data heterogeneity across devices, it is highly unlikely that one global model would work well for all devices. This has been highlighted in \citep{li2019fair}, where a large spread in terms of performance was noted across devices.
When devices' local objectives are distinct, solving \eqref{eq_main_problem} can potentially lead to local models which are far away from models obtained by solving:
\begin{align}\label{eq_local_problem}
	\min_{{\bbx\in\mathcal{X}}} f_i(\bbx),
\end{align}
For instance, consider the problem of learning ``language models" for a cellphone keyboard, where the goal is to predict the next word. FL can be used in such a case to learn a common global model, but a global model might fail to capture distinctive writing styles, as well as cultural nuances of different users. In such a case, specific personalized local model [cf. \eqref{eq_local_problem}] for each device is required which utilizes model updates from other devices. Therefore, there are two competing criteria: \emph{global performance} in terms of  \eqref{eq_main_problem} as compared with  \emph{local perspective} [cf. \eqref{eq_local_problem}]. {Hence we ask this question:
\emph{``\textbf{Can we automate the balance between global and local model performance in federated learning?}"}}

We answer this question affirmatively in this work. We note that balancing between the global objective in \eqref{eq_main_problem} and personalization \eqref{eq_local_problem} in general involves multi-objective optimization \citep{tan2022towards}. This class of problems cannot be solved exactly, especially in the case of non-convex local objectives with heterogeneous data, and may be approximately solved with the scalarized composite criterion $\alpha f_i(\bbx) + (1 - \alpha) (1/N)\sum_{i} f_i(\bbx)$ \citep{jahn1985scalarization}.
Current approaches in literature tend to focus only on the performance of global model  \citep{mcmahan2017communication,li2018federated,karimireddy2020scaffold}, or the local model \citep{fallah2020personalized,hanzely2020lower}, but do not quantitatively calibrate the trade-off between them. Prioritizing global model performance only amongst the individual devices admits a reformulation as a consensus optimization problem \citep{nedic2009distributed,nedic2010constrained}, which gives rise to FedAvg \citep{mcmahan2017communication}. In this context, it is well-known that averaging steps approximately enforce consensus \citep{shi2015extra}, whereas one can enforce the constraint exactly by employing Lagrangian relaxations, namely, ADMM \citep{boyd2011distributed}, saddle point methods \citep{nedic2009subgradient}, and dual decomposition \citep{terelius2011decentralized}. This fact has given rise to efforts to improve the constraint violation of federated learning algorithms, as in FedPD \citep{zhang2021fedpd} and FedADMM \citep{wang2022fedadmm}.

Other approaches involve using model-agnostic meta-learning \citep{fallah2020personalized}, in which one executes one gradient step as an approximation for \eqref{eq_local_problem} as input for solving  \eqref{eq_main_problem} with objective $\frac{1}{N}\sum_{i=1}^{N} f_i(\bbx-\alpha \nabla f_i(\bbx))$. However, it does not explicitly allow one to trade off local and global performance. Several works have sought to balance between these competing local and global criteria based upon regularization \citep{li2018federated,hanzely2020lower,t2020personalized,li2021ditto}. Alternatives prioritize the performance of the global model amidst heterogeneity via control variate corrections \citep{karimireddy2020scaffold,acar2021federated}. Please refer to Appendix \ref{related_work} for additional related work context.

In this work, we reinterpret augmentation (which may be interpreted as an extension of the proximal augmentation approaches in \citep{hanzely2020lower,li2018federated,t2020personalized}) as (i) a relaxation of the consensus equality constraint, and (ii) as a constrained optimization alternative to scalarizing the multi-criterion problem of finding a minimizer simultaneously of \eqref{eq_main_problem} and \eqref{eq_local_problem}.
We consider a problem in which the global objective \eqref{eq_main_problem} is the primal, which owing to node-separability, allows each node to only prioritize its local objective  \eqref{eq_local_problem}. Then, the deviation from the global model is imposed as a constraint with a personalization hyper-parameter $\gamma_i$ for each agent $i$ to be enforced via Lagrange multipliers. The particular constraint formulation is derived via a relaxation of a consensus reformulation of \eqref{eq_main_problem}. Therefore, our \emph{main contributions} are as follows:\\
(1) We rigorously define the tradeoff between the performance of local and global models across devices in federated learning as a multi-criterion problem to be approximated through a constrained non-convex program (cf. \ref{eq_consensus_reformulation}).
This formulation itself is novel for the FL settings.\\
(2) We derive the Lagrangian relaxation of this problem and an instantiation of primal-dual method, which, owing to node-separability of the Lagrangian, admits a federated algorithm called Federated Beyond Consensus (\texttt{FedBC}) (cf. Algorithm  \ref{alg:alg_FedBC}). \\
(3) We not only establish theoretically the convergence of the global and the local models and that its rates are at par with the state of the art but also illustrate the efficacy of \texttt{FedBC} through a suite of datasets (Synthetic, MNIST, CIFAR-10, Shakespeare).

\section{Problem Formulation}\label{sec:prob}
In this section, we formulate the constrained non-convex program underlying the algorithm by drawing connections to consensus optimization \citep{nedic2009distributed}.
Consider a consensus reformulation of \eqref{eq_main_problem}, where each node $i$ is now only responsible for its local copy $\bbx_i$ of the global variable $\bbz$:
\begin{align}
	\min_{\{{(\bbz, \bbx_i)\in\mathcal{X}}\}} \sum_{i=1}^{N} f_i(\bbx_i)
	~~\text{s.t. } \bbx_i=\bbz, \ \ \forall i.\label{assumption}
\end{align}
\begin{figure}[t]

     \begin{subfigure}[b]{0.6\textwidth}
     \centering
    \includegraphics[scale=0.36]{idea2.png}
         \caption{ }
         \label{fig:loss_syn_e5}

      \end{subfigure}

     %
          \begin{subfigure}[b]{0.30\textwidth}
    \includegraphics[scale=0.14]{gamma1.pdf}
    \caption{}
    \label{fig:gamma_example}
     \end{subfigure}
     \caption{\textbf{(a)} In the left figure, note that the consensus in standard FL results in averaging at server, which doesn't allow to converge to optimal. For the right figure in (a), we introduced  proximity constraints with parameter $\gamma_i$ which allows server model to converge to optimal. \textbf{(b)} $\gamma$ establishes a balance between local and global performance. \vspace{-5mm} }
     \label{figure_1}
     \end{figure}
The linear equality constraints $\bbx_i=\bbz$ for all $i$ [cf. \eqref{assumption}] enforces consensus among all the devices in the network. To solve this in a federated manner, one may employ techniques from multi-agent optimization \citep{nedic2009distributed,nedic2010constrained} and consider localized gradient updates followed by averaging steps, as in FedAvg \citep{mcmahan2017communication}.
Setting aside the issue of how sharply one enforces the constraints for the moment, observe that in \eqref{assumption}, each node must balance between the two competing global and local objectives. These quantities only coincide when the set of minimizers of the sum is contained inside the set of minimizers of each cost function in the sum. This holds only under the case that the sampling distributions $\mathbb{P}(\zeta_i)$ coincide. Efforts to deal with the gap between the global objective \eqref{eq_main_problem} and personalization  \eqref{eq_local_problem} have relied upon various augmentations of the local objective, e.g.,
\begin{align}\label{eq:fedprox}
    &f_i(\bbx) + (\mu/2)\|\bbx-\bbz\|^2  \ \ \ \ \ \ &\text{in FedProx\citep{li2018federated,hanzely2020lower}}
    \\
    \label{eq:pfedme}
    &\arg\min_{\bbtheta}f_i(\bbtheta) + (\mu/2)\|\bbtheta-\bbz\|^2  \ \ \ \ \ \ &\text{in pFedMe\citep{t2020personalized}}
    \\
    \label{eq:perfedavg}

    &f_i(\bbx-\nabla f_i(\bbx))  \ \ \ \ \ \ &\text{in Per-FedAvg\citep{fallah2020personalized}}
.\end{align}

In the above objectives, observe that a penalty coefficient is introduced, which then requires its own hyper-parameter tuning in order to obtain a suitable tradeoff between global performance and personalization. This relationship is even more opaque in meta-learning, as the tradeoff then depends upon mixed first-order partial derivatives of the local objective with respect to the global model -- see \citep{fallah2020personalized}. Therefore, it makes sense to discern whether it is possible to obtain a methodology to solve for the suitable trade off between local and global performance while solving for the model parameters themselves. To do so, we reinterpret the penalization in \eqref{eq:fedprox} as a constraint, which gives rise to the following problem:
\begin{align}\label{eq_consensus_reformulation}
	\min_{\{{(\bbz, \bbx_i)\in\mathcal{X}}\}} \sum_{i=1}^{N} f_i(\bbx_i)~~\text{s.t. }\|\bbx_i-\bbz\|^2\leq \gamma_i, \ \ \forall i\; .
\end{align}
In settings where there is a degree of correlation among devices' estimates, this formulation is known to exhibit numerical advantages \citep{koppel2017proximity,cao2020decentralized}. We note that $\gamma_i$ calibrates the trade off between the performance of the local and global model, as may be observed in Fig. \ref{figure_1}. This result for a synthetic dataset is explained in full detail with corroborating instances on real data, in Sec. \ref{sec:experiments}. Therefore, algorithms to solve \eqref{eq_consensus_reformulation} explicitly allow one to determine the degree of personalization. Next, we derive an algorithmic tool to solve  \eqref{eq_consensus_reformulation}.

\section{\texttt{FedBC}: Federated Learning Beyond Consensus}\label{sec:algorithm}
To solve \eqref{eq_consensus_reformulation}, one could consider the primal-dual method \citep{nedic2009subgradient} or ADMM \citep{boyd2011distributed}. However, as the constraints [cf \eqref{eq_consensus_reformulation} are nonlinear, ADMM requires a nonlinear optimization in the inner-loop. Thus, we consider the primal-dual method, which may be derived by Lagrangian relaxation of \eqref{eq_consensus_reformulation} as:
\begin{align}\label{min_max}
	\mathcal{L}(\bbz, \{\bbx_i, \lambda_i\}_{i=1}^N) =   \frac{1}{N}\sum_{i=1}^{N}  \mathcal{L}_i(\bbz, \bbx_i, \lambda_i),
\end{align}
where $\mathcal{L}_i(\bbz, \bbx_i, \lambda_i):= \big[f_i(\bbx_i) +\lambda_i\left(\|\bbx_i-\bbz\|^2-\gamma_i\right)\big]$. Then we alternate between primal minimization and dual maximization. To do so, ideally one would minimize the Lagrangian  \eqref{min_max} with respect to $\bbx_i$ while keeping $\bbz$ and $\{\lambda_i\}_{i=1}^N$ constant, i.e., at give instant $t$, we solve for $\bbx_i$ as
	\begin{align}\label{primal_vari_1}
		\bbx_i^{t+1}=& \operatornamewithlimits{argmin}_{\bbx } \mathcal{L}_i(\bbx,\bbz^t,  \lambda_i^t).
\end{align}
As local objectives may be non-convex, solving \eqref{primal_vari_1} is not simpler than solving \eqref{eq_consensus_reformulation} for given $\bbz^t$ and $\lambda_i^t$. To deal with this, we consider an oracle which provides an $\epsilon_i$-approximated solution of the form
\begin{align}\label{Proposed_primal}
		\bbx_i^{t+1}=&\text{Oracle}_i (\mathcal{L}_i(\bbx_i^t,\bbz^t,  \lambda_i^t),K_i);  &[\text{{$K_i$-local updates}}].
\end{align}
{where the $\epsilon_i$-approximate solution $\bbx_i^{t+1}$ is a stationary point of the Lagrangian in the sense of $\|\nabla_{\bbx_i}\mathcal{L}_i(\bbx_i^{t+1},\bbz^t,  \lambda_i^t)\|^2\leq \epsilon_i$. In the case of a stochastic gradient oracle being available, this condition instead may be stated as $\mathbb{E}\left[\|\nabla_{\bbx_i}\mathcal{L}_i(\bbx_i^{t+1},\bbz^t,  \lambda_i^t)\|^2\right]\leq \epsilon_i$. We note that any iterative optimization algorithm can be used to perform the $K_i$ local updates. The number of local updates $K_i$ depends upon the accuracy parameter $\epsilon_i$. For instance, in case of non-convex local objective, a gradient descent based oracle would need $K_i=\mathcal{O}\left(1/\epsilon_i\right)$ and a stochastic gradient descent based oracle would require $K_i=\mathcal{O}\left(1/\epsilon_i^2\right)$ number of local steps -- see \citep{wright1999numerical}. A gradient descent based iteration as an instance of \eqref{Proposed_primal} is given in Algorithm \ref{oracle_1}.}
Next, we present the Lagrange multiplier updates initially under the hypothesis that all devices communicate, which we will subsequently relax. In particular, after collecting the locally updated variables at the server $\bbx_i^{t+1}$,  the dual variable is updated via a gradient ascent step given by:
\begin{align}\label{Proposed_dual}
	{\lambda_i^{t+1}=}& \mathcal{P}_{\Lambda}\left[{\lambda_i^{t}+\alpha (\|\bbx_i^{t+1}-\bbz^{t}\|^2-\gamma_i)}\right].  &[\text{{Dual-local update}}],
\end{align}
where the dual variable $\lambda_i^{t+1}$ is projected ($\mathcal{P}_{\Lambda}$ denotes projection operation) onto a compact domain given by $\Lambda:=[\lambda_{\min},\lambda_{\max}]$, where the values of $\lambda_{\min}$ and $\lambda_{\max}$ will be derived later from the analysis. Then, we shift to minimization with respect to global model variable $\bbz$, which by the strong convexity of the Lagrangian [cf. \eqref{min_max}] in this variable, may be obtained as
\begin{align}\label{server_udpate0}
			\bbz{^{t+1}}=&\frac{1}{\sum_{i=1}^{N}(\lambda_i^{t+1}) }\sum_{i=1}^{N}(\lambda_i^{t+1} )\bbx_i^{t+1}\ \ \ \ \ \ &[\text{{Server update}}],
\end{align}
which we obtain by equating $\nabla_{\bbz} \mathcal{L}(\bbz, \{\bbx_i^{t+1}, \lambda_i^{t+1}\}_{i=1}^N) =0$ and then solving for $\bbz$.
The server update requires the access to all local models $\bbx_i$ and Lagrange multipliers $\lambda_i$  in order to execute \eqref{server_udpate0}.
To perform the update \eqref{server_udpate0} we use device selection as is common in FL, we uniformly sample a set of $|\mathcal{S}_t|$ devices from $N$ total devices.
	\begin{algorithm}[t]
		\caption{\textbf{Fed}erated Learning \textbf{B}eyond \textbf{C}onsensus (\texttt{FedBC})}
		\label{alg:alg_FedBC}
		\begin{algorithmic}[1]
			\STATE \textbf{Input}: $T$, $K_i$ for each device $i$, $\gamma_i$ step size parameters $\alpha$ and $\beta$.
			\STATE \textbf{Initialize}: $\bbx_i^{0}$, $\bbz^0$, and $\lambda_i^0$ for all $i$.
			\FOR{$t=0$ to $T-1$}
			\STATE \textbf{Select} a subset of $S$ devices uniformly from $N$ devices , we get $\mathcal{S}_t\in \{1,2,\cdots,N \}$
			\STATE Send $\bbz^t$ to each $j\in\mathcal{S}_t$
			\STATE \textbf{Parallel} loop for each device  $j\in\mathcal{S}_t$

		\STATE	\hspace{0.6cm}	Primal update:	$\bbx_j^{t+1}=\text{Oracle}_j (\mathcal{L}_j(\bbx_j^t,\bbz^t,  \lambda_j^t),K_j)$ {according to Algorithm~\ref{oracle_1}}

		\STATE \hspace{0.6cm} Dual update: $ {\lambda_j^{t+1}=} \mathcal{P}_{\Lambda}\left[{\lambda_i^{t}+\alpha
		(\|\bbx_i^{t+1}-\bbz^{t}\|^2-\gamma_i)}\right] $
		%
		\STATE \hspace{0.6cm} Each device $j$ sends  $\bbx_j^{t+1}, \lambda_j^{t+1}$ back to server
			%
		\STATE \textbf{Server}  updates
		\begin{align}\label{eq:server_update}
		    \bbz{^{t+1}}=\frac{1}{\sum_{j\in \mathcal{S}_t }\lambda_j^{t+1} }\sum_{j\in \mathcal{S}_t }(\lambda_j^{t+1} \bbx_j^{t+1})
		\end{align}
			%
			\ENDFOR
			%
		\STATE		\textbf{Output}: $\bbz{^{T}}$
		\end{algorithmic}
	\end{algorithm}
	%
		%

\textbf{Reduction to Existing Approaches:} FL algorithms alternate between localized updates and server-level information aggregation. The most common is FedAvg \citep{mcmahan2017communication}, which is an instance of  \texttt{FedBC} with  $\lam_i^t=0$ for all $i$. Furthermore,
FedProx is an augmentation of FedAvg with an additional proximal term in the device loss function. Observe that \texttt{FedBC} algorithm with $\lam_i^t=\mu$ for all $i$ and $t$ reduces to FedProx \citep{li2018federated} for \eqref{eq_main_problem}. Furthermore, for $\gamma_i=0$ and without device sampling, the algorithm would become a version of FedPD \citep{zhang2021fedpd}. For constant $\lambda_i=c$ and with $K_i=1$ local GD step, our algorithm reduces to L2GD \citep{hanzely2020federated}, which is limited to convex settings. \citep{li2021ditto} similarly mandates constant Lagrange multipliers and $K_i=1$.

In prior approaches, an inherent assumption is that all local devices would like to converge to the common global model. By contrast, in this work, we introduce personalization via allowing the local device model $\bbx_i$ to be different but close to the global model $\bbz$ at the server in an explicitly tunable manner through the tolerance parameter $\gamma_i$ in \eqref{eq_consensus_reformulation} which is automatically solved through Lagrange multiplier updates. Different from meta learning where there are a distinct collection of tasks and one would like a model that performs well on all subtasks, the goal in this paper is to perform well on the local objective and remain close to the global model. Through the node-separability of the Lagrangian, we rigorously formalize that this approach yields performance guarantees in terms of the global model and the personalization in terms of the local models as a function of $\gamma_i$ in the next subsection.

\section{Convergence Analysis}\label{sec:convergence}
In this section, we establish performance guarantees of Algorithm \ref{alg:alg_FedBC} in terms of trading off solving the global [cf. \eqref{eq_main_problem}] and the local problem \eqref{eq_local_problem}, i.e., personalization, by linking this tradeoff to primal stationarity versus constraint violation [cf. \eqref{eq_consensus_reformulation}]. We first state the assumptions:
\begin{assumption}\label{assum_0}
The domain $\mathcal{X}$  of functions $f_i$ in \eqref{eq_local_problem} is compact with diameter $R$.
\end{assumption}
\begin{assumption}[Lipschitz gradients]\label{assum_1}
The gradient of the local objective $\nabla f_i(\bbx)$ of each device is Lipschitz continuous, i.e.,
    $\|\nabla f_i(\bbx)-\nabla f_i(\bby)\|\leq L_i \|\bbx-\bby\|$, $\forall\bbx,\bby\in\mathcal{X} \;$ .
\end{assumption}%
\begin{assumption}[{Bounded} Heterogeneity]\label{assum_3}
For any device pair $(i,j)$, it holds that
  $ \max_{(a,b)\in\Lambda} \|a\nabla f_i(\bbx)-b\nabla f_j(\bbx) \| \leq \delta$,
for all $\bbx\in\mathcal{X}$.
\end{assumption}
Next, we describe the assumption required when we use stochastic gradients instead of the actual gradients. If we denote the stochastic gradient for agent $i$ as $\bbg_i$, then it satisfies the following assumption.
		\begin{algorithm}[t]
		\caption{Oracle$_i$ {in Equation~\eqref{Proposed_primal}  [\text{{$K_i$-local updates}}]}}
		\label{oracle_1}
		\begin{algorithmic}[1]
			\STATE \textbf{Input}: $K_i$, $\gamma_i$, $\beta$, $\bbx_i^{t},\bbz^{t},\lambda_i^t$
			\STATE \textbf{Initialize}: $\bbw_i^{0}=\bbx_i^{t}$
			\FOR{$k=0$ to $K_i-1$ for each device $i$ }
			\STATE		Update the local model via any optimizer

			\hspace{2cm} \textbf{GD optimizer} :
			$	\bbw_i^{k+1}=\bbw_i^{k} - \beta \left(\nabla_{\bbw} f_i(\bbw_i^{k}) +(2\lambda_i^t)\left(\bbw_i^{k}-\bbz^t\right)\right)$

			\hspace{2cm} \textbf{SGD optimizer} :
			$	\bbw_i^{k+1}=\bbw_i^{k} - \beta \left(\bbg_i^k +(2\lambda_i^t)\left(\bbw_i^{k}-\bbz^t\right)\right)$

			\ENDFOR

		\STATE		\textbf{Output}: $\bbw_i^{K_i}$
		\end{algorithmic}
	\end{algorithm}
	%
\begin{assumption}[Stochastic Gradient Oracle]\label{assum_2}
If a stochastic gradient oracle is used at device $i$, then $\bbg_i$ satisfies
 $   \mathbb{E}[\bbg_i~|~\mathcal{H}_k] = \nabla f(\bbx_i)$, \ \ \text{and} \ \  $\mathbb{E}[\|\bbg_i-\nabla f(\bbx_i)\|^2~|~\mathcal{H}_k] \leq \sigma^2$, \ \ $\forall{i}$,
 where $\mathcal{H}_k$ is defined as filtration or $\sigma$-algebra generated by past realizations $\{\zeta_i^u\}_{u<k}$.
\end{assumption}
We note that the Assumptions \ref{assum_0}-\ref{assum_2} are standard \citep{nemirovski2009robust}.
Assumption \ref{assum_1} makes sure that the local non-convex objective are smooth with parameter $L_i$. Assumption \ref{assum_3} is a version of the heterogeneity assumption considered in the literature (Assumption 3 in \citep{t2020personalized}). Assumption \ref{assum_2} imposes conditions on the stochastic gradient oracle, particularly, unbiasedness and finite variance, which are standard.

 We are now  ready to present the main results of this work in the form of Theorem \ref{main_theorem}.
For the convergence analysis, we consider the performance metric $\frac{1}{T}\sum_{t=1}^{T}\|\nabla f(\bbz^t)\|^2$ which is widely used in the literature \citep{mcmahan2017communication,li2018federated,karimireddy2020scaffold,t2020personalized}. Under these conditions, we have the following convergence result.

\begin{theorem}\label{main_theorem}
	Under Assumption \ref{assum_0}-\ref{assum_3}, for the iterates of proposed Algorithm \ref{alg:alg_FedBC}, we establish that the global performance satisfies:
		%
		\begin{align}\label{main_result_1}
	\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}[\|\nabla f (\bbz^t)\|^2]
	= & \mathcal{O}\left(\frac{B_0}{T}\right)+\mathcal{O}(\epsilon)+ \mathcal{O}(\delta^2)+ \mathcal{O}(\alpha)+ \mathcal{O}\left(\frac{1}{ NT}\sum_{i=1}^N \gamma_i\right),
		\end{align}
		%
		where $B_0$ is the initialization dependent constant, $\epsilon=\max_{i}\epsilon_i$ is the accuracy with each agent solves the local optimization problem in the algorithm, $\delta$ is the heterogeneity parameter (sf. Assumption \ref{assum_3}), $\alpha>0$ is the step size, and $\gamma_i$'s are the local personalization parameters.
\end{theorem}

The proof of Theorem \ref{main_theorem} is provided in Appendix \ref{proof_theorem}.The expectation in \eqref{main_result_1} is with respect to the randomness in the stochastic gradients and device sampling. The first terms is initialization dependent term, and as long as the initialization $B_0$ is bounded, the first term reduces linearly with respect to $T$ and goes to zero in the limit as $T\rightarrow\infty$.  This term is present in any state of the art FL algorithm \citep{mcmahan2017communication,li2018federated,karimireddy2020scaffold,t2020personalized}. The second term is $\mathcal{O}(\epsilon)$ which depends upon the worst local approximated solution across all the devices. Note that the individual approximation errors $\epsilon_i$ depend on the number of local iterations $K_i$ as well. This term is also present in most of the analysis of FL algorithms for non-convex objectives. The third term is due to the heterogeneity across the devices and is a specific feature of the FL problem. The fourth terms is the step size dependent term. The last term is important here because that appears due to the introduction of $\gamma_i$ in the problem formulation in \eqref{eq_consensus_reformulation}, and is completely novel to the analysis in this work. This term decays linearly even if $\gamma_i>0$ for all $i$. The $\gamma_i$'s are directly affecting the global performance because they are allowing client models to move away from each other, and hence affecting the global performance. We remark that for the special case of $\gamma_i=0$ for all $i$, our result in \eqref{main_theorem} is equivalent to FedPD \citep{zhang2021fedpd}, pFedMe \citep{t2020personalized} except for the fact that there is no device sampling in FedPD.

The technical points of departure in the analysis of \texttt{FedBC} (cf. Algorithm \ref{alg:alg_FedBC}) from prior work are associated with the fact that we build out from an ADMM-style analysis. \citep{zhang2021fedpd} However, due to non-linear constraint \eqref{eq_consensus_reformulation}, one cannot solve the $\operatornamewithlimits{argmin}$ exactly. This introduces an additional $\mathcal{O}(\epsilon)$ error term that we relate to $K_i$ in \eqref{Proposed_primal}. This issue also demands we constrain the dual variables to a compact set in \eqref{Proposed_dual}. Moreover, device sampling for the server update (cf. \eqref{eq:server_update}) is introduced here for the first time in a primal-dual framework, which does not appear in \citep{zhang2021fedpd}.

Our nonlinear proximity constraints [cf. \eqref{eq_consensus_reformulation}] additionally permit us to relate the performance in terms of the local objective [cf. \eqref{eq_local_problem}] to the proximity to the global model defined in \eqref{eq_consensus_reformulation} as a function of tolerance parameter $\gamma_i$. We formalize their interconnection in the following corollary.

\begin{corollary}[Local Performance]\label{corollary} 	Under Assumption \ref{assum_0}-\ref{assum_3}, for the iterates of proposed Algorithm \ref{alg:alg_FedBC}, we establish that
\begin{align}
     \frac{1}{T}\sum_{t=1}^T\|\nabla f_i(\bbx_i^{t})\|^2  \leq  \mathcal{O}\left(\epsilon_i\right)+\alpha^2 \mathcal{O}\left(  \bigg[\sum_{k=0}^{T}\mathbb{I}_{\{\|\bbx_i^{k}-\bbz^{k-1}\|^2\leq \gamma_i)\}}\bigg]_+\right)^2. \label{zero0000000}
\end{align}
\end{corollary}
The proof is provided in Appendix \ref{proof_corollary}. We note that the local stationarity of each client $i$ actually depends on the local $\epsilon_i$ approx error and $\gamma_i$ via a complicated term present in the second term in \eqref{zero0000000}, where $\mathbb{I}$ is an indicator function that is $1$ if condition is not satisfied, and $-1$ otherwise. We note that the term inside the big bracket is larger (worse local performance) for lower $\gamma_i$, and vice versa.  Hence we have a trade-off between global and local model performance in terms of $\gamma_i$.

\section{Experiments}\label{sec:experiments}
In this section, we aim to address the following questions with our experiments:
\textit{\ding{172} How does \texttt{FedBC} compare to other federated learning algorithms in heterogeneous environments?} \textit{\ding{173} Does \texttt{FedBC} allow users to have their own localized models and to what extent?} Incidentally, we came across an additional phenomenon that may be of independent interest: Lagrange multiplier updates tend to more equitably weigh the importance of each device's variables, dependent on the number of locally available training samples.
\begin{table}[t]
  \centering
  \caption{
  \footnotesize{
    Synthetic dataset classification accuracy for different number of local training epochs ($E$), i.e. $K_i = E$ for each device $i$ in Algorithm \ref{alg:alg_FedBC}. The $\pm$ shows standard deviation.
  }
  }

\resizebox{0.8\textwidth}{!}{\begin{tabular}{|l|lllll|}
\hline
\multirow{2}{*}{Algorithm} & \multicolumn{5}{c|}{Epochs}                                                                                \\ \cline{2-6}
                           & \multicolumn{1}{c|}{1}                         & \multicolumn{1}{c|}{5}                         & \multicolumn{1}{c|}{10}                        & \multicolumn{1}{c|}{25}                        & \multicolumn{1}{c|}{50}   \\ \hline
FedAvg                     & \multicolumn{1}{l|}{83.61 $\pm$ 0.43}          & \multicolumn{1}{l|}{83.42 $\pm$ 0.58}          & \multicolumn{1}{l|}{83.49 $\pm$ 0.70}          & \multicolumn{1}{l|}{83.73 $\pm$ 0.57}          & 82.94 $\pm$ 0.66          \\ \hline
q-FedAvg                   & \multicolumn{1}{l|}{87.12 $\pm$ 0.25}          & \multicolumn{1}{l|}{86.76 $\pm$ 0.15}          & \multicolumn{1}{l|}{86.46 $\pm$ 0.28}          & \multicolumn{1}{l|}{86.76 $\pm$ 0.12}          & 86.71 $\pm$ 0.12          \\ \hline
FedProx                    & \multicolumn{1}{l|}{86.23 $\pm$ 0.42}          & \multicolumn{1}{l|}{85.59 $\pm$ 0.37}          & \multicolumn{1}{l|}{85.34 $\pm$ 0.61}          & \multicolumn{1}{l|}{85.00 $\pm$ 0.63}          & 85.34 $\pm$ 0.48          \\ \hline
Scaffold                   & \multicolumn{1}{l|}{83.84 $\pm$ 0.09}          & \multicolumn{1}{l|}{82.95 $\pm$ 0.30}          & \multicolumn{1}{l|}{83.48 $\pm$ 0.21}          & \multicolumn{1}{l|}{83.60 $\pm$ 0.21}          & 82.80 $\pm$ 0.62          \\ \hline
\cellcolor{LightCyan}\texttt{FedBC}                      & \multicolumn{1}{l|}{\textbf{87.83 $\pm$ 0.35}} & \multicolumn{1}{l|}{\textbf{87.48 $\pm$ 0.18}} & \multicolumn{1}{l|}{\textbf{87.43 $\pm$ 0.20}} & \multicolumn{1}{l|}{\textbf{86.99 $\pm$ 0.11}} & \textbf{87.26 $\pm$ 0.12} \\ \hline
\end{tabular}
}
\label{table: synthetic}
\end{table}

\textbf{Experiment Setup.} The synthetic dataset is associated with a $10$-class classification task, and is adapted from \citep{li2018federated} with parameters $\alpha$ and $\beta$ controlling model and data variations across users (see Appendix \ref{experimental_Setup} for details). For real datasets, we use MNIST and CIFAR-$10$ for image classification. MNIST and CIFAR-10 datasets consist of handwritten digits and color images from $10$ different classes respectively \citep{krizhevsky2009learning} \citep{lecun1998gradient}. We denote $C$ to represent the most common number of classes in users' local data (see Appendix \ref{data_partition} for details). To evaluate the global performance of \texttt{FedBC}, we compare it with $4$ other federated learning algorithms, i.e., FedAvg \citep{mcmahan2017communication}, q-FedAvg \citep{li2019fair}, FedProx \citep{li2018federated}, and Scaffold \citep{karimireddy2020scaffold}. Since we do not know the optimal $\gamma_i$ for each user $i$ apriori, we initialize them to $0$, and let the user decide its \textbf{own} $\gamma_i$'s magnitude (see Appendix \ref{experimental_Setup} for details).  To further enhance  personalization, we fine-tune the local models of \texttt{FedBC}  by performing $1$-step gradient descent using a batch of test data, and name this approach as \texttt{FedBC}-FineTune. In the end, we also incorporate MAML-type training into \texttt{FedBC}, and derive another personalization algorithm Per-\texttt{FedBC} (see Appendix \ref{per_fedbc}). We use the term global accuracy when evaluating the performance of global model using the entire test dataset. We use the term local accuracy when evaluating the performance of each users' local model using its own test data and take the average across all users. Similar to \texttt{FedBC}-FineTune, we also perform the $1$ fine-tuning step using test data for Per-\texttt{FedBC} and Per-FedAvg following \citep{fallah2020personalized} (see Appendix \ref{hyper_tuning} for details about hyperparameter tuning). \footnote{Code is available at \url{https://github.com/chenfan95/FedBC}}

\textbf{Proof of Concept:} Table \ref{table: synthetic} presents results synthetic dataset classification. We observe that \texttt{FedBC} consistently outperforms other algorithms for different number of local training epochs $E$. Figure \ref{fig:p1_syn} provides empirical justifications for such remarkable performance. We first observe that \texttt{FedBC} can achieve a lower training loss than others, whereas algorithm such as FedAvg plateaus at an early stage as shown in Figure \ref{fig:loss_syn_e5}. Next, we compare the contributions from min and max users' (defined in Figure \ref{fig:p1_syn} caption) local models in updating the global model $\bbz^t$ via \eqref{eq:server_update}. To this end, we plot their $\lambda_i$s at the end of each communication round in Figure \ref{fig:coef_syn_e5}. We observe that \texttt{FedBC} is significantly less biased towards the min user than other algorithms. The min user's coefficient eventually catches that of max user for \texttt{FedBC}, and the difference between them is one-order of magnitude less than that of data-size based coefficient. Figure \ref{fig:coef_syn_magnitude} further shows the distribution of $\lambda$-based coefficients at the end of training for users of different data sizes. Interestingly, the max user's coefficient is almost the same as those of users of small data sizes. In fact, the coefficients of users of data sizes less than $300$ for \texttt{FedBC} are consistently larger than their data-size-based counterparts, and vice versa for data sizes greater than $300$. Lastly, Figure \ref{fig:coef_syn_variance} shows the model of \texttt{FedBC} achieves a high uniformity in test accuracy over users' local data at different $E$s.
\begin{figure}[t]
     \centering
     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{loss_syn_e5.pdf}
         \caption{Train loss}
         \label{fig:loss_syn_e5}
     \end{subfigure}
     \hfill
     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{coef_syn_e5.pdf}
         \caption{Max/min coefficient}
         \label{fig:coef_syn_e5}
     \end{subfigure}
     \hfill
     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{coef_syn_magnitude.pdf}
         \caption{Coefficient vs $n_i$}
         \label{fig:coef_syn_magnitude}
     \end{subfigure}
    \hfill
     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{variance_syn.pdf}
         \caption{Test accuracy variance}
         \label{fig:coef_syn_variance}
     \end{subfigure}
        \caption{ We use $30$ as the total number of users and $E = 5$. (a) We plot train loss vs number of rounds, and observe that \texttt{FedBC} achieves the lowest train loss. (b) We track the user with the smallest (min user) or largest (max user) number of data at each round, and plot min or max user's coefficient in computing the global model based on either its local dataset size ($n_i$) or $\lambda_i$. We observe that the magnitude difference between min and max user's coefficient based on $\lambda_i$ is consistently smaller than that based on $n_i$. (c) We plot $\lambda_i$ against $n_i$ for each user at the end of training, and observe that users of small dataset sizes (e.g. $<200$) can contribute significantly to the global model in \texttt{FedBC}. (d) We show the variance of test accuracy of global model on each user's local data for different $E$s (shaded area shows standard deviation), and observe that the model of \texttt{FedBC} achieves high uniformity in terms of this performance measure.}

        \label{fig:p1_syn}
        \vspace{-6mm}
\end{figure}
In addition to the results in Figure \ref{fig:p1_syn}, we track and plot the test accuracy of global model on min and max users' local data throughout the entire communication in Figure \ref{fig:p2_syn}. We observe that the global model performs better on max user's data than on min user's data in general for all algorithms. Most importantly, \texttt{FedBC} demonstrates superior advantage in reducing this performance gap. After $100$ rounds of communication, test accuracy for min and max user are nearly the same as shown in Figure \ref{fig:user_bc5}. Whereas for FedAvg (Figure \ref{fig:user_avg5}), this gap can be as large as $100 \%$ even after nearly $200$ rounds of communication. When the local model differs more from the global model as $E$ increases to $50$, FedBC is still able to make good predictions on the min user's local data despite an increase in the performance gap when compared with $E = 5$ (Figure \ref{fig:user_bc50_main} versus Figure \ref{fig:user_bc5}). This is significantly different from the case of FedAvg in which its model fails to make any correct predictions on the min user's local data for the majority of times as shown in Figure \ref{fig:user_avg50_main}.   Compared to FedProx (Figure \ref{fig:user_prox5}), q-Fedavg (Figure \ref{fig:user_q5}), or Scaffold (Figure \ref{fig:user_sca5}), \texttt{FedBC} has a much higher fraction of points at which test accuracy for min and max user overlap, which indicates that they are being treated equally well (see Appendix \ref{exp_syn_appendix} for more results). In essence, \texttt{FedBC} has the best performance because it allows all users to fairly participate in updating the global model.
\begin{figure}[t]
     \centering
     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{user_bc5.pdf}
         \caption{\texttt{FedBC} ($E = 5$)}
         \label{fig:user_bc5}
     \end{subfigure}

     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{user_avg5.pdf}
         \caption{FedAvg ($E = 5$)}
         \label{fig:user_avg5}
     \end{subfigure}

     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{user_bc50.pdf}
         \caption{\texttt{FedBC} ($E = 50$)}
         \label{fig:user_bc50_main}
     \end{subfigure}

     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{user_avg50.pdf}
         \caption{FedAvg ($E = 50$)}
         \label{fig:user_avg50_main}
     \end{subfigure}

     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{user_qavg5.pdf}
         \caption{q-FedAvg ($E = 5$)}
         \label{fig:user_q5}
     \end{subfigure}

     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{user_prox5.pdf}
         \caption{FedProx ($E = 5$)}
         \label{fig:user_prox5}
     \end{subfigure}

     \begin{subfigure}[b]{0.24\textwidth}
         \centering
\includegraphics[width=\textwidth]{user_sca5.pdf}
         \caption{Scaffold ($E = 5$)}
         \label{fig:user_sca5}
     \end{subfigure}
        \caption{ Test accuracy of global model on \textcolor{blue}{min} and \textcolor{orange}{max} (defined in Figure \ref{fig:p1_syn}) users' local data at each round of communication.
        The global model of \texttt{FedBC} (a) initially has a performance gap on min and max users' data, but the gap is largely eliminated in the end. For FedAvg (b), q-FedAvg (e), FedProx (f), and Scaffold (g), this gap is more apparent and persistent throughout training.}

        \label{fig:p2_syn}
                \vspace{-4mm}
\end{figure}

\textbf{Real Dataset Experiments:} The experiments on real datasets are inline with our previous observations in Figure \ref{fig:p1_syn}. We first observe the unbiasedness in computing the coefficient for the global model of FedBC when $C = 1$, $2$, or $3$ as shown in Figure \ref{fig:coef_cifar} in the appendix. We also observe the the high uniformity in test accuracy over users' local data for all classes with different power law exponents as shown in Figure \ref{fig:var_cifar_e5} in the appendix. These nice features of \texttt{FedBC} are reflected in its good test performance on CIFAR-10 as shown in Table \ref{table: table_cifar}. For example, FedBC outperforms FedAvg by $7.89 \%$ when $C = 3$. In the most challenging situation when $C = 1$, it also outperforms all other baselines (see Table \ref{table: cifar_5_appendix} in the appendix for the results of other power law exponents). Table \ref{table: mnist_1} and Table \ref{table: mnist_5} in the appendix show classification results on MNIST, and we observe similar trends as in the case of CIFAR-10.
We also notice the global performance of pFedMe and Per-FedAvg is worse than that of FedAvg, FedProx or \texttt{FedBC} when $C = 1$. Personalized algorithms are designed to optimize the local objective of \eqref{eq_local_problem}. However, this may create conflicts to the global objective \eqref{eq_main_problem} and lead to a poor global performance. Table \ref{table: pers_cifar} presents the results of local performance of \texttt{FedBC} and other personalization algorithms. We observe that \texttt{FedBC} performs better and worse than pFedMe and Per-FedAvg respectively. By doing $1$-step fine-tuning on the test data, i.e. \texttt{FedBC}-FineTune, we are able to improve the local performance of \texttt{FedBC}. For example, \texttt{FedBC}-FineTune achieves a $7.55 \%$ performance increase over \texttt{FedBC} when $C = 3$. By incorporating MAML-type training into \texttt{FedBC}, i.e. Per-\texttt{FedBC}, it achieves the overall best performance for both $C = 1$ and $C = 3$.

 \textbf{Takeaways:} In general, we have shown that \texttt{FedBC} has the best global performance when compared against all baselines. This addresses question \textit{\ding{172}}. It has reasonably good local performance, but can be improved with either fine-tuning or performing MAML-type training. This addresses question \textit{\ding{173}}. We leave the question of how to fully exploit the advantage in the freedom of choosing $\gamma_i$ to achieve personalization in a future investigation.

\begin{table}[t
    \centering
    \caption{ Global (a) and local (b) performance of \texttt{FedBC} and other baselines on CIFAR-10 classification. The power law exponent is $1.5$ and $E = 5$. All colored cells denote the proposed algorithms (see Algorithm \ref{alg:perfedbc} in the appendix for Per-FedBC).}
    \begin{subtable}[t]{1.0\textwidth}
        \centering
  \caption{Top-1 classification accuracy of global model}
\resizebox{\textwidth}{!}{\begin{tabular}{|l|ccccccc|}
\hline
                          & \multicolumn{7}{c|}{Algorithm}                                                                                \\ \cline{2-8}
\multirow{-2}{*}{Classes} & \multicolumn{1}{c|}{FedAvg}           & \multicolumn{1}{c|}{q-FedAvg}         & \multicolumn{1}{c|}{FedProx}          & \multicolumn{1}{c|}{Scaffold}         & \multicolumn{1}{c|}{Per-FedAvg}       & \multicolumn{1}{c|}{pFedMe}            & \cellcolor{LightCyan}FedBC \\ \hline
C = 1                     & \multicolumn{1}{c|}{60.82 $\pm$ 0.85} & \multicolumn{1}{c|}{60.56 $\pm$ 1.19} & \multicolumn{1}{c|}{60.57 $\pm$ 0.65} & \multicolumn{1}{c|}{36.08 $\pm$ 3.25} & \multicolumn{1}{c|}{56.31 $\pm$ 5.52} & \multicolumn{1}{c|}{50.19 $\pm$ 2.24}  & \textbf{61.12 $\pm$ 1.24}     \\ \hline
C = 3                     & \multicolumn{1}{c|}{58.91 $\pm$ 2.76} & \multicolumn{1}{c|}{64.01 $\pm$ 1.96} & \multicolumn{1}{c|}{59.85 $\pm$ 2.81} & \multicolumn{1}{c|}{33.04 $\pm$ 9.85} & \multicolumn{1}{c|}{59.27 $\pm$ 2.46} & \multicolumn{1}{c|}{53.26 $\pm$  1.24} & \textbf{66.80 $\pm$ 1.07}     \\ \hline
\end{tabular}
}
\label{table: table_cifar}
\end{subtable}
\begin{subtable}[t]{.8\textwidth}
  \centering
  \caption{Top-1 classification accuracy of local model}
\resizebox{\columnwidth}{!}{\begin{tabular}{|l|lllcc|}
\hline
                          & \multicolumn{5}{c|}{Algorithm}                                                                                \\ \cline{2-6}
\multirow{-2}{*}{Classes} & \multicolumn{1}{c|}{Per-FedAvg}       & \multicolumn{1}{c|}{pFedMe}           & \multicolumn{1}{c|}{\cellcolor{LightCyan}FedBC} & \multicolumn{1}{c|}{\cellcolor{LightCyan}FedBC-FineTune} & \cellcolor{LightCyan}Per-FedBC \\ \hline
C = 1                     & \multicolumn{1}{l|}{85.09 $\pm$ 1.42} & \multicolumn{1}{l|}{76.94 $\pm$ 0.82} & \multicolumn{1}{l|}{81.31 $\pm$ 0.22}              & \multicolumn{1}{c|}{82.01 $\pm$ 0.17}                       & \textbf{85.18 $\pm$ 1.08}         \\ \hline
C = 3                     & \multicolumn{1}{l|}{81.63 $\pm$ 3.32} & \multicolumn{1}{l|}{59.71 $\pm$ 0.50} & \multicolumn{1}{l|}{65.05 $\pm$ 1.00}              & \multicolumn{1}{c|}{72.60 $\pm$ 0.34}                       & \textbf{83.40 $\pm$ 1.41}         \\ \hline
\end{tabular}
}
\label{table: pers_cifar}
    \end{subtable}
\end{table}

\section{Conclusions}\label{sec:conclusions}

In this work, we cast the tradeoff between local and global model performance in federated learning as a multi-criterion optimization problem that well-motivates a constrained optimization approximation. Then, we considered the Lagrangian relaxation of this problem, to which we applied primal-dual method to derive \texttt{FedBC}. We established that \texttt{FedBC} exhibits favorable performance in practice across a variety of settings, automating the tradeoff of personalization and generalization across devices more effectively than several state of the art alternatives.

\bibliographystyle{plain}