# Federated Deep Equilibrium Learning: A Compact Shared Representation for
Edge Communication Efficiency

IEEE Publication Technology This paper was produced by the IEEE Publication
Technology Group. They are in Piscataway, NJ.Manuscript received April 19,
2021; revised August 16, 2021.

###### Abstract

Federated Learning (FL) is a prominent distributed learning paradigm
facilitating collaboration among nodes within an edge network to co-train a
global model without centralizing data. By shifting computation to the network
edge, FL offers robust and responsive edge-AI solutions and enhance privacy-
preservation. However, deploying deep FL models within edge environments is
often hindered by communication bottlenecks, data heterogeneity, and memory
limitations. To address these challenges jointly, we introduce FeDEQ, a
pioneering FL framework that effectively employs deep equilibrium learning and
consensus optimization to exploit a compact shared data representation across
edge nodes, allowing the derivation of personalized models specific to each
node. We delve into a unique model structure composed of an equilibrium layer
followed by traditional neural network layers. Here, the equilibrium layer
functions as a global feature representation that edge nodes can adapt to
personalize their local layers. Capitalizing on FeDEQ’s compactness and
representation power, we present a novel distributed algorithm rooted in the
alternating direction method of multipliers (ADMM) consensus optimization and
theoretically establish its convergence for smooth objectives. Experiments
across various benchmarks demonstrate that FeDEQ achieves performance
comparable to state-of-the-art personalized methods while employing models of
up to 4 times smaller in communication size and 1.5 times lower memory
footprint during training.

###### Index Terms:

Federated Learning, Distributed Optimization, Equilibrium Models, Edge
Networks.

##  I Introduction

With the rapid proliferation of Internet-connected devices and the consequent
vast data they generate, the shortcomings of centralized computing are
becoming increasingly evident. Edge networks have emerged as a compelling
countermeasure, processing data closer to its source. Such networks capably
manage the escalating data volumes, inherently reducing latency, optimizing
bandwidth usage and improving real-time responsiveness. In tandem with this
shift towards decentralized processing, Federated Learning (FL) has been
recognized as an apt machine learning paradigm for the network edge.
Contrasting with traditional methods that necessitate transferring all data to
a central server, FL enables multiple edge nodes to collaborate on training a
unified machine learning model. Here, each node refines the model using its
local data and transmits only the model’s updates to a central server, where
they are aggregated to improve the global model [1]. This not only aligns with
privacy requirements since data never leaves the network edge, but also
provides scalable solutions tapping into the distributed processing power of
edge devices.

Despite offering a scalable, privacy-preserving solution that taps into the
distributed processing power of edge devices, fully realizing FL’s potential
within complex edge environments is contingent upon overcoming three major
challenges. First, there are communication bottlenecks, resulting from the
need to transmit large model parameters between parties. These constraints can
lead to significant bandwidth limitations, causing inefficiencies and delays
in training. Second, data heterogeneity arises from the variety and diversity
of the data held by edge nodes, leading to client drift effects where
individual model performances diverge. Finally, the limitations in memory and
storage constrain the complexity and size of the models that can be deployed,
requiring innovative solutions to maximize the potential of FL in edge network
scenarios.

Recent literature has presented diverse approaches to address the foregoing
challenges, as outlined in Fig. 1. The communication and memory challenges has
been tackled using model compression techniques [2, 3] to reduce model size
for facilitating transmission and adaptive gradients [4] for efficient
parameter updates. On the other hand, personalization has emerged as a key
strategy to address the issue of data heterogeneity and client drift. Full
model personalization methods, such as meta-learning [5, 6] and multi-task
learning [7, 8, 9], have shown efficiency in capturing unique local patterns.
However, these techniques can be resource-intensive and may necessitate
considerable communication overhead. Conversely, partial model personalization
approaches [10, 11] offer a balanced solution by employing a shared
representation complemented by personalized local models, enabling both
customization and conservation of resources. While these methods typically
tackle individual challenges, a comprehensive approach that can addressing
both communication and memory constraints, as well as data heterogeneity
within FL in edge networks remains yet to be thoroughly explored in research.

Figure 1: FeDEQ aims to simultaneously address data heterogeneity and
communication and resource constraints, while other methods focus on solving
individual challenge.

In this paper, we address a pivotal question: ”How do we leverage FL to
exploit a compact shared representation under data heterogeneity in edge
networks that not only substantially reduces communication bandwidth and
memory footprint but effectively enhance learning performance?”. To this end,
we propose a novel framework, namely _FEderated Deep EQuilibrium Learning_
(FeDEQ), that leverages consensus optimization to learn a shared
representation via deep equilibrium models (DEQs), combined with personalized
explicit layers for each edge node. The compact shared layer can capture
diverse patterns across edge nodes, while the explicit layers are fine-tuned
for personalization using local data. This combination enables us to
effectively handle the aforementioned challenges: FeDEQ performs on par with
state-of-the-art (SOTA) approaches while reducing communication by 2–4 times
and memory footprint by 1.5 times. The main contributions of this work are
summarized as follows.

  * •

We introduce FeDEQ, a novel framework tailored for resource-efficient and
personalized FL under data heterogeneity in edge networks. Central to FeDEQ’s
design is its ability to learn a compact shared representation with constant
memory footprint via equilibrium models, thereby adapting this shared base to
swiftly fine-tune personalized layers at individual edge node to achieve local
personalization.

  * •

To collaboratively learn shared and personalized layers across edge nodes, we
devise a consensus optimization scheme grounded in the alternating directions
method of multipliers (ADMM). This strategically decomposes the original
consensus problem into the primal problem – seeking optimal personalized and
shared parameters, and the dual problem – controlling discrepancies between
the local and shared representation to mitigate the drifting effect in local
training.

  * •

We theoretically provide the convergence guarantees for the proposed algorithm
for smooth objectives. Under mild assumptions, we show that FeDEQ converges to
a stationary point, which achieves a shared representation across all edge
nodes and optimal personalized parameters for each node.

  * •

Through extensive experimentation on several FL benchmarks, we substantiate
that FeDEQ achieves personalization with communication and memory efficiency.
Further, the shared layer is highly adaptable to new edge nodes in
generalization experiments.

To the best of our knowledge, this is the first work that employs deep
equilibrium learning and ADMM consensus optimization for simultaneously
tackling communication bottlenecks, memory constraints and data heterogeneity
in FL for edge networks.

##  II Related Work

###  II-A Federated learning at the Network Edge.

The emergence of FL comes in response to three significant challenges in
large-scale machine learning, particularly pertinent to edge networks:
processing massive amounts of edge data, managing the heavy communication load
within networks, and preserving privacy without resorting to centralized
storage. The standard federated algorithm - FedAvg [1], which utilizes local
stochastic gradient descent (SGD) and averaging, has established the
foundational framework for much of the FL research in edge networks.

The primary challenges of FL in the context of edge networks can be broadly
classified into communication and resource bottlenecks and data heterogeneity.
Communication and memory bottlenecks pertain to the large number of model
parameters that need to be exchanged and large models that need to be trained
on edge nodes. These challenges have prompted prior works to create more
communication and computation-efficient strategies including: quantization
[12, 13, 14], where local computations reduce the precision of model updates
to reduce the size of data transferred; model compression [2], which involves
reducing the size of the model to make it more manageable to transmit; and
adaptive communication [4], where the amount of data exchanged is adjusted
dynamically based on various factors like network conditions.

The second category, data heterogeneity, is concerned with the non-identically
distributed (non-i.i.d.) data that are often present across edge nodes. This
aspect of the data can affect the convergence and performance of the global
model [15], leading to a variety of approaches to mitigate its impact. Studies
have delved into methods for managing the diversity of data across edge nodes,
incorporating techniques like data sharing [16], data augmentation [17] or
variance reduction [18] to recognize and adapt to the variations in data
across different nodes in the network, ensuring that the globally aggregated
model accurately reflects the underlying patterns of the entire distributed
dataset.

###  II-B Personalized Federated Learning

Personalized Federated Learning (PFL) has recently gained significant
attention to address data heterogeity [19]. Regarding full model
personalization, a simple yet effective personalization approach involves
fine-tuning the global model’s parameters for each client [20, 21]. Various
works adopted multi-task learning (MTL), enabling clients to learn their
distinct data patterns while benefiting knowledge shared by others [7, 8, 9].
Meta-learning frameworks are also employed to develop an effective initial
model to rapidly adapted to new heterogeneous tasks [5, 6], while model
interpolation techniques are used to personalize models using a mixture of
global and local models [22, 23]. Recent years, partial model personalization
approaches [24, 25, 26, 10, 27, 28, 29, 11] have brought up the idea of
exploiting the common representation in heterogeneous settings to learn a
personalized model tailored for each device, enabling simple formulation and
layer-wise flexibility; however, their applicability is limited to specific
settings such as linear representations [10] or pre-trained representations
[28].

###  II-C Deep Equilibrium Learning

Equilibrium models has emerged as a fascinating topic in the realm of implicit
deep learning in recent times. In contrast to the conventional layer stacking
models, deep equilibrium models (DEQ) implicitly define layers through an
equilibrium point (a.k.a. fixed-point) of an infinite sequence of computation
[30]. A general DEQ consists of a single implicit layer $f_{\theta}$ modeled
as a fixed-point system $z^{\star}=f_{\theta}(z^{\star};x)$. Here, $\theta$ is
the model parameter, x is the input data, and the equilibrium $z^{\star}$ is
approximated by fixed-point iterations or root-finding. Several works [31, 32]
delved into the theoretical underpinnings of DEQ, providing valuable insights
into the stability and approximation properties of implicit layers. Recent
works have successfully applied DEQ on large-scale tasks such as sequential
modeling, and semantic segmentation [30, 33]. Despite potential drawbacks such
as slow training time and computational complexity, DEQ possesses the
capability to mitigate communication overhead and memory-intensive deep
models, thereby presenting a compelling solution for FL.

###  II-D ADMM-based learning

ADMM has made significant advancements in both theory and application over
recent decades, particularly in distributed learning [34, 35, 36]. In FL,
several ADMM-based decentralized methods have been developed to enhance
communication efficiency [37, 38]. However, employing this framework for
tackling data heterogeneity is still uncharted.

In this work, we also inherit the ideas of partial personalization that
efficiently learn non-linear representations with a reduced model size using
only edge local data.

##  III Problem Formulation

###  III-A Federated Deep Equilibrium Learning

Consider a federated learning setting within an edge network consisting of a
central server and $n$ edge nodes. Each node $i$ owns a dataset of $n_{i}$
independent and identically distributed (i.i.d.) observations, denoted as
$\\{(x_{ij},y_{ij})\in{\mathcal{X}}_{i}\times{\mathcal{Y}}_{i}:j=1,\ldots,n_{i}\\}$.
This local dataset represents a crucial piece of the learning puzzle,
capturing unique patterns and characteristics inherent to each specific node’s
region or domain. To distill meaningful insights from this data, we introduce
a predictive function $f_{i}:{\mathcal{X}}_{i}\rightarrow{\mathcal{Y}}_{i}$,
tailored to the unique statistical properties of each node’s local data, and a
loss function $\ell(f_{i}(x_{ij}),y_{ij})$ measuring the discrepancy between
$f_{i}(x_{ij})$ and $y_{ij}$. The local empirical risk on node $i$, defined as
${\mathcal{L}}_{i}=\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}\ell(f_{i}(x_{ij}),y_{ij})$,
reflects how well the model is performing within the individual local context.
The primary objective is to minimize the global risk, equal to the sum of all
local risks.

Conventional FL approaches, such as FedAvg [1], learn a global model by
averaging the model parameters obtained from edge nodes after local SGD
rounds. While this method promotes collaboration among nodes, it ignores the
intricate of heterogeneous data which leads to suboptimal performance. A
prominent way to rectify this drawback is through the implementation of
partial personalization, where a model decouples into a _representation
module_ and a _personalized module_. The former serves as a collaborative
learning ground for edge nodes to acquire a common representation (either
linear or non-linear) that captures the shared patterns and regularities
across all nodes’ data. This common learning phase helps maintain a unified
and consistent model view. On the other hand, the _personalized module_ adapts
and refines itself to the unique characteristics of each node’s local data,
thereby enabling better local generalization [24, 10, 28].

(a)

(b)

Figure 2: Using explicit and implicit layers as a shared representation module
for edge nodes in FL. (a) Each node has $M$ explicit layers in its model. (b)
Instead of having multiple explicit layers, each node uses only one
equilibrium layer parametrized by $\theta_{i}$.

While this personalization approach offers a more nuanced handling of the data
and can lead to improved model performance, it is not without challenges. An
effective representation is typically deep and complex, often consisting of
numerous parameters and layers - as illustrated by the representation $r$ in
Fig. 2(a), which includes $M$ layers. This complexity translates into
substantial communication costs and $O(M)$ memory complexity during training.
These costs can prove to be significant barriers, particularly in resource-
constrained edge networks where efficiency is paramount.

In this work, we concurrently tackle these challenges of data heterogeneity,
communication overhead, and memory constraints by leveraging partial
personalization and model compactness through the utilization of deep
equilibrium, as shown in Fig. 1. Specifically, for every node $i$ we consider
a special predictive function $f_{i}$ of the form

$\displaystyle f_{\theta_{i},w_{i}}(x_{ij})$
$\displaystyle\vcentcolon=h_{w_{i}}\big{(}g_{\theta_{i}}(z_{ij}^{\star};x_{ij})\big{)}$
(1) $\displaystyle\text{s.t.}\quad z_{ij}^{\star}$
$\displaystyle=g_{\theta_{i}}(z_{ij}^{\star};x_{ij}),\quad\forall i,j.$

By this definition, $f_{i}$ is a composite of two functions: $g$ parametrized
by $\theta_{i}$ and $h$ by $w_{i}$. A schematic illustration is presented in
Fig. 2(b). While $h_{w_{i}}$ can be any conventional layer (or stack of
layers), such as a fully-connected network, $g_{\theta_{i}}$ is defined
implicitly with the constraint
$z_{ij}^{\star}=g_{\theta_{i}}(z_{ij}^{\star};x_{ij})$. Such a representation
for $g_{\theta_{i}}$ is called a _deep equilibrium model_ [30], and the key
advantage lies in its compactness and expressive power: a single layer can
represent an infinitely deep weight-tied network while maintaining a constant
memory footprint [30]. This allows the model to learn a sophisticated non-
linear representation at a low memory capacity. We describe this in detail in
Sec. III-B below.

We then formulate a global consensus optimization problem for FeDEQ as
follows:

$\displaystyle\min_{\\{\theta_{i}\\},\\{w_{i}\\},\theta}\quad\sum_{i=1}^{n}\Bigg{[}{\mathcal{L}}_{i}\left(\theta_{i},w_{i}\right)$
$\displaystyle\vcentcolon=\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}\ell\left(f_{\theta_{i},w_{i}}\left(x_{ij}\right),y_{ij}\right)\Bigg{]}$
(2) $\displaystyle\text{s.t.}\quad\theta_{i}$
$\displaystyle=\theta,\quad\forall i.$

This mathematical representation embodies a critical balance between global
coherence and local customization. The consensus constraints
$\theta_{i}=\theta,\forall i$ are to enforce that the local $\theta_{i}$ are
learned through sharing among edge nodes. This encourages the representation
layer to benefit from the structure of all nodes’ data—the objective of FL. On
the other hand, all $w_{i}$ are fine-tuned locally, adapting from the shared
representation to achieve personalization. In Sec. IV below, we present a
novel algorithm to solve this optimization problem based on ADMM consensus
optimization.

###  III-B Equilibrium Layers: Forward Pass and Backward Pass

A single equilibrium layer in the constraint of Equation 1 embodies an
“infinitely deep” network through a fixed-point system. Here we examine a
common non-linear mapping of the form:

$z=g_{\theta}(z;x)\vcentcolon=\phi(Bz+Cx+b),$ (3)

where $x\in\mathbb{R}^{d}$ is the input, $z\in\mathbb{R}^{d_{1}}$ is the
output, $\theta=(B\in\mathbb{R}^{d_{1}\times
d_{1}},C\in\mathbb{R}^{d_{1}\times d},b\in\mathbb{R}^{d_{1}})$ is the
parameters and $\phi:\mathbb{R}^{d_{1}}\rightarrow\mathbb{R}^{d_{1}}$ is some
nonlinear activation functions. Equation Equation 3 implies that $z^{\star}$
is a fixed point of the function $z\mapsto\phi(Bz+Cx+b)$ for a given data
point (or a mini-batch) $x$: an equilibrium layer of the form Equation 3 is
equivalent to feeding the input through an infinite sequence of weight-tied
layers of a feedforward network. It has been observed that such a single
layer, while much smaller in size, can perform on par with deep networks
containing many explicit layers. This makes implicit representations
particularly appealing for FL since the communication of model parameters is
often a bottleneck.

####  III-B1 Forward Pass – Fixed-Point Representation

In the general problem Equation 1 or the special case Equation 3, we aim to
find an equilibrium state $z^{\star}$ such that
$z^{\star}=g_{\theta}(z^{\star};x)$, where $g_{\theta}$ is a neural network
parameterized by $\theta$ and $x$ is the input to the model. The fixed-point
formulation allows a single equilibrium layer to capture long-range
dependencies as all parts of the input contribute to the equilibrium state
$z^{\star}$. Moreover, in the process of finding $z^{\star}$, the model
implicitly emphasizes features that are stable under $g_{\theta}$, promoting
robust and consistent representations. Additionally, if the input distribution
changes, the DEQ can adapt by finding a new equilibrium, enabling dynamic
adaptation to data changes. However, there remain two subproblems that affect
the forward pass of DEQ optimization including fixed-point convergence and
fixed-point solver.

Fixed-point convergence. From a theoretical perspective, a unique fixed point
$z^{\star}$ offers an equilibrium representation under $g_{\theta}$. The
convergence of fixed-point iteration is strongly tied to the properties of the
function $g_{\theta}$, where a crucial condition for the guaranteed
convergence is that $g_{\theta}$ needs to be a contraction mapping. This means
that there exists a Lipschitz constant $L<1$ such that
$\|g_{\theta}(x)-g_{\theta}(y)\|\leq L\|x-y\|$ for all $x,y$ in the domain of
$g_{\theta}$. If $g_{\theta}$ is a contraction mapping, the Banach fixed-point
theorem guarantees the convergence to a unique fixed point [39]. However,
ensuring that $g_{\theta}$ is a contraction mapping can be challenging in
practice, given that $g_{\theta}$ is typically a deep neural network, capable
of representing highly complex and non-linear functions.

Therefore, we make two assumptions following the theory in [31]. First, $\phi$
is a component-wise non-expansive (CONE) mapping, i.e., (i) the $k$th
component of its output only depends on the $k$th component of its input and
(ii) when operating on scalar input, $\phi$ is $1$-Lipschitz continuous:
$\forall u,v\in\mathbb{R},\lvert\phi(u)-\phi(v)\rvert\leq\lvert u-v\rvert$.
Many activation functions in deep learning, such as tanh, sigmoid, and (leaky)
ReLU, satisfy this property. Second, the infinity norm of $B$, denoted by
$\lVert B\rVert_{\infty}$, is strictly bounded above by $1$. Together, these
assumptions imply that $g_{\theta}$ in Equation 3 is a contraction mapping,
and a fixed point $z^{\star}$ exists for all input $x$. The use of
$l_{\infty}$ projection can promote fixed-point convergence in certain
iterative methods by limiting the size of the updates. Specifically, after
each update, the resulting point is ”projected” onto a set defined by an
$\infty$-norm ball. For a matrix $B\in\mathbb{R}^{d_{1}\times d_{1}}$ in
Equation 3 and radius $\kappa>0$ we can employ $l_{\infty}$ projection by
solving the following problem.

$\min_{A}\frac{1}{2}\lVert A-B\rVert_{F}^{2}\leavevmode\nobreak\
\leavevmode\nobreak\ \text{s.t.}\leavevmode\nobreak\ \leavevmode\nobreak\
\lVert A\rVert_{\infty}\leq\kappa.$ (4)

This problem can be solved by decomposing across rows and applying existing
methods such as Bisection [31]. We demonstrate these findings to an FL context
that involves 100 edge nodes, each possessing a portion of the FEMNIST dataset
for a 62-class classification task. To benchmark, we employ the widely used
FedAvg algorithm and compare two local model configurations:

Figure 3: FedAvg using DEQ-MLP with and without $l_{\infty}$ projection

(i) a DEQ-MLP network with a single equilibrium layer of size 512 (DEQ-MLP),
wherein the parameters associated with $z$ are projected onto the $l_{\infty}$
norm, succeeded by a linear layer of size 128, and ultimately, a final layer
mapping to 62 classes; and (ii) a similar network architecture but without the
projection step. As illustrated in Fig. 3, the DEQ-MLP network with the
projection mechanism achieves convergence, whereas the network without
projection appears to be diverging.

Fixed-point Solver. Recent works leverage Broyden [40, 30] and Anderson
Acceleration (AA)[41, 33, 42] for solving fixed-point systems in DEQs. In this
work, we employ a variant of AA to expedite the convergence of fixed-point
iterations. The concept behind AA is to form the next iterate as a linear
combination of past iterates, rather than merely relying on the last iterate.
Specifically, AA minimizes the differences between the iterates and their
mapped values over a history of $m$ previous iterations. In terms of
representation learning in FeDEQ, we suggest a warm-start strategy for AA in
which the fixed point of the previous iterations is reused as the initial
point for the next. Since the changes in $\theta$ and $x$ are small from one
iteration to the next (i.e., $\theta^{(i)}\approx\theta^{(i-1)}$ and
$x^{(i)}\approx x^{(i-1)}$), the fixed points $z^{{}^{(i)}}$ and
$z^{{}^{(i-1)}}$ are likely to be close to each other as well. This could
bring several benefits:

  * •

Faster Convergence: As the input data changes slightly from one iteration to
the next, reusing the fixed point could lead to faster convergence of the
fixed-point iteration, which could in turn speed up the overall training
process.

  * •

Continuity of Representations: The learned representations could exhibit a
form of continuity or smoothness from one iteration to the next. This could
potentially make the learned representations more robust and stable,
especially in the face of small changes in the input data or the underlying
function.

  * •

Improved Robustness: If the model is trained on a stream of data that changes
over time, reusing the fixed point could make the model more robust to these
changes, as it provides a form of context from the previous data.

####  III-B2 Backward Pass – Implicit Differentiation

Once the equilibrium state $z^{\star}$ has been found in the forward pass, the
model computes the loss ${\mathcal{L}}$ based on this state and the target
output. The goal in the backward pass is to update the parameters $\theta$ to
minimize this loss. The gradient of the loss ${\mathcal{L}}$ w.r.t. $\theta$
is calculated by employing the chain-rule as follows:

$\frac{\partial{\mathcal{L}}}{\partial\theta}=\frac{\partial{\mathcal{L}}}{\partial
z^{\star}}\frac{\partial z^{\star}(\theta)}{\partial\theta}$ (5)

In this context, the gradient of ${\mathcal{L}}$ w.r.t. $z^{\star}$ can be
derived through backpropagation. The notation $z^{\star}(\theta)$ is also
utilized to signify situations where $z^{\star}$ is perceived as an implicit
function of the variables that we’re differentiating w.r.t. the shared
parameter $\theta$, and we use $z^{\star}$ on its own when we’re just
mentioning the equilibrium value. Because of the self-consistent nature of the
equilibrium state $z^{\star}$, the gradient $\partial
z^{\star}(\theta)/{\partial\theta}$ cannot be computed directly. Instead,
implicit differentiation is used. By differentiating the fixed-point equation
$z^{\star}=g_{\theta}(z^{\star};x)$ w.r.t. $\theta$, we get:

$\frac{\partial z^{\star}(\theta)}{\partial\theta}=\frac{\partial
g_{\theta}(z^{\star};x)}{\partial\theta}=\frac{\partial
g_{\theta}(z^{\star};x)}{\partial z^{\star}}\frac{\partial
z^{\star}(\theta)}{\partial\theta}+\frac{\partial
g_{\theta}(z^{\star},x)}{\partial\theta}$ (6)

Upon reorganizing the terms, we attain an explicit formulation for the
Jacobian as follows.

$\frac{\partial z^{\star}(\theta)}{\partial\theta}=\Bigg{(}I-\frac{\partial
g_{\theta}(z^{\star};x)}{\partial z^{\star}}\Bigg{)}^{-1}\frac{\partial
g_{\theta}(z^{\star};x)}{\partial\theta}$ (7)

where $I$ is the identity matrix, and $\frac{\partial g_{\theta}}{\partial
z^{\star}}$ and $\frac{\partial g_{\theta}}{\partial\theta}$ are the Jacobians
of $g_{\theta}$ w.r.t. $z^{\star}$ and $\theta$, respectively. Although the
right-hand side terms can be explicitly computed using conventional automatic
differentiation, solving the inverse of Jacobian is computationally intensive,
especially when $z^{*}$ is high-dimensional. To solve this equation more
efficiently, we can employ the following vector-Jacobian product (VJP) [43].

$y^{T}\left(\frac{\partial
z^{\star}(\theta)}{\partial\theta}\right)=y^{T}\left(I-\frac{\partial
g(z^{\star};x)}{\partial z^{\star}}\right)^{-1}\left(\frac{\partial
g(z^{\star};x)}{\partial\theta}\right).$ (8)

which is equivalent to the following linear system:

$u^{T}=u^{T}\left(\frac{\partial g_{\theta}(z^{\star};x)}{\partial
z^{\star}}\right)+y^{T}.$ (9)

where $u^{T}=y^{T}\left(I-\frac{\partial g_{\theta}(z^{\star},x)}{\partial
z^{\star}}\right)^{-1}$ is the solution of the linear system, which can be
solved efficiently by using linea solvers such as conjugate gradient descent
[44]. To this end, by substituting Equation 7 to Equation 5, we can have a
form of VJP as follows.

$\frac{\partial{\mathcal{L}}}{\partial\theta}=\frac{\partial{\mathcal{L}}}{\partial
z^{\star}}\left(I-\frac{\partial g(z^{\star};x)}{\partial
z^{\star}}\right)^{-1}\frac{\partial g(z^{\star};x)}{\partial\theta}$ (10)

This can be solved using the above implicit differentiation techniques. The
gradient ${\partial{\mathcal{L}}}$ is then used in gradient-based optimization
methods to update the model parameters $\theta$ and minimize the loss.

One of the significant advantages of DEQs’ backward pass is their memory
efficiency. In traditional deep learning models, the activations of all layers
need to be stored for backpropagation. In contrast, DEQs only need to store
the final equilibrium state, as all layers are the same and share the same
parameters. This results in substantial memory savings, especially for very
deep networks. Furthermore, DEQs allow for end-to-end training with gradient-
based optimization methods [30]. Despite the complexity introduced by the
fixed-point iteration and implicit differentiation, the backward pass is still
differentiable, and gradients can be computed and used for parameter updates.

Jacobian-Free Backpropagation. The main challenge in the backward pass of DEQs
is the extensive computational time, even though existing methods can avoid
computing the inversion of a Jacobian matrix directly. Surprisingly, recent
research has demonstrated implicit differentiation can be performed by using
the zeroth-order approximation of the Neumann series for the inverse of the
Jacobian matrix [45] instead of explicitly computing of this expensive term.
This approach still provides satisfactory performance with significantly
reducing computation time. Hence, in this work, we employ this Jacobian-Free
backpropagation approach to estimate the implicit gradients.

Figure 4: FedAvg with 5-layer MLP model vs. DEQ-MLP

###  III-C Federated Learning with DEQs vs. Explicit Models

Previous work has shown that equilibrium layers can achieve similar
performance to explicit layers while maintaining a compact model size [30]. We
extend this finding to an FL context with 100 edge nodes, each holding a
portion of the FEMNIST dataset for a 62-class classification task. We employ
the de-facto FedAvg to compare two local model settings: (i) an MLP network
with 4 explicit layers of size 512, followed by a linear layer of size 128,
and a final layer mapping to 62 classes; and (ii) the same network with the
four explicit layers replaced by a single equilibrium layer of size 512 (DEQ-
MLP) as in Equation 3. Fig. 4 shows that the same test accuracy is achieved in
both settings. However, since setting (ii) requires a model that is $41\%$
smaller in size, its communication cost is saved by nearly half per round. In
Sec. V, we provide experiments with much more sophisticated models like DEQ-
ResNet and DEQ-Transformer that also highlight this benefit.

##  IV Consensus Optimization for Federated Learning with Equilibrium Layers

With the objective Equation 2, we introduce an ADMM consensus optimization
variant to collaboratively learn the shared parameters $\theta_{i}$ and the
personalized parameters $w_{i}$. We detail the optimization procedure in
Algorithm 1 and provide some remarks on its convergence at the end of this
section.

###  IV-A Learning shared representation via ADMM Consensus Optimization

In FL for edge networks, the data heterogeneity across diverse edge nodes
often culminates in a pronounced challenge known as client drift [18]. Such
drift emerges when edge nodes, optimizing their model based on uniquely
distributed data, converge toward a local objective minimum. This local
convergence may differ considerably from the global optimization objective
desired for an aggregated model. Consequently, when the server integrates
models from myriad nodes, it often encounters a compounded effect of these
varied local minima. The resulting global model, unfortunately, can become an
aggregate that fails to optimally represent any individual edge node’s data,
leading to overall suboptimal performance.

Motivation for using ADMM in Combatting Client Drift Effects. To mitigate the
challenges posed by data heterogeneity and client drift, we have formulated
the consensus optimization problem Equation 2 to find a shared representation
across all edge nodes and propose to employ ADMM [34] to solve it. This
algorithm is based on alternating updates of primal and dual variables to
achieve two goals: $\theta_{i}$ across edge nodes become closer and the
objective function decreases. Applying this to our problem, we first introduce
a dual variable $\lambda_{i}$ for each $\theta_{i}$ and construct the
following augmented Lagrangian for each client $i$:

$\tilde{{\mathcal{L}}_{i}}(\theta,\theta_{i},w_{i},\lambda_{i})\vcentcolon={\mathcal{L}}_{i}(\theta_{i},w_{i})+\left\langle\lambda_{i},\theta_{i}-\theta\right\rangle+\frac{\rho}{2}\lVert\theta_{i}-\theta\rVert^{2}$
(11)

where $\rho>0$ is the penalty parameter.

Upon solving the local augmented Lagrangian, clients will then collaboratively
minimize the global objective $\widetilde{{\mathcal{L}}}$, which is defined as
follows.

$\widetilde{{\mathcal{L}}}(\theta,\\{\theta_{i}\\},\\{w_{i}\\},\\{\lambda_{i}\\})\vcentcolon=\sum_{i=1}^{n}\left[\tilde{{\mathcal{L}}_{i}}(\theta,\theta_{i},w_{i},\lambda_{i})\right],$
(12)

Delving deeper into the augmented Lagrangian Equation 11, two essential terms
come to the forefront. The inner product term,
$\left\langle\lambda_{i},\theta_{i}-\theta\right\rangle$, reflects the
alignment between the deviation of a local model from the global model and its
associated dual variable, $\lambda_{i}$. Essentially, it acts as a controller
of discrepancies between local and global models. The quadratic term,
$\frac{\rho}{2}\lVert\theta_{i}-\theta\rVert^{2}$, directly penalizes
deviations of the local model from the global standard. The parameter $\rho$
adjusts the stringency of this penalty, ensuring a balance between local
adaptability and global consensus. These terms result in the new correction
term in the gradient
$\nabla_{\theta_{i}}\tilde{{\mathcal{L}}_{i}}=\nabla_{\theta_{i}}{\mathcal{L}}_{i}+\lambda_{i}+\rho(\theta_{i}-\theta)$
aligning client updates with server directions to overcome client drifts.

As described in detail in Algorithm 1, the ADMM consensus optimization
proceeds by iteratively updating the local parameters $\theta_{i}$ and
$w_{i}$, the dual variables $\lambda_{i}$, and the shared global parameters
$\theta$. This procedure consists of three key steps per communication round.
First, the server performs averaging to update the shared parameters $\theta$
based on the set of $\theta_{i}$ obtained from the previous round (line 5) and
then sends it to a new subset of edge nodes (line 7). Second, each selected
node $i$ minimizes the augmented Lagrangian $\tilde{{\mathcal{L}}}_{i}$ w.r.t.
the primal variables (lines 8–9). This local minimization makes use of the
Rep_Update procedure (lines 13–20), which involves solving Problems Equation 1
and Equation 10 to obtain the gradient
$\nabla_{\theta_{i}}\tilde{{\mathcal{L}}_{i}}$ via implicit differentiation.
The dual variable $\lambda_{i}$ is then updated by an ascent step to enforce
the consensus constraint (lines 10–11). This process is repeated within a
number of iterations until convergence, at which point the local
representation $\theta_{i}$ and the consensus variable $\theta$ agree,
resulting in a shared representation learned from the distributed data.

###  IV-B Personalization through Explicit Layers

With the shared representation $\theta$ learned via ADMM consensus
optimization, we proceed to the second aim of FeDEQ, which focuses on
achieving personalized model for each node to better address data
heterogeneity. Particularly, in every communication round, node $i\in S$ fine-
tunes its personalized parameters $w_{i}$ on top of the shared representation
$\theta$ with its local data, i.e., minimizing the original loss
${\mathcal{L}}_{i}(\theta,w_{i})$ w.r.t. $w_{i}$ (Line 8 of Algorithm 1). This
problem can be efficiently solved by using traditional gradient-based
optimization methods such as SGD, where the personalized parameter $w_{i}^{t}$
obtained from the previous round $t$ is used as the initial value for updating
the next $w_{i}^{t+1}$, providing edge nodes the flexibility adapt to their
unique data distributions to different personalization settings.

By decomposing the global problem into local subproblems and a consensus step,
ADMM consensus procedure seeks a balance between local objectives and the
global objective. The augmented Lagrangian Equation 11 introduces a penalty
for deviations from the consensus. Thus, even if a node’s local data pushes
its model towards a distinct local minimum, the consensus penalty nudges it
closer to the global consensus. Furthermore, with the penalty parameter
$\rho$, ADMM can be adjusted in terms of how strictly it enforces consensus.
This provides flexibility to tune the degree of consensus based on the extent
of client drift observed.

Furthermore, combining equilibrium models with the ADMM consensus optimization
framework enhances communication efficiency in key ways. Equilibrium models,
with their implicitly ”infinite-depth” single layer, simplify model complexity
and reduce communication overhead by requiring fewer parameters. Moreover,
they maintain a constant, often lower, memory footprint than traditional deep
models, thus enabling resource-constrained edge devices to manage complex
models without straining hardware resources. Additionally, the ADMM framework
permits asynchronous updates of local parameters, thereby optimizing network
resource utilization and convergence speed.

Algorithm 1 FeDEQ

1:Parameters: Communication rounds $T$; penalty $\rho$

2:Initialize $S_{0}=\\{1,\ldots,n\\}$, and $\theta^{0}=\theta_{i}^{0}$,
$w_{i}^{0}$ and $\lambda_{i}^{0}$ randomly for $i=1,2,\ldots,n$

3:for $t=0,1,\ldots$, $T-1$ do $\triangleright$ Global iteration

4: Server aggregates the shared parameters:

5: $\theta^{t+1}=\frac{1}{\lvert S_{t}\rvert}\sum_{i\in S_{t}}\theta_{i}^{t}$

6: Server randomly samples a subset $S_{t+1}$ of edge nodes

7: for each edge node $i\in S_{t+1}$ in parallel do

8: Receive $\theta^{t+1}$ from server $\triangleright$ Communication

9: Update primal variables:

10: $\left\\{\rule{0.0pt}{14.39996pt}\right.$
$w_{i}^{t+1}\leftarrow\text{arg\,min}_{w_{i}}{\mathcal{L}}_{i}(\theta^{t+1},w_{i})$

11:
$\theta_{i}^{t+1}\leftarrow\textsc{Rep\\_Update}(\tilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t},w_{i}^{t+1},\lambda_{i}^{t}))$

12: Update dual variable:
$\lambda_{i}^{t+1}=\lambda_{i}^{t}+\rho\left(\theta_{i}^{t+1}-\theta^{t+1}\right)$

13: Send $\theta_{i}^{t+1}$ to server $\triangleright$ Communication

14:

15:

16:

17:function
Rep_Update($\tilde{{\mathcal{L}}_{i}}(\theta,\theta_{i},w_{i},\lambda_{i}$))

18: Parameters: local dataset (${\mathcal{X}}_{i},{\mathcal{Y}}_{i}$),
learning rate $\eta$

19: for each local epoch do

20: for each sample $(x,y)\in({\mathcal{X}}_{i},{\mathcal{Y}}_{i})$ do

21: Find the fixed-point $z^{\star}=g_{\theta_{i}}(z^{\star};x)$

22: Compute the loss
${\mathcal{L}}_{i}(h_{w_{i}}(g_{\theta_{i}}(z^{\star};x)),y)$

23: Compute the gradient $\nabla_{\theta_{i}}{\mathcal{L}}_{i}$ using Eq.
Equation 10

24: Update gradients
$\nabla_{\theta_{i}}\tilde{{\mathcal{L}}_{i}}=\nabla_{\theta_{i}}{\mathcal{L}}_{i}+\lambda_{i}+\rho(\theta_{i}-\theta)$

25: Update:
$\theta_{i}\leftarrow\theta_{i}-\eta\nabla_{\theta_{i}}\tilde{{\mathcal{L}}_{i}}$
return $\theta_{i}$

###  IV-C FeDEQ: Convergence Analysis

We present a theoretical analysis for FeDEQ’s convergence in a smooth, non-
convex setting. In comparison with the original consensus optimization problem
in [34, §7.1], our problem is different in four key aspects.

First, on edge node $i$ the loss function ${\mathcal{L}}_{i}$ is non-convex
w.r.t. $(\theta_{i},w_{i})$. Therefore, Algorithm 1 is based on a more recent
treatment of ADMM for non-convex objectives in [46, Algorithm 3]. In
particular, we assume that ${\mathcal{L}}_{i}$ is differentiable and
$L$-Lipschitz smooth w.r.t. $(\theta_{i},w_{i})$. Based on this, the update of
these variables is done using gradient descent steps, which allows us to bound
the decrease in $\widetilde{{\mathcal{L}}}_{i}$ after lines 8–9. Further, the
gradient w.r.t. $w_{i}$ can be found via the traditional backward propagation
procedure, as $w_{i}$ is the parameter of an explicit module. On the other
hand, the gradient w.r.t. $\theta_{i}$, the parameters of an implicit module,
is found using the implicit function theorem, which we have previously
detailed in Sec. III-B.

Second, as is typical of FL algorithms, we must allow for node sampling, i.e.,
only a subset of edge nodes is required to participate in training in each
global iteration $t$. This is done via choosing the subset $S_{t+1}$ (line 5)
and only aggregating the parameters $\theta_{i}$ from the sampled nodes (line
4). To assist the convergence analysis, we make an assumption that, regardless
of the sampling method, there exists a period of $T>0$ global rounds after
which _all edge nodes_ will have participated at least once. In [46], this is
called the period-$T$ essentially cyclic update rule.

Third, we employ _variable splitting_ for our ADMM variant. In other words,
the primal variable is split into $\theta_{i}$ and $w_{i}$ on each node, and
only $\theta_{i}$ is subject to the consensus constraint in Equation 2.
Moreover, unlike the original ADMM version which minimizes
$\widetilde{{\mathcal{L}}}_{i}$ jointly w.r.t. $\theta_{i}$ and $w_{i}$, we
split the update of these variables in Algorithm 1. For $w_{i}$ (line 8),
since it is not constrained, minimizing the augmented Lagrangian is equivalent
to minimizing the loss function ${\mathcal{L}}_{i}$. Then, $\theta_{i}$ is
updated by minimizing $\widetilde{{\mathcal{L}}}_{i}$.

Fourth and finally, as is common in large-scale machine learning, parameters
are learned using optimization with _mini-batches_. This gives us stochastic
gradients w.r.t. the parameters $\theta_{i}$ and $w_{i}$, and as a result, the
augmented Lagrangian is not guaranteed to decrease after every SGD round (it
is only guaranteed in expectation).

We make the following assumption for our convergence analysis.

######  Assumption IV.1.

For every node $i$, the local loss function ${\mathcal{L}}_{i}$ is bounded
below. Further, ${\mathcal{L}}_{i}$ is differentiable everywhere w.r.t.
$(\theta_{i},w_{i})$ and has Lipchitz gradients with parameter $L>0$.

######  Remark IV.2.

This assumption is common in deep learning’s non-convex optimization analysis,
ensuring gradient stability and the existence of an optimal solution.

To establish a convergence guarantee in a non-convex setting of ADMM, we aim
to show the following. Start with the variables
$(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})$.
After $T$ rounds (i.e., after all edge nodes have participated at least once
since iteration $t$), the new variables
$(\theta^{t+T},\\{\theta_{i}^{t+T}\\},\\{w_{i}^{t+T}\\},\\{\lambda_{i}^{t+T}\\})$
lead to decrease in the global augmented Lagrangian
$\widetilde{{\mathcal{L}}}$. Then we show these variables will converge to a
local stationary point of $\widetilde{{\mathcal{L}}}$. The following result
establishes the convergence of FeDEQ.

######  Lemma IV.3 (Decrease in augmented Lagrangian after every $T$
iterations).

Suppose $\rho$ is chosen large enough that $\rho>6L$. After $T$ global
iterations since $t$, when all edge nodes have participated in training at
least once, we have

$\displaystyle\begin{split}&\widetilde{{\mathcal{L}}}(\theta^{t+T},\\{\theta_{i}^{t+T}\\},\\{w_{i}^{t+T}\\},\\{\lambda_{i}^{t+T}\\})-\widetilde{{\mathcal{L}}}(\theta_{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})\\\
&\leq\frac{1}{T}\sum_{i=1}^{n}\left[-\left(\frac{\rho}{2}-\frac{3L^{2}}{\rho}\right)\lVert\theta^{t+T}-\theta^{t}\rVert^{2}\right.\\\
&\leavevmode\nobreak\
\qquad\qquad\qquad\quad\left.-\left(\frac{L+\rho}{2}-\frac{12L^{2}}{\rho}\right)\lVert\theta_{i}^{t+T}-\theta_{i}^{t}\rVert^{2}\right.\\\
&\leavevmode\nobreak\
\qquad\qquad\qquad\qquad\quad\left.-\left(\frac{L}{2}-\frac{3L^{2}}{\rho}\right)\lVert
w_{i}^{t+T}-w_{i}^{t}\rVert^{2}\right].\end{split}$ (13)

###### Proof.

The detailed proof can be found in Sec. A-A. ∎

######  Remark IV.4.

Lemma IV.3 shows that if we choose a large enough $\rho$, the augmented
Lagrangian decreases by a sufficient amount after all edge nodes have
participated at least once. However, the smoothness parameter is difficult to
estimate in practice. We therefore fine-tune $\rho$ from a predefined range of
values.

######  Lemma IV.5 (Lower bound for augmented Lagrangian).

Suppose that the global loss function ${\mathcal{L}}(\theta,\\{w_{i}\\})$ is
lower bounded, i.e.,

$\underline{{\mathcal{L}}}\vcentcolon=\min_{\theta,\\{w_{i}\\}}{\mathcal{L}}(\theta,\\{w_{i}\\})>-\infty$

If $\rho$ is chosen large enough that $\rho>6L$, then the augmented Lagrangian
is lower bounded by $\underline{{\mathcal{L}}}$.

###### Proof.

The detailed proof can be found in Sec. A-B. ∎

Now we combine Lemmas IV.3 and IV.5 to establish the following theorems for
the convergence of Equation 12.

######  Theorem IV.6 (Convergence of the augmented Lagrangian).

Suppose the following are true:

  * •

After $T$ iterations, all clients have participated in training at least once.

  * •

The global loss function ${\mathcal{L}}(\theta,\\{w_{i}\\})$ is bounded below
by a finite quantity $\underline{{\mathcal{L}}}$.

  * •

The hyperparameter $\rho$ for the augmented Lagrangian is chosen such that
$\rho>6L$.

Then the augmented Lagrangian $\widetilde{{\mathcal{L}}}$ will monotonically
decrease and is convergent to a quantity of at least
$\underline{{\mathcal{L}}}$. Further, for all $i=1,\dots,n$ we have
$\lim_{t\rightarrow\infty}\lVert\theta^{t+T}-\theta_{i}^{t+T}\rVert=0$.

###### Proof.

Lemma IV.3 implies that the augmented Lagrangian decreases by a non-negative
amount after every $T$ global iterations. Given Lemma IV.5 and the fact that
every client will be updated at least once in the interval $[t,t+T]$, we
conclude that the limit
$\lim_{t\rightarrow\infty}\widetilde{{\mathcal{L}}}(\theta^{t+T},\\{\theta_{i}^{t+T}\\},\\{w_{i}^{t+T}\\},\\{\lambda_{i}^{t+T}\\})$
exists and is at least $\underline{{\mathcal{L}}}$.

Now we prove the second statement. From Equation 27 and the fact that
$\widetilde{{\mathcal{L}}}$ converges, we conclude that as
$t\rightarrow\infty$,

$\displaystyle\lVert\theta^{t+T}-\theta^{t}\rVert\rightarrow
0,\quad\lVert\theta_{i}^{t+T}-\theta_{i}^{t}\rVert\rightarrow 0,\quad\lVert
w_{i}^{t+T}-w_{i}^{t}\rVert\rightarrow 0.$

Combining this with Equation 25, we have
$\lVert\lambda_{i}^{t+1}-\lambda_{i}^{t}\rVert\rightarrow 0$. Based on the
definition of $\lambda_{i}^{t+1}$, this implies that
$\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert\rightarrow 0$. ∎

######  Remark IV.7.

Theorem IV.6 establishes that the sequence of primal and dual variables
updated after each $T$ global iterations of Algorithm 1 converges.
Furthermore, we have $\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert\rightarrow 0$,
i.e., the consensus constraint is satisfied.

In the below theorem we present another guarantee that the limit of the
sequence
$(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})$ is
the stationary solution for Problem Equation 2.

######  Theorem IV.8 (Convergence to stationary point).

Suppose the assumptions in Theorem IV.6 hold. Then the limit point
$(\theta^{*},\\{\theta_{i}^{*}\\},\\{w_{i}^{*}\\},\\{\lambda_{i}^{*}\\})$ of
the sequence
$(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})$ is a
stationary solution to Problem Equation 2. That is, for all $i$,

$\displaystyle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta^{*},w_{i}^{*})+\lambda_{i}^{*}$
$\displaystyle=0,$
$\displaystyle\nabla_{w_{i}}{\mathcal{L}}_{i}(\theta^{*},w_{i}^{*})$
$\displaystyle=0,$ $\displaystyle\theta^{*}$ $\displaystyle=\theta_{i}^{*}.$

###### Proof.

See the proof to [46, Theorem 2.4]. ∎

Here we make some remarks. First, because Algorithm 1 allows for node sampling
on line 5, the proof relies on an assumption that all edge nodes will
participate after some fixed number of rounds (in particular, $T$ rounds).
Second, since local objectives are non-convex and smooth, the primal variables
$\theta_{i}$ and $w_{i}$ are updated using first-order methods like SGD.
Further, the gradient w.r.t. $\theta_{i}$ is found through IFT (lines 12–13),
while that w.r.t. $w_{i}$ only involves a simple backward pass. Finally, the
proof aims to show that after every $T$ rounds, we obtain a sufficient
decrease in the augmented Lagrangian; coupled with the fact that the augmented
Lagrangian is bounded below, the sequence of variables converges to a
stationary point of $\widetilde{{\mathcal{L}}}$.

##  V Experimental Results

In this section, we present the experimental results and insights into the
performance and effectiveness of FeDEQ, demonstrating its advantages compared
to existing approaches under various scenarios.

###  V-A Experimental Setup

TABLE I: Details on Datasets and Non-i.i.d. Partition Methods

Dataset |  Partition Method |  Original Nodes |  Sample Nodes |  Labels/ Node |  Total Train |  Total Test |  Mean |  Std
---|---|---|---|---|---|---|---|---
FEMNIST |  Natural non-i.i.d. |  3400 |  100
200 |  Various |  19,357
39,139 |  2,237
4,521 |  194
195 |  72.9
75.7
CIFAR-10 |  By-label non-i.i.d. |  100 |  100 |  2 & 5 |  50,000 |  10,000 |  500 |  –
CIFAR-100 |  By-label non-i.i.d. |  100 |  100 |  5 & 20 |  50,000 |  10,000 |  500 |  –
SHAKESPEARE |  Natural non-i.i.d. |  715 |  200 |  Various |  92,834 |  11,805 |  465 |  506.2

TABLE II: Summary on Shared Representation and Personalization Layers

Datasets | Methods |  Shared Representation |  Params (M)/ Size (MB) |  Personalized Layer |  Params (M)/ Size (MB)
---|---|---|---|---|---
FEMNIST | Baselines |  ResNet-20 (19 Conv + 2 FC) |  2.97 / 11.87 |  1 FC |  0.008 / 0.03
| FeDEQ |  DEQ-ResNet-S (3 Conv + 2 FC) |  0.98 / 3.94 |  1 FC |  0.008 / 0.03
CIFAR-10 | Baselines |  Resnet-34 (33 Conv + 2 FC) |  7.78 / 31.11 |  1 FC |  0.001/0.01
| FeDEQ |  DEQ-Resnet-M (3 Conv + 2 FC) |  2.73 / 10.90 |  1 FC |  0.001/0.01
CIFAR-100 | Baselines |  Resnet-34 (33 Conv + 2 FC) |  10.88 / 43.51 |  1 FC |  0.01/0.05
| FeDEQ |  DEQ-Resnet-M (3 Conv + 2 FC) |  2.73 / 10.90 |  1 FC |  0.01/0.05
ShakeSpeare | Baselines |  Transformer-8 (8 UT, 4 H) |  0.35 / 1.43 |  1 UT + 1 FC |  0.06/0.22
| FeDEQ |  DEQ-Transformer (4 UT, 4 H) |  0.21 / 0.83 |  1 UT + 1 FC |  0.06/0.22

  * Acronym: Conv – Convolutional layer; FC – Fully Connected layer; UT – Universal Transformer layer; H – Head-attention; Params – The number of model parameters; M – Millions; Size – The size of model; MB – Megabytes

####  V-A1 Datasets and Non-i.i.d. Partition Methods.

We evaluate FeDEQ on four diverse federated datasets: FEMNIST [47], CIFAR-10
[48], CIFAR-100 [48], and SHAKESPEARE [49]. These datasets cater to various
application domains, such as image classification and natural language
processing, and are suitable for evaluating FL algorithms in different
settings for edge networks. Table I presents the details for the datasets.
Here we list some key characteristics of each dataset.

FEMNIST: A variant of the handwritten EMNIST dataset [50] designed
specifically for FL. In this work, we use the version obtained from FedJax’s
database [51], which simulates real-world non-i.i.d data distributions by
writers. The original dataset contains 3,400 writers, providing a total of
780,000 images, with each writer contributing at least one image of each of 62
labels (10 digits and 52 letters). Due to a lack of computing resources, we
randomly sample 100 and 200 writers from the original to assign to edge nodes
for our experiments.

CIFAR-10: An object recognition dataset [48] including $50,000$ training and
$10,000$ test colored images belonged to $10$ classes. The dataset is equally
partitioned into 100 nodes in a non-i.i.d. manner, with each node either
having data from 2 or 5 distinct classes, depending on the desired scenario.
This partitioning simulates real-world FL scenarios facilitating the
evaluation of algorithms under heterogeneous data distributions.

CIFAR-100: An object recognition dataset [48] including $50,000$ training and
$10,000$ test colored images belonged to $100$ classes. The dataset is divided
equally into 100 edge nodes non-i.i.d., with each node having data from either
5 or 20 distinct classes, based on the chosen scenario. This partitioning
strategy provides a realistic FL environment, allowing for comprehensive
algorithm evaluation with diverse data distributions.

SHAKESPARE: A federated dataset obtained from FedJax’s database [49]
comprising 715 users, representing characters from Shakespeare’s plays, with
each example being a series of consecutive lines spoken by a character in a
specific play. In this paper, we randomly select 200 edge nodes from the
original dataset which each has more than 100 training samples with a sequence
length of 16.

####  V-A2 Baselines.

We benchmark FeDEQ against several comparison methods, including locally
trained models (Local), the de-facto FedAvg [1], the fine-tuning approach
(FedAvg+FT)[20, 21], and four SOTA personalized algorithms including: (1)
Ditto [9], an effective full model personalization approach via MTL; (2)
FedPer [24] and (3) FedRep [10], emerging partial personalization approaches
enabling nodes to collaboratively learn an explicit shared representation and
subsequently personalize the local model; and (4) kNN-Per [28], a semi-
parametric learning approach that employs kNN as a personalized model on top
of pre-trained models. Notably, the latter three are closely related to ours.
We apply settings consistent with FeDEQ for all methods to ensure a fair
comparison, and fine-tune their respective hyperparameters to reach their best
performance on the settings.

####  V-A3 Models and Representations.

For vision tasks, we construct DEQ-ResNet models inspired by ResNet [52],
comprising an equilibrium residual block followed by a 3-layer fully connected
(FC) network with the last serving for personalization and using Softplus
activations [53], a smooth version of ReLU, to enforce the smoothness
assumption in theoretical analysis. We offer two versions of our DEQ-ResNet:
”S” (small) and ”M” (medium), varying by the size of equilibrium layers. DEQ-
ResNet-S is used for FEMNIST and DEQ-ResNet-M for CIFAR-10/CIFAR-100, with
FeDEQ. As for other baselines, we employ the architectures of ResNet, using
their residual blocks with the same FC network as DEQ-ResNet. Especially,
ResNet-20 and ResNet-34 are then employed for experiments on FEMNIST and
CIFAR-10/CIFAR-100, correspondingly. For sequence tasks, we design a DEQ-
Transformer based on Universal Transformer (UT)[54], consisting of an
equilibrium UT block succeeded by a personalized module with one last UT block
and one linear layer; while other methods use a UT with 8 layers, 4-head
attention (Transformer-8). Specific details regarding the shared parameters
and personalized parameters, including the number of model parameters and the
model sizes, are presented in Table II and Sec. B-A.

For training DEQs, we use Anderson Acceleration [41] with warm-start strategy
for solving fixed-point system in the forward pass, and Jacobian-free
backpropagation [45] for solving the implicit differentiation in the backward
pass. It is worth noting that when using non-smooth activations like ReLU,
FeDEQ still achieves convergence which reaches similar levels of performance
compared to those models using smooth activations.

####  V-A4 Training Details.

Training is proceeded for $T=150$ communication rounds in the experiments
involving FEMNIST, CIFAR-10 and CIFAR-100, whereas for Shakespeare, it is
conducted for $T=400$ rounds. For all methods and experiments, we uniformly
sample 10% of nodes without replacement for each round and employ SGD
optimizer with a constant learning rate. Unless specified in certain
instances, we use $5$ local epochs for training shared representation and $3$
for personalization for all methods. The final accuracy is determined by
averaging local accuracies from all nodes over the last $10$ rounds. More
details about the training processes are presented in Sec. B-B.

All experiments are conducted in a coding environment powered by an AMD Ryzen
3970X Processor [55] with 64 cores and 256GB of RAM. Additionally, fours
NVIDIA GeForce RTX 3090 GPUs [56] are employed to accelerate the training
process. For software, we rely on Python3 as our primary language,
supplemented by advanced machine learning and optimization frameworks: Jax
[57], Haiku [58], Optax [59], Jaxopt [60], and FedJax [51]. Furthermore, our
implementation is also inspired by some personalized FL benchmarking works
such as Motley [61].

Figure 5: Effects of $\rho$ on the convergence of the shared representation
$\theta$

###  V-B Results and Evaluation

####  V-B1 Effects of $\rho$ on the performance of FeDEQ.

The penalty parameter $\rho$ in FeDEQ is a tuning parameter that controls the
trade-off between the dual constraint satisfaction (i.e., how closely the
local representation matches the shared representation) and the minimization
of the objective function. Therefore, the selection of $\rho$ is crucial for
the performance of FeDEQ and it can vary significantly based on the specific
optimization problem. To investigate the impact of $\rho$ FeDEQ’s performance
and convergence, we vary $\rho$ values within
$\\{0.001,0.01,0.1,1.0,5.0,10.0\\}$ on the FEMNIST dataset.

The results in Fig. 5 show that FeDEQ achieves convergence with any $\rho$
within the range; however, if $\rho$ is too small, the constraint
$\theta_{i}=\theta$ is weakly enforced, meaning that the local variables have
more freedom to minimize their local objective function, which might cause the
consensus variable to deviate from the optimum. This might increase the number
of iterations needed for convergence (e.g., $\rho=0.001$) or may even prevent
convergence in certain cases. On the other hand, it can provide a better local
performance as each subproblem can be more optimized according to its local
data. Conversely, if $\rho$ is too large, the algorithm will strongly enforce
the consensus constraint, which can lead to a swift alignment of the local
variables to the consensus variable (e.g., $\rho=5.0$ or $\rho=10.0$).
However, this might cause suboptimal solutions because the strong enforcement
of the consensus might prevent local variables from reaching their local
optima. In other words, it imposes high ”pressure” to force local solutions to
agree, potentially at the expense of solution quality.

Through empirical evaluation, we found that $\rho=0.01$ provides the best
balance for our datasets. In practice, selecting a suitable $\rho$ often
involves trial and error or a validation process. It is also possible to use
an adaptive strategy, where $\rho$ is updated during the optimization process.

Figure 6: Test accuracy of FeDEQ across the datasets.

Figure 7: Average loss of FeDEQ across the datasets.

####  V-B2 The convergence of FeDEQ.

To empirically assess the convergence of FeDEQ across all the datasets, we
provide both the test accuracy and the training loss (using the function
${\mathcal{L}}$) over global communication rounds in Fig. 6,7. Those figures
reveal a consistent trend: across all datasets, FeDEQ’s loss decreased and
stabilized within a dataset-specific range of communication rounds, signaling
a stable and reliable convergence trajectory. Specially, for the FEMNIST
dataset, convergence was observed by the 50th communication round, achieving a
test accuracy of approximately 87%. In the case of CIFAR-10, despite its
inherent challenges, FeDEQ exhibited convergence by the 80th round, with test
accuracies reaching up to 90%. The CIFAR-100 dataset, with its augmented label
complexity, presented a more intricate challenge. Nevertheless, convergence
was achieved around the same 80th round, with test accuracies ranging between
58% and 80%, contingent on the label distribution. The Shakespeare dataset,
distinct in its data distribution, witnessed convergence by the 120th round
and realized an accuracy of approximately 48%. These empirical findings not
only attest to the adaptability and resilience of FeDEQ but also position it
as a promising solution for a myriad of federated learning scenarios.

####  V-B3 Model Size and Communication Efficiency.

In edge networks, the implementation of FL often confronts formidable
communication challenges. Such challenges arise from the recurring necessity
to update and transmit large and deep model parameters during the training
process. Here, we highlight the prowess of FeDEQ in achieving efficient
communication, even with compact model size, without compromising on
performance excellence. Specially, we examine the efficacy of FeDEQ with DEQs,
in comparison to FedRep with explicit models. For our experiments with FeDEQ,
we utilize DEQ-ResNet-S for FEMNIST, DEQ-ResNet-M for both CIFAR-10 and
CIFAR-100 and DEQ-Transformer for the Shakespeare dataset. In contrast, for
FedRep, we delve into a range of models from more compact architectures (e.g.,
ResNet-14, Transformer-4) to their intricate counterparts (e.g., ResNet-34,
Transformer-12), across various datasets.

The results shown in Fig. 8 reveal that FeDEQ consistently maintains a smaller
model size for all datasets while achieving competitive accuracy levels. In
FEMNIST, FeDEQ with DEQ-ResNet-S outperforms FedRep using the ResNet-14 model,
which is 1.5 times larger in size. Remarkably, even with a model size that is
$2-4$ times smaller, FeDEQ achieves performance metrics comparable to deeper
architectures, such as ResNet-20 and ResNet-34. This trend of maintaining
compactness without compromising accuracy is consistently observed in CIFAR-10
and CIFAR-100, where FeDEQ either matches or slightly exceeds the performance
of other benchmarked models. Furthermore, in the Shakespeare dataset, FeDEQ’s
implementation with DEQ-Transformer consistently presents a model that is up
to $3$ times more compact than standard Transformers, yet the accuracy remains
competitive.

These findings underscore FeDEQ’s potential as a robust solution for edge
environments that prioritize both model efficiency and performance. A smaller
model size implies more efficient update exchanges in FL scenarios, an
essential aspect in edge environments characterized by limited bandwidth and
computational resources. The compact nature of FeDEQ ensures that edge devices
engage in communication processes that are both resource-efficient and rapid,
thereby facilitating expedited decision-making.

Figure 8: Comparison the model size and accuracy of FeDEQ and FedRep across
the datasets. TABLE III: Test accuracy of different algorithms across edge
nodes

Datasets ($c$, $p$)† | FEMNIST | CIFAR-10 | CIFAR-100 | Shakespeare
---|---|---|---|---
Algorithms | (100, -) | (200, -) | (100, 2) | (100, 5) | (100, 5) | (100, 20) | (200, -)
Local | 68.80 | 68.77 | 83.85 | 65.78 | 66.67 | 34.93 | 28.70
FedAvg | 83.50 | 82.12 | 38.91 | 63.78 | 33.05 | 33.06 | 48.07
FedAvg + FT | 84.98 | 85.40 | 88.75 | 79.88 | 77.14 | 55.71 | 46.60
Ditto | 87.93 | 88.43 | 90.27 | 82.18 | 80.37 | 57.18 | 47.56
FedPer | 85.10 | 85.42 | 87.89 | 80.38 | 78.88 | 55.07 | 44.66
FedRep | 85.29 | 86.18 | 87.81 | 80.25 | 79.53 | 56.30 | 48.01
kNN-Per | 84.94 | 85.36 | 87.43 | 80.18 | 76.19 | 56.27 | 48.77
FeDEQ (Ours) | 87.13 | 87.46 | 90.54 | 82.84 | 80.42 | 57.49 | 47.62

  * † ($c$, $p$) corresponding to $c$ nodes and $p$ classes per node). For FEMNIST and Shakspeare, $p$ might be various denoted by ”-”.

####  V-B4 Performance of Personalized Models.

To investigate how FeDEQ addresses the data heterogeneity issue through
personalization, we evaluate the performance of FeDEQ and various baselines on
the local test datasets of edge nodes (unseen during training) across a range
of tasks, as detailed in Table III. As anticipated, FeDEQ consistently
outperforms non-personalized methods (Local, FedAvg, and FedAvg+FT), which
generally exhibit lower performance compared to personalized approaches in the
context of data heterogeneity. Taking the FEMNIST dataset as an example, while
the best non-personalized algorithm, FedAvg+FT, achieves accuracies of 84.98%
for 100 nodes and 85.40% for 200 nodes, FeDEQ impressively scores 87.13% and
87.46% respectively, demonstrating a clear advantage in the realm of
approximately 18%.

When pitted against personalized FL baselines, FeDEQ showcases competitive
performance with Ditto and faintly surpasses FedPer, FedRep, and kNN-Per. On
FEMNIST dataset, FeDEQ achieves accuracies faintly less than Ditto but
outshines the other methods. In the CIFAR-10 dataset with various degrees of
heterogeneity, FeDEQ achieves top accuracies of 90.54% and 82.84%
respectively. These results marginally surpass the performance of other
personalized algorithms like Ditto, which scores 90.27% and 82.18% under the
same conditions. Similarly, for the CIFAR-100 dataset, FeDEQ’s performance
peaks at 80.42% and 57.49% for configurations (100, 5) and (100, 20),
respectively, again nudging past the closest competitor, Ditto, which posts
figures of 80.37% and 57.18%. Furthermore, in the next-character prediction
task, while kNN-Per achieves the highest accuracy, FeDEQ closely follows with
a small gap, exhibiting competitiveness. Consequently, FeDEQ consistently
delivers comparable or superior performance with SOTA algorithms across all
datasets, while offering substantial efficiency in terms of communication with
up to $4$ times smaller in the size of representation.

(a)

(b)

(c)

(d)

Figure 9: Effects of Shared Representation on Personalization in FEMNIST (a,
b) and CIFAR-10 (c, d) TABLE IV: Test accuracy across unseen edge nodes (test
accuracy of training nodes between parentheses)

Datasets ($c$, $u$, $p$)† | FEMNIST | CIFAR-10 | CIFAR-100 | Shakespeare
---|---|---|---|---
Algorithms | (180, 20, -) | (90, 10, 5) | (90, 10, 20) | (180, 20, -)
FedAvg | 80.04 (82.30) | 50.20 (57.76) | 28.40 (28.90) | 46.92 (47.68)
Ditto | 81.65 (88.70) | 50.90 (82.34) | 29.60 (55.86) | 47.20 (47.24)
FedPer | 83.47 (85.20) | 78.20 (79.63) | 59.20 (52.82) | 45.47 (44.53)
FedRep | 84.68 (86.12) | 78.90 (79.11) | 57.40 (52.97) | 47.89 (48.00)
kNN-Per | 84.48 (85.40) | 75.70 (80.72) | 50.70 (53.73) |  48.26 (48.59)
FeDEQ (Ours) |  85.69 (87.00) |  79.20 (83.18) |  60.10 (54.22) |  47.25 (47.46)

  * † ($c$, $u$, $p$) corresponding to $c$ training nodes, $u$ unseen nodes and $p$ classes per node. For FEMNIST and Shakspeare, $p$ might be various denoted by ”-”.

####  V-B5 Generalization to new edge nodes.

We demonstrate the benefits of the shared representation in generalizing to
new, unseen edge nodes. To achieve this, we split the edge nodes into two
groups: 90% training nodes, and 10% new, unseen nodes. Initially, we train the
shared representation for FeDEQ, FedPer FedRep, and kNN-Per using the
designated training group, then use it to refine the personalization layers
for new nodes using their training data, which has not been incorporated
during the representation training process. For FedAvg and Ditto, a full model
will be trained and adapted to new participants. Table IV presents the
performance of these new edge nodes.

Notably, kNN-Per underperforms FedPer and FedRep in vision tasks as its shared
representation is trained from scratch rather than using pre-trained models as
in the original proposal, but excels all in the sequence task. Despite Ditto
showcases effective personalization on training nodes, it falls short when
dealing with new nodes. Meanwhile, FeDEQ demonstrates remarkable
competitiveness, consistently surpassing the other algorithms on FEMNIST,
CIFAR-10, CIFAR-100 while keeping acceptable performance on Shakespeare. This
implies that the shared representation facilitated by FeDEQ possesses the
capacity for effective adaptation to newly joined nodes, even with a
significantly reduced size.

####  V-B6 Effects of Shared Representation on Personalization

Our investigation turns towards assessing the robustness and adaptability of
the shared representation learned by FeDEQ, FedPer, and FedRep, particularly
on the FEMNIST and CIFAR-10 datasets. For a comprehensive assessment, two
tailored scenarios have been meticulously designed. In the first scenario, we
vary the local epochs for updating personalized parameters while keeping a
constant number of epochs for training the shared representation. As depicted
in Fig. 9(a),9(c), in FEMNIST, all methods only require a few personalized
epochs to reach peak performance due to the simplicity of the dataset;
however, overfitting becomes a concern shortly thereafter. Meanwhile, in
CIFAR-10, FeDEQ achieves desirable results with fewer personalized updates
than the others.

In the second scenario, the emphasis is shifted towards the epochs for
training the shared representation while the epochs for personalized parameter
updates are kept constant. We observe that a satisfactory shared
representation could be achieved within a few epochs across all methods on
both datasets, but FeDEQ potentially needs fewer to attain a good one, as
shown in Fig. 9(b),9(d). These results imply that FeDEQ can learn
representations more effectively than its explicit counterparts. Furthermore,
FeDEQ’s performance consistently surpasses the others in both scenarios,
highlighting its ability to maintain stable performance on personalization.

Figure 10: Memory usage vs. training time of DEQ-ResNets and ResNets.

####  V-B7 Memory and Time Complexity.

As FL to be deployed the network edge, efficient memory management has become
a non-negotiable criterion, especially in scenarios with resource constraints.
DEQs which are recognized for their consistent memory footprints during
training phases, emerge as a promising solution in this context [30]. To
demonstrate that, we measure the average GPU memory consumption during
training FedAvg with DEQ-ResNets on edge nodes using the CIFAR-100 dataset, in
comparison to FedAvg with explicit ResNet models. Fig. 10 reveals that DEQ-
ResNet-S has the lowest memory footprint due to its smaller size. DEQ-
ResNet-M’s memory usage is roughly on par with ResNet-14 and ResNet-20, but
nearly $1.5$ times lower than ResNet-34 for training each mini-batch. This can
be attributed to DEQs’ reliance on implicit differentiation, which eliminates
the need to store intermediate values of all layers for backpropagation.

Limitations. While the memory benefits of FeDEQ are significant, the biggest
drawback of FeDEQ lies in the training time, which is approximately $2-3$
times slower than explicit models (Fig. 10) due to the iterative nature of
solving DEQ’s forward and backward pass. To mitigate this, we can explore
acceleration methods for solving fixed-point iterations or implicit
differentiation.

##  VI Conclusion

The challenges of data heterogeneity, communication bottlenecks, and memory
constraints are well-known in FL at the network edge. While existing works
have extensively addressed each problem—such as using model personalization
for data heterogeneity or model compression to reduce communication cost—a
comprehensive solution to address these challenges is underexplored. In this
paper, we introduce FeDEQ, a novel approach for FL that aims to tackle these
challenges concurrently. The design of FeDEQ benefits from the compactness and
representation power of deep equilibrium models and the personalization
capabilities of local fine-tuning. We formulate an ADMM consensus optimization
scheme that allows edge nodes to learn a shared representation and then adapt
to their personalized parameters. Extensive experiments on various benchmarks
showcase FeDEQ’s ability to achieve comparable performance to state-of-the-art
methods with explicit models of up to four times larger in size, highlighting
FeDEQ’s lower memory footprint during training and hence a direct benefit to
communication in FL for edge environments.

## References

  * [1] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Arcas, “Communication-Efficient Learning of Deep Networks from Decentralized Data,” in _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics_ , ser. Proceedings of Machine Learning Research, A. Singh and J. Zhu, Eds., vol. 54. PMLR, 20–22 Apr 2017, pp. 1273–1282.
  * [2] F. Sattler, S. Wiedemann, K.-R. Müller, and W. Samek, “Robust and communication-efficient federated learning from non-i.i.d. data,” _IEEE Transactions on Neural Networks and Learning Systems_ , vol. 31, no. 9, pp. 3400–3413, 2020.
  * [3] A. Reisizadeh, A. Mokhtari, H. Hassani, A. Jadbabaie, and R. Pedarsani, “Fedpaq: A communication-efficient federated learning method with periodic averaging and quantization,” in _Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics_ , ser. Proceedings of Machine Learning Research, S. Chiappa and R. Calandra, Eds., vol. 108. PMLR, 26–28 Aug 2020, pp. 2021–2031.
  * [4] Y. Wang, L. Lin, and J. Chen, “Communication-efficient adaptive federated learning,” in _Proceedings of the 39th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, Eds., vol. 162. PMLR, 17–23 Jul 2022, pp. 22 802–22 838.
  * [5] A. Fallah, A. Mokhtari, and A. Ozdaglar, “Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach,” in _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , ser. NIPS’20. Red Hook, NY, USA: Curran Associates Inc., 2020.
  * [6] C. T. Dinh, N. H. Tran, and T. D. Nguyen, “Personalized federated learning with moreau envelopes,” in _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , ser. NIPS’20. Red Hook, NY, USA: Curran Associates Inc., 2020.
  * [7] V. Smith, C.-K. Chiang, M. Sanjabi, and A. Talwalkar, “Federated Multi-Task Learning,” _arXiv:1705.10467 [cs, stat]_ , Feb. 2018, arXiv: 1705.10467.
  * [8] O. Marfoq, G. Neglia, A. Bellet, L. Kameni, and R. Vidal, “Federated multi-task learning under a mixture of distributions,” 2021.
  * [9] T. Li, S. Hu, A. Beirami, and V. Smith, “Ditto: Fair and robust federated learning through personalization,” in _International Conference on Machine Learning_ , 2020.
  * [10] L. Collins, H. Hassani, A. Mokhtari, and S. Shakkottai, “Exploiting shared representations for personalized federated learning,” in _Proceedings of the 38th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, M. Meila and T. Zhang, Eds., vol. 139. PMLR, 18–24 Jul 2021, pp. 2089–2099.
  * [11] K. Pillutla, K. Malik, A.-R. Mohamed, M. Rabbat, M. Sanjabi, and L. Xiao, “Federated learning with partial model personalization,” in _Proceedings of the 39th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, Eds., vol. 162. PMLR, 17–23 Jul 2022, pp. 17 716–17 758.
  * [12] J. Konečný, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon, “Federated Learning: Strategies for Improving Communication Efficiency,” _arXiv:1610.05492 [cs]_ , Oct. 2017, arXiv: 1610.05492.
  * [13] A. T. Suresh, F. X. Yu, S. Kumar, and H. B. McMahan, “Distributed Mean Estimation with Limited Communication,” _arXiv:1611.00429 [cs]_ , Sep. 2017, arXiv: 1611.00429.
  * [14] A. Reisizadeh, A. Mokhtari, H. Hassani, A. Jadbabaie, and R. Pedarsani, “Fedpaq: A communication-efficient federated learning method with periodic averaging and quantization,” in _Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics_ , ser. Proceedings of Machine Learning Research, S. Chiappa and R. Calandra, Eds., vol. 108. PMLR, 26–28 Aug 2020, pp. 2021–2031.
  * [15] Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra, “Federated Learning with Non-IID Data,” _arXiv:1806.00582 [cs, stat]_ , Jun. 2018, arXiv: 1806.00582.
  * [16] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the Convergence of FedAvg on Non-IID Data,” _arXiv:1907.02189 [cs, math, stat]_ , Jun. 2020, arXiv: 1907.02189.
  * [17] T. Yoon, S. Shin, S. J. Hwang, and E. Yang, “Fedmix: Approximation of mixup under mean augmented federated learning,” in _International Conference on Learning Representations_ , 2021.
  * [18] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, “SCAFFOLD: Stochastic controlled averaging for federated learning,” in _Proceedings of the 37th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, H. D. III and A. Singh, Eds., vol. 119. PMLR, 13–18 Jul 2020, pp. 5132–5143.
  * [19] A. Z. Tan, H. Yu, L. zhen Cui, and Q. Yang, “Towards personalized federated learning,” _IEEE Transactions on Neural Networks and Learning Systems_ , vol. PP, 2021.
  * [20] Y. Jiang, J. Konecný, K. Rush, and S. Kannan, “Improving federated learning personalization via model agnostic meta learning,” _ArXiv_ , vol. abs/1909.12488, 2019.
  * [21] L. Collins, H. Hassani, A. Mokhtari, and S. Shakkottai, “Fedavg with fine tuning: Local updates lead to representation learning,” 2022.
  * [22] F. Hanzely and P. Richtárik, “Federated learning of a mixture of global and local models,” _ArXiv_ , vol. abs/2002.05516, 2020.
  * [23] Y. Deng, M. M. Kamani, and M. Mahdavi, “Adaptive personalized federated learning,” _CoRR_ , vol. abs/2003.13461, 2020.
  * [24] M. G. Arivazhagan, V. Aggarwal, A. K. Singh, and S. Choudhary, “Federated learning with personalization layers,” _CoRR_ , vol. abs/1912.00818, 2019\.
  * [25] F. Hanzely, B. Zhao, and M. Kolar, “Personalized federated learning: A unified framework and universal optimization techniques,” _CoRR_ , vol. abs/2102.09743, 2021.
  * [26] K. Singhal, H. Sidahmed, Z. Garrett, S. Wu, J. K. Rush, and S. Prakash, “Federated reconstruction: Partially local federated learning,” in _Advances in Neural Information Processing Systems_ , A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, Eds., 2021.
  * [27] S. Lin, Y. Han, X. Li, and Z. Zhang, “Personalized federated learning towards communication efficiency, robustness and fairness,” in _Advances in Neural Information Processing Systems_ , A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, Eds., 2022.
  * [28] O. Marfoq, G. Neglia, R. Vidal, and L. Kameni, “Personalized federated learning through local memorization,” in _Proceedings of the 39th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, Eds., vol. 162. PMLR, 17–23 Jul 2022, pp. 15 070–15 092.
  * [29] I. Achituve, A. Shamsian, A. Navon, G. Chechik, and E. Fetaya, “Personalized federated learning with gaussian processes,” in _Advances in Neural Information Processing Systems_ , A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, Eds., 2021.
  * [30] S. Bai, J. Z. Kolter, and V. Koltun, “Deep equilibrium models,” in _Advances in Neural Information Processing Systems_ , H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, Eds., vol. 32. Curran Associates, Inc., 2019.
  * [31] L. El Ghaoui, F. Gu, B. Travacca, A. Askari, and A. Tsai, “Implicit deep learning,” _SIAM Journal on Mathematics of Data Science_ , vol. 3, no. 3, pp. 930–958, 2021.
  * [32] K. Kawaguchi, “On the theory of implicit deep learning: Global convergence with implicit layers,” in _International Conference on Learning Representations_ , 2021.
  * [33] S. Bai, V. Koltun, and J. Z. Kolter, “Multiscale deep equilibrium models,” in _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 5238–5250.
  * [34] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein, “Distributed optimization and statistical learning via the alternating direction method of multipliers,” _Found. Trends Mach. Learn._ , vol. 3, no. 1, p. 1–122, jan 2011.
  * [35] X. Zhang, M. M. Khalili, and M. Liu, “Improving the privacy and accuracy of ADMM-based distributed algorithms,” in _Proceedings of the 35th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, J. Dy and A. Krause, Eds., vol. 80. PMLR, 10–15 Jul 2018, pp. 5796–5805.
  * [36] A. Elgabli, J. Park, S. Ahmed, and M. Bennis, “L-fgadmm: Layer-wise federated group admm for communication efficient decentralized deep learning,” _2020 IEEE Wireless Communications and Networking Conference (WCNC)_ , pp. 1–6, 2019.
  * [37] S. Zhou and G. Y. Li, “Communication-efficient admm-based federated learning,” _CoRR_ , vol. abs/2110.15318, 2021.
  * [38] S. Yue, J. Ren, J. Xin, S. Lin, and J. Zhang, “Inexact-admm based federated meta-learning for fast and continual edge learning,” in _Proceedings of the Twenty-Second International Symposium on Theory, Algorithmic Foundations, and Protocol Design for Mobile Networks and Mobile Computing_ , ser. MobiHoc ’21. New York, NY, USA: Association for Computing Machinery, 2021, p. 91–100.
  * [39] K. Ciesielski, “On Stefan Banach and some of his results,” _Banach Journal of Mathematical Analysis_ , vol. 1, no. 1, pp. 1 – 10, 2007.
  * [40] C. G. Broyden, “A class of methods for solving nonlinear simultaneous equations,” _Mathematics of Computation_ , vol. 19, pp. 577–593, 1965.
  * [41] H. F. Walker and P. Ni, “Anderson acceleration for fixed-point iterations,” _SIAM Journal on Numerical Analysis_ , vol. 49, no. 4, pp. 1715–1735, 2011\.
  * [42] S. Bai, V. Koltun, and J. Z. Kolter, “Neural deep equilibrium solvers,” in _International Conference on Learning Representations_ , 2022.
  * [43] M. Z. K. Johnson, David Duvenaud, “Chapter 4: Deep Equilibrium Models.”
  * [44] M. R. Hestenes and E. Stiefel, “Methods of conjugate gradients for solving linear systems,” _Journal of research of the National Bureau of Standards_ , vol. 49, pp. 409–435, 1952.
  * [45] S. W. Fung, H. Heaton, Q. Li, D. Mckenzie, S. J. Osher, and W. Yin, “Jfb: Jacobian-free backpropagation for implicit networks,” in _AAAI Conference on Artificial Intelligence_ , 2021.
  * [46] M. Hong, Z.-Q. Luo, and M. Razaviyayn, “Convergence analysis of alternating direction method of multipliers for a family of nonconvex problems,” _SIAM Journal on Optimization_ , vol. 26, no. 1, pp. 337–364, 2016.
  * [47] S. Caldas, S. M. K. Duddu, P. Wu, T. Li, J. Konečný, H. B. McMahan, V. Smith, and A. Talwalkar, “Leaf: A benchmark for federated settings,” in _Workshop on Federated Learning for Data Privacy and Confidentiality_ , 2019\.
  * [48] A. Krizhevsky, “Learning Multiple Layers of Features from Tiny Images,” p. 60, 2009.
  * [49] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-efficient learning of deep networks from decentralized data,” in _Artificial Intelligence and Statistics_. PMLR, 2017, pp. 1273–1282.
  * [50] G. Cohen, S. Afshar, J. Tapson, and A. Van Schaik, “Emnist: Extending mnist to handwritten letters,” in _2017 International Joint Conference on Neural Networks (IJCNN)_. IEEE, 2017, pp. 2921–2926.
  * [51] G. Wang, H. Yu, S. Wu, W. Dai, J. Feng, S. Li, H. Yu, T. Li, and J. Konečný, “Fedjax: Federated learning simulation with jax,” 2021.
  * [52] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” _CoRR_ , vol. abs/1512.03385, 2015.
  * [53] C. Dugas, Y. Bengio, F. Bélisle, C. Nadeau, and R. Garcia, “Incorporating second-order functional knowledge for better option pricing,” in _Advances in Neural Information Processing Systems_ , T. Leen, T. Dietterich, and V. Tresp, Eds., vol. 13. MIT Press, 2000.
  * [54] M. Dehghani, S. Gouws, O. Vinyals, J. Uszkoreit, and L. Kaiser, “Universal transformers,” in _International Conference on Learning Representations_ , 2019.
  * [55] “AMD Ryzen™ Threadripper™ 3970X Drivers & Support | AMD,” May 2022.
  * [56] “Geforce rtx 3090 family - 3090 & 3090 Ti Graphics Cards.”
  * [57] J. Bradbury, R. Frostig, P. Hawkins _et al._ , “JAX: composable transformations of Python+NumPy programs,” 2018.
  * [58] T. Hennigan, T. Cai, T. Norman, and I. Babuschkin, “Haiku: Sonnet for JAX,” 2020.
  * [59] I. Babuschkin, K. Baumli, A. Bell _et al._ , “The DeepMind JAX Ecosystem,” 2020.
  * [60] M. Blondel, Q. Berthet, M. Cuturi, R. Frostig, S. Hoyer, F. Llinares-Lopez, F. Pedregosa, and J.-P. Vert, “Efficient and modular implicit differentiation,” in _Advances in Neural Information Processing Systems_ , S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, Eds., vol. 35. Curran Associates, Inc., 2022, pp. 5230–5242.
  * [61] S. Wu, T. Li, Z. Charles, Y. Xiao, Z. Liu, Z. Xu, and V. Smith, “Motley: Benchmarking heterogeneity and personalization in federated learning,” 2022.

##  Appendix A Convergence of FeDEQ

First, we restate the original optimization problem. Let the local loss
function on each edge node be
${\mathcal{L}}_{i}\left(\theta,w_{i}\right)\vcentcolon=\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}\ell\left(f_{\theta,w_{i}}\left(x_{ij}\right),y_{ij}\right)$,
where $\theta$ is the parameters to be shared in the implicit module and
$w_{i}$ is the local parameters w.r.t. to the personalized module. The
objective is to minimize the the global loss function as follows

$\displaystyle\min_{\theta,\\{w_{i}\\}}\left\\{{\mathcal{L}}(\theta,\\{w_{i}\\})\vcentcolon=\sum_{i=1}^{n}{\mathcal{L}}_{i}(\theta,w_{i})\right\\}.$

We reframe this problem as an equivalent consensus optimization scheme:

$\displaystyle\min_{\\{\theta_{i}\\},\\{w_{i}\\},\theta}\quad\sum_{i=1}^{n}{\mathcal{L}}_{i}\left(\theta_{i},w_{i}\right)\quad\text{s.t.}\quad\theta_{i}=\theta,\quad\forall
i,$ (14)

where $\theta_{i}$ and $w_{i}$ are the local parameters on edge node $i$. The
enforcement of the constraints $\theta_{i}=\theta,\forall i$ is to encourage
the learning of a common representation module $\theta$ that benefits from all
edge nodes’ data. On the other hand, $w_{i}$ is not shared as it is used
specifically for personalization on edge node $i$.

Recognizing that Equation 14 is a consensus optimization problem, we employ
alternating direction method method of multipliers (ADMM) to solve it. This
algorithm is based on alternating updates of primal and dual variables to
achieve two goals: $\theta_{i}$ across edge nodes become closer and the
objective function decreases. An introduction to ADMM can be found in [34].
Applying this to our problem, we first construct the _augmented Lagrangian_
for each edge node $i$ as

$\displaystyle\widetilde{{\mathcal{L}}}_{i}(\theta,\theta_{i},w_{i},\lambda_{i})\vcentcolon={\mathcal{L}}_{i}(\theta_{i},w_{i})+\left\langle\lambda_{i},\theta_{i}-\theta\right\rangle+\frac{\rho}{2}\lVert\theta_{i}-\theta\rVert^{2},$
(15)

where $\lambda_{i}$ is the dual variable associated with the constraint
$\theta_{i}=\theta$, and $\rho>0$ is a hyperparameter. (We will discuss the
effect of $\rho$ in the following discussion.) All inner products and norms
are Euclidean. Note that we do not have a dual variable for $w_{i}$ because it
is not shared. Given this, we define the global augmented Lagrangian as

$\displaystyle\widetilde{{\mathcal{L}}}(\theta,\\{\theta_{i}\\},\\{w_{i}\\},\\{\lambda_{i}\\})\vcentcolon=\sum_{i=1}^{n}\widetilde{{\mathcal{L}}}_{i}(\theta,\theta_{i},w_{i},\lambda_{i}).$
(16)

Algorithm 1 in the main text describes the ADMM procedure. First, all
variables ($\theta^{0},\theta_{i}^{0},w_{i}^{0},\lambda_{i}^{0}$) are
initialized randomly. In each _global iteration_ $t$ (line 3), we have the
current variables $\theta^{t},\theta_{i}^{t},w_{i}^{t},\lambda_{i}^{t}$ and
will update them as follows:

  * •

The server aggregates $\theta_{i}^{t}$ from a selected subset of edge nodes
$S_{t}$: $\theta^{t+1}=\frac{1}{\left|S_{t}\right|}\sum_{i\in
S_{t}}\theta_{i}^{t}$ (line 4). Then the server samples a new edge node subset
$S_{t+1}$ and sends the new $\theta^{t+1}$ to them (lines 5–7).

  * •

Each chosen node in $S_{t+1}$ aims to minimize its local augmented Lagrangian
$\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i},w_{i},\lambda_{i}^{t})$
w.r.t. its parameters $\theta_{i}$ and $w_{i}$ (lines 8–9). This gives the new
variables $\theta_{i}^{t+1}$ and $w_{i}^{t+1}$.

  * •

Finally, each chosen node updates its dual variable $\lambda_{i}$ using the
ascent step:
$\lambda_{i}^{t+1}=\lambda_{i}^{t}+\rho(\theta_{i}^{t+1}-\lambda^{t+1})$ (line
10).

Here, we proceed to show that the sequence of variables
$(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})$
converges to the set of stationary points of the augmented Lagrangian,
concluding our proof. In particular, we aim to show that the augmented
Lagrangian sufficiently decreases over time (Sec. A-A). In addition,
$\widetilde{{\mathcal{L}}}$ is lower bounded by the minimum of the global loss
function ${\mathcal{L}}$ (Sec. A-B).

###  A-A Proof of Lemma IV.3

###### Proof.

First, the change in the augmented Lagrangian after one global round can be
decomposed as follows

$\displaystyle\begin{split}&\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t+1}\\},\\{w_{i}^{t+1}\\},\\{\lambda_{i}^{t+1}\\})-\widetilde{{\mathcal{L}}}(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})\\\
&=\left(\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})-\widetilde{{\mathcal{L}}}(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})\right)\\\
&\quad+\left(\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t+1}\\},\\{w_{i}^{t+1}\\},\\{\lambda_{i}^{t}\\})\right.\\\
&\qquad\qquad\qquad\qquad\left.-\
\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})\right)\\\
&\quad+\left(\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t+1}\\},\\{w_{i}^{t+1}\\},\\{\lambda_{i}^{t+1}\\})\right.\\\
&\qquad\qquad\qquad\qquad\left.-\
\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t+1}\\},\\{w_{i}^{t+1}\\},\\{\lambda_{i}^{t}\\})\right).\end{split}$
(17)

Each of the three differences on the right-hand side corresponds to the
sequential update rule for the variables in Algorithm 1: $\theta$, then
$\theta_{i}$ and $w_{i}$, and finally $\lambda_{i}$.

Bounding the first term of Equation 17. Note that by the definitions in
Equations 15 and 16, $\widetilde{{\mathcal{L}}}_{i}$ is $(n\rho)$-strongly
convex w.r.t. $\theta$. The averaging step to obtain $\theta^{t+1}$ (line 4 of
Algorithm 1) is basically performing the minimization of
$\widetilde{{\mathcal{L}}}$ w.r.t. $\theta$. In other words,
$\theta^{t}=\text{arg\,min}_{\theta}\widetilde{{\mathcal{L}}}(\theta,\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})$.
Therefore,

$\begin{split}&\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})-\widetilde{{\mathcal{L}}}(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})\\\
&\quad\leq-\frac{n\rho}{2}\lVert\theta^{t+1}-\theta^{t}\rVert^{2}.\end{split}$
(18)

Bounding the second term of Equation 17. The second difference quantifies the
changes in the augmented Lagrangian after updating the local variables
$\theta_{i}$ and $w_{i}$. Because we perform node sampling, for those edge
nodes who are not chosen at each round, their augmented Lagrangian does not
change:

$\displaystyle
w_{i}^{t+1}=w_{i}\quad\text{and}\quad\theta_{i}^{t+1}=\theta_{i}^{t},\quad\forall
i\notin S_{t+1}.$ (19)

For the chosen node, analyze the change as follows. Since ${\mathcal{L}}_{i}$
is assumed to be $L$-smooth w.r.t. $w_{i}$, we have

$\displaystyle{\mathcal{L}}_{i}(\theta_{i}^{t},w_{i})$
$\displaystyle\quad\leq{\mathcal{L}}_{i}(\theta_{i}^{t},w_{i}^{t})+\left\langle\nabla_{w_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t},w_{i}^{t}),w_{i}-w_{i}^{t}\right\rangle+\frac{L}{2}\lVert
w_{i}-w_{i}^{t}\rVert^{2}.$

By minimizing the right-hand side, we have the GD rule:
$w_{i}^{t+1}=w_{i}^{t}-\frac{1}{L}\nabla_{w_{i}}{\mathcal{L}}(\theta_{i}^{t},w_{i}^{t}).$
This leads to the following sufficient decrease to the local loss

$\displaystyle{\mathcal{L}}_{i}(\theta_{i}^{t},w_{i}^{t})-{\mathcal{L}}_{i}(\theta_{i}^{t},w_{i}^{t})\leq-\frac{L}{2}\lVert
w_{i}^{t+1}-w_{i}^{t}\rVert^{2}.$

Since the augmented Lagrangian only differs from the local loss function by a
constant w.r.t. $w_{i}$, this implies that

$\begin{split}&\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t},w_{i}^{t+1},\lambda_{i}^{t})-\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t},w_{i}^{t},\lambda_{i}^{t})\\\
&\quad\leq-\frac{L}{2}\lVert w_{i}^{t+1}-w_{i}^{t}\rVert^{2}.\end{split}$ (20)

To bound the difference but w.r.t. $\theta_{i}$, note that as per the
definition in Equation 15, $\widetilde{{\mathcal{L}}}_{i}$ is
$(L+\rho)$-smooth w.r.t. $\theta_{i}$. Using the same logic as above, this
time updating $\theta_{i}$ using the gradient of the augmented Lagrangian, we
have the following bound

$\displaystyle\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t})-\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t},w_{i}^{t+1},\lambda_{i}^{t})$
$\displaystyle\leq-\frac{L+\rho}{2}\lVert\theta_{i}^{t+1}-\theta_{i}^{t}\rVert^{2}.$
(21)

Combining Equations 19, 20 and 21, we have $\forall i$,

$\displaystyle\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t})-\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t},w_{i}^{t},\lambda_{i}^{t})$
$\displaystyle\leq-\frac{L}{2}\lVert
w_{i}^{t+1}-w_{i}^{t}\rVert^{2}-\frac{L+\rho}{2}\lVert\theta_{i}^{t+1}-\theta_{i}^{t}\rVert^{2}.$
(22)

Bounding the third term of Equation 17. Based on the update of
$\theta_{i}^{t+1}$ on line 10 of Algorithm 1, the change in the augmented
Lagrangian can be written as

$\displaystyle\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t+1})-\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t})$
$\displaystyle=\left\langle\lambda_{i}^{t+1}-\lambda_{i}^{t},\theta_{i}^{t+1}-\theta^{t+1}\right\rangle=\frac{1}{\rho}\lVert\lambda_{i}^{t+1}-\lambda_{i}^{t}\rVert^{2}.$
(23)

Notice that the first-order condition for $\theta_{i}^{t+1}$ implies that

$\displaystyle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+\lambda_{i}^{t}+(L+\rho)(\theta_{i}^{t+1}-\theta^{t+1})=0.$

Using the fact that
$\lambda_{i}^{t+1}=\lambda_{i}^{t}+\rho(\theta_{i}^{t+1}-\theta_{t})$, we have

$\lambda_{i}^{t+1}=-\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})-L(\theta_{i}^{t+1}-\theta^{t+1}).$
(24)

Therefore,

$\displaystyle\lVert\lambda_{i}^{t+1}-\lambda_{i}^{t}\rVert$
$\displaystyle\quad=\left\|\left[\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1}-\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t},w_{i}^{t}))\right]\right.$
$\displaystyle\qquad\qquad\qquad\qquad\left.+\
L\left[\theta_{i}^{t+1}-\theta_{i}^{t}\right]+L\left[\theta^{t}-\theta^{t+1}\right]\right\|$
$\displaystyle\quad\leq
2\left\|\theta_{i}^{t+1}-\theta_{i}^{t}\right\|+\left\|w_{i}^{t+1}-w_{i}^{t}\right\|+\left\|\theta^{t+1}-\theta^{t}\right\|,$

where on the second row, we use the triangle inequality and the fact that
${\mathcal{L}}_{i}$ is $L$-smooth w.r.t. $(\theta_{i},w_{i})$. This implies
that

$\displaystyle\lVert\lambda_{i}^{t+1}-\lambda_{i}^{t}\rVert^{2}$
$\displaystyle\leq 12L^{2}\left\|\theta_{i}^{t+1}-\theta_{i}^{t}\right\|^{2}$
$\displaystyle\quad+3L^{2}\left\|w_{i}^{t+1}-w_{i}^{t}\right\|^{2}+3L^{2}\left\|\theta^{t+1}-\theta^{t}\right\|^{2}.$
(25)

This result also holds for edge nodes who do not participate in round $t$,
because the left-hand side is simply zero.

Substituting Equation 25 into the right-hand side of Equation 23 gives

$\displaystyle\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t+1})-\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t})$
$\displaystyle\quad\leq\frac{12L^{2}}{\rho}\left\|\theta_{i}^{t+1}-\theta_{i}^{t}\right\|^{2}+\frac{3L^{2}}{\rho}\left\|w_{i}^{t+1}-w_{i}^{t}\right\|^{2}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad+\frac{3L^{2}}{\rho}\left\|\theta^{t+1}-\theta^{t}\right\|^{2}.$
(26)

Combining the differences. Adding together Equations 18, 22 and 26, summing
the result over all edge nodes and then the unrolling the result over $T$
global iteration gives

$\displaystyle\begin{split}&\widetilde{{\mathcal{L}}}(\theta^{t+T},\\{\theta_{i}^{t+T}\\},\\{w_{i}^{t+T}\\},\\{\lambda_{i}^{t+T}\\})-\widetilde{{\mathcal{L}}}(\theta^{t},\\{\theta_{i}^{t}\\},\\{w_{i}^{t}\\},\\{\lambda_{i}^{t}\\})\\\
&\leq\sum_{k=1}^{T}\sum_{i=1}^{n}\left[-\left(\frac{\rho}{2}-\frac{3L^{2}}{\rho}\right)\lVert\theta^{t+k+1}-\theta^{t+k}\rVert^{2}\right.\\\
&\quad\quad\quad\quad\quad\quad\left.-\left(\frac{L+\rho}{2}-\frac{12L^{2}}{\rho}\right)\lVert\theta_{i}^{t+k+1}-\theta_{i}^{t+k}\rVert^{2}\right.\\\
&\quad\quad\quad\quad\quad\quad\quad\left.-\left(\frac{L}{2}-\frac{3L^{2}}{\rho}\right)\lVert
w_{i}^{t+k+1}-w_{i}^{t+k}\rVert^{2}\right].\end{split}$ (27)

If we choose $\rho>6L$, then
$\left(\frac{\rho}{2}-\frac{3L^{2}}{\rho}\right)$,
$\left(\frac{L+\rho}{2}-\frac{12L^{2}}{\rho}\right)$ and
$\left(\frac{L}{2}-\frac{3L^{2}}{\rho}\right)$ are all positive. This allows
us to move the first sum into each squared norm, we have Equation 13. ∎

###  A-B Proof of Lemma IV.5

###### Proof.

For each edge node $i\in S_{t+1}$, we have

$\displaystyle\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t+1})$
$\displaystyle={\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+\frac{\rho}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}+\left\langle\lambda_{i}^{t+1},\theta_{i}^{t+1}-\theta^{t+1}\right\rangle$
$\displaystyle={\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+\frac{\rho}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}$
$\displaystyle\quad-\left\langle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+L(\theta_{i}^{t+1}-\theta^{t+1}),\theta_{i}^{t+1}-\theta^{t+1}\right\rangle$
$\displaystyle={\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+\frac{\rho-2L}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}$
$\displaystyle\quad+\left\langle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1}),\theta^{t+1}-\theta_{i}^{t+1}\right\rangle,$
(28)

where in the second inequality we substitute expression of $\lambda_{i}^{t+1}$
in Equation 24. Now, since ${\mathcal{L}}_{i}$ is $L$-smooth w.r.t.
$\theta_{i}$, we have

$\displaystyle{\mathcal{L}}_{i}(\theta^{t+1},w_{i}^{t+1})$
$\displaystyle\quad\leq{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+\frac{L}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}$
$\displaystyle\quad\quad+\left\langle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1}),\theta^{t+1}-\theta_{i}^{t+1}\right\rangle$
$\displaystyle\quad={\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+\frac{L}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}$
$\displaystyle\quad\quad+\left\langle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})-\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta^{t+1},w_{i}^{t+1}),\theta^{t+1}-\theta_{i}^{t+1}\right\rangle$
$\displaystyle\quad\quad+\left\langle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta^{t+1},w_{i}^{t+1}),\theta^{t+1}-\theta_{i}^{t+1}\right\rangle$
$\displaystyle\quad\leq{\mathcal{L}}_{i}(\theta_{i}^{t+1},w_{i}^{t+1})+\frac{3L}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}$
$\displaystyle\quad\quad+\left\langle\nabla_{\theta_{i}}{\mathcal{L}}_{i}(\theta^{t+1},w_{i}^{t+1}),\theta^{t+1}-\theta_{i}^{t+1}\right\rangle,$
(29)

where in the last line we use the Cauchy-Schwarz inequality and again the fact
that ${\mathcal{L}}_{i}$ is $L$-smooth. From Equations 28 and 29, we have

$\displaystyle\widetilde{{\mathcal{L}}}_{i}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t+1})$
$\displaystyle\quad\geq{\mathcal{L}}_{i}(\theta^{t+1},w_{i}^{t+1})+\frac{\rho-5L}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}.$
(30)

This implies that

$\displaystyle\widetilde{{\mathcal{L}}}(\theta^{t+1},\\{\theta_{i}^{t+1}\\},\\{w_{i}^{t+1}\\},\\{\lambda_{i}^{t+1}\\})$
$\displaystyle\quad\geq\sum_{i=1}^{n}\left({\mathcal{L}}_{i}(\theta^{t+1},w_{i}^{t+1})+\frac{\rho-5L}{2}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}\right)$
$\displaystyle\quad={\mathcal{L}}(\theta^{t+1},\\{w_{i}^{t+1}\\})+\frac{\rho-5L}{2}\sum_{i=1}^{n}\lVert\theta_{i}^{t+1}-\theta^{t+1}\rVert^{2}.$
(31)

Given the assumptions that (i) ${\mathcal{L}}$ is be bounded below by
$\underline{{\mathcal{L}}}$ and (ii) $\rho>5L$, we establish that
$\widetilde{{\mathcal{L}}}(\theta^{t+1},\theta_{i}^{t+1},w_{i}^{t+1},\lambda_{i}^{t+1})$
is lower bounded by $\underline{{\mathcal{L}}}$. ∎

##  Appendix B Experimental Details

###  B-A Model and Training Details

####  B-A1 Model Details

TABLE V: Models employed for FEMNIST, CIFAR-10 and CIFAR-100 Datasets (The
model size is shown in megabytes (MB)

Dataset | Model |  1 Initial Conv | Residual Blocks |  3 FC Layers |  Activation Function |  Total Model Size
---|---|---|---|---|---|---
|  Output Channel |  Blocks/ Group |  Ouput Channels/ Group |  Output Dimension |  |
FEMNIST | ResNet-14 |  32 |  (1, 2, 2, 1) |  (32, 64, 128, 256) |  (256, 128, 62) |  ReLU |  6.81
| ResNet-20 |  |  (2, 3, 2, 2) |  |  |  |  11.90
| ResNet-34 |  |  (3, 4, 5, 3) |  |  |  |  20.54
| DEQ-ResNet-S |  64 | 1 Block with 2 Conv: 64, 128 |  (256, 128, 62) |  Softplus |  3.97
CIFAR-10 | ResNet-14 |  64 |  (2, 2, 2) |  (64, 128, 256) |  (256, 128, 10) |  ReLU |  11.34
| ResNet-20 |  |  (3, 3, 3) |  |  |  |  17.54
| ResNet-34 |  |  (5, 6, 5) |  |  |  |  31.12
| DEQ-ResNet-M |  128 | 1 Block with 2 Conv: 128, 256 |  (256, 128, 10) |  Softplus |  10.91
CIFAR-100 | ResNet-14 |  64 |  (1, 2, 3) |  (64, 128, 256) |  (256, 128, 100) |  ReLU |  15.81
| ResNet-20 |  |  (2, 3, 4) |  |  |  |  22.01
| ResNet-34 |  |  (3, 5, 8) |  |  |  |  43.56
| DEQ-ResNet-M |  128 | 1 Block with 2 Conv: 128, 256 |  (256, 128, 100) |  Softplus |  10.95

TABLE VI: Models employed for SHAKESPEARE (The model size is shown in
megabytes (MB))

Model |  Embedding Length |  Transformer Block |  1 Linear Layer Dimension |  Total Model Size
---|---|---|---|---
Transformer-4 |  64 |  5 UT Layers, 4 Multi-Head Attention |  90 |  1.05
Transformer-8 |  |  9 UT Layers, 4 Multi-Head Attention |  |  1.85
Transformer-12 |  |  13 UT Layers, 4 Multi-Head Attention |  |  2.45
DEQ-Transformer |  64 |  5 UT Layers, 4 Multi-Head Attention |  90 |  1.05

In this section, we provide details on the models used for the experiment
parts. Particularly, the models associated with each dataset are presented as
follows.

  * •

FEMNIST, CIFAR-10, and CIFAR-100: As for baselines, we employ the
convolutional and pooling stages of ResNet-14, ResNet-20, and ResNet-34 [52],
each followed by two fully connected (FC) layers for extracting the
representation, and an additional FC layer for personalization. For FeDEQ, we
introduce DEQ-ResNet models which are designed with a similar structure to
ResNet counterparts but utilize a single equilibrium Resnet block design and
Softplus activation functions for smoothness objectives. All models, including
the specific architectures, the activation functions used, and their
respective sizes in megabytes, are detailed in Table V.

  * •

SHAKESPEARE: As detailed in Table VI, we design three explicit Transformer-
like models based on Universal Transformer [54] for the baselines, featuring
4, 8, and 12 layers respectively. Concurrently, we construct a DEQ-Transformer
model for FeDEQ. Each of these models uses an embedding length of 64, and a
linear layer with the dimensiona equal to the vocab size. The Transformer
models differ in their number of Universal Transformer (UT) layers and their
use of multi-head attention, with the DEQ-Transformer model also employing a
similar configuration. The last UT layer and the linear layer are subsequently
employed for personalization.

In our experiment setup for personalized tasks, we utilize ResNet-20 for
FEMNIST, ResNet-34 for CIFAR-10 and CIFAR-100, and Transformer-8 for the
Shakespeare dataset as baseline models. FeDEQ is implemented with DEQ variants
of the baseline models. Specific details regarding the shared parameters and
personalized parameters, including the number of model parameters and the
model sizes, are presented in Table II. From these tables, it’s evident that
DEQ-ResNet models are $2-4$ times smaller in size compared to their
traditional ResNet counterparts for each corresponding dataset. This suggests
that DEQ-ResNet models are more parameter-efficient, providing similar
functionality with fewer parameters (see Sec. V). This compactness might lead
to less memory consumption and potentially faster computation, which could be
beneficial in practical applications, especially where resources are
constrained.

Figure 11: Effects of $\rho$ on the consensus condition $\theta_{i}=\theta$
(The residual error $\|\theta_{i}-\theta\|_{2}$ are the average $L_{2}$ norm
of the errors between $\theta_{i}$ for $i=1,\ldots,n$ and $\theta$ over
iterations)

Figure 12: Effects of $\rho$ on the convergence of the shared representation
$\theta$

###  B-B Training Details

In our research, all experiments are conducted in a coding environment powered
by an AMD Ryzen 3970X Processor [55] with 64 cores and 256GB of RAM.
Additionally, an NVIDIA GeForce RTX 3090 GPU [56] is to accelerate the
training process. For softwares, we rely on Python 3 as our primary language,
supplemented by advanced machine learning and optimization frameworks: Jax
[57], Haiku [58], Optax [59], Jaxopt [60], and FedJax [51]. Furthermore, our
implementation is also inspired by some personalized FL benchmarking works
such as Motley [61].

  * •

FEMNIST: We employ ResNet-20 as the base model. We use Stochastic Gradient
Descent (SGD) as the optimization algorithm, with a learning rate of 0.01. The
model is trained using a batch size of 10 across 150 communication rounds.

  * •

CIFAR-10: We utilize ResNet-34 as the base model. The model is trained using
SGD with a learning rate of 0.01. The training proceeds with a batch size of
20 over 150 communication rounds.

  * •

CIFAR-100: We use ResNet-34 for the baseline models and DEQ-ResNet-M for
FeDEQ. The learning rate for the baseline models is set at 0.05, while for
FeDEQ, it is set at 0.01. The training is conducted with a batch size of 20
across 150 communication rounds.

  * •

Shakespeare: We employ Transformer-8 as the base model for baselines, and DEQ-
Transformer for FeDEQ. We use SGD with a learning rate of 0.1 for the training
process. The model is trained using a batch size of 16 across 400
communication rounds.

These configurations ensure a fair comparison between the baselines and FeDEQ
while maintaining a balance between computational efficiency and model
performance.

Memory Footprint: To estimate GPU memory usage, we use the system-level
nvidia-smi command and the py3nvml library, which is a Python 3 wrapper around
the NVIDIA Management Library. The memory usage is monitoring during the
training phases.

###  B-C Additional experiments on the effects of $\rho$

Here we extend more insights about the effects of $\rho$ on the performance
and convergence of FeDEQ. As discussed in Sec. V-B, the parameter $\rho$ in
FeDEQ is a tuning parameter that controls the trade-off between the dual
constraint satisfaction (i.e., how closely the local representation match the
shared representation) and the minimization of the objective function.
Therefore, the selection of $\rho$ is crucial for the performance of FeDEQ and
it can vary significantly based on the specific optimization problem.

As shown in Fig. 11 and 12, if $\rho$ is too small, the constraint
$\theta_{i}=\theta$ is weakly enforced, meaning that the local variables have
more freedom to minimize their local objective function, which might cause the
consensus variable to deviate from the optimum. This might increase the number
of iterations needed for convergence (e.g., $\rho=0.001$) or may even prevent
convergence in certain cases (e.g., $\rho=1.0$ in CIFAR-10 and CIFAR-100). On
the other hand, it can provide a better local performance as each subproblem
can be more optimized according to its local data. Conversely, if $\rho$ is
too large, the algorithm will strongly enforce the consensus constraint, which
can lead to a swift alignment of the local variables to the consensus variable
(e.g., $\rho=5.0$ or $\rho=10.0$). However, this might cause suboptimal
solutions because the strong enforcement of the consensus might prevent local
variables from reaching their local optima. In other words, it imposes high
”pressure” to force local solutions to agree, potentially at the expense of
solution quality. Through empirical evaluation, we found that $\rho=0.01$
provides the best balance for our datasets. In practice, selecting a suitable
$\rho$ often involves trial and error or a validation process. It is also
possible to use an adaptive strategy, where $\rho$ is updated during the
optimization process.