# Simplified Quantum Algorithm for the Oracle Identification Problem

Leila Taghavi<EMAIL_ADDRESS>QuOne Lab, Phanous Research and Innovation
Centre, Tehran, Iran

###### Abstract

In the oracle identification problem we have oracle access to bits of an
unknown string $x$ of length $n$, with the promise that it belongs to a known
set $C\subseteq\\{0,1\\}^{n}$. The goal is to identify $x$ using as few
queries to the oracle as possible. We develop a quantum query algorithm for
this problem with query complexity $O\left(\sqrt{\frac{n\log M}{\log(n/\log
M)+1}}\right)$, where $M$ is the size of $C$. This bound is already derived by
Kothari in 2014, for which we provide a more elegant simpler proof.

##  I Introduction

In the oracle identification problem, we have query access to bits of an
unknown $n$-bit string $x$, with the promise that it belongs to a known set
$C\subseteq\\{0,1\\}^{n}$ of size $M$. We want to determine $x$ while
minimizing the number of queries to bits of $x$. We denote this problem by
OIP$(C)$. Here, we are interested in quantum algorithms for this problem in
which case we assume we can query bits of $x$ in superposition.

We are interested in the query complexity of OIP$(C)$ in the worst case with
the promise that $|C|=M$. To this end, for a given $M$ and $n$, we say that
$Q\big{(}{\rm OIP}(M,n)\big{)}=q$ if there exists a bounded error quantum
query algorithm that solves any OIP$(C)$ where $|C|=M$ with at most $q$
queries. We note that for some sets $C$ the quantum query complexity might be
less than $Q(\text{OIP}(M,n))$, yet here we consider the worst case query
complexity over such choices of $C$.

Characterization of $Q(\text{OIP}(M,n)\big{)}$ for $M\leq n$ is easy; by
reducing the problem to Grover’s search on $M$ elements we find that
$Q(\text{OIP}(M,n)\big{)}=O(\sqrt{M})$ [3]. For the hard case of $n<M\leq
2^{n}$, Kothari [5] characterized the quantum query complexity of OIP$(M,n)$
by proposing an algorithm beyond a simple Grover’s search. This algorithm is
based on ideas from classical learning theory in combination with a
composition property of the so called filtered $\gamma_{2}$-norm. Here, we
give a simpler direct proof of the same result that eliminates the need for
the filtered $\gamma_{2}$-norm. Our proof is based on the framework of [2]
that converts a classical algorithm into an improved quantum one.

Assume that we have a classical algorithm that computes a function
$f:[\ell]^{n}\to[m]$ with query complexity $T$. Moreover, assume that we have
a _guessing algorithm_ that tries to predict values of queried bits, making at
most $G$ mistakes. Using these classical algorithms, we can design a quantum
query algorithm for computing $f$ with query complexity $O(\sqrt{GT})$ [2, 6].
This result in [2] is proven based on the framework of non-binary span
programs (NBSP) [1], which is a generalization of the span programs [9] for
functions with non-binary input and/or output alphabets. In this paper, once
again using the framework of NBSPs we prove a generalization of the
aforementioned result of [2]. Next, using that generalization we give a simple
proof of the following bound on $Q(\text{OIP}(M,n)\big{)}$.

######  Theorem 1.

Suppose that we are given a set $C\subseteq\\{0,1\\}^{n}$, where $|C|=M$ and
$n<M\leq 2^{n}$. Also suppose that we have query access to bits of a string
$x\in\\{0,1\\}^{n}$ with the promise that $x\in C$. Then quantum query
complexity of identifying $x$ is $O\left(\sqrt{\frac{n\log M}{\log(n/\log
M)+1}}\right)$.

Note that the bound of this theorem is tight; it is shown by Kothari [5] that
there exists a set $C$ of size $M$ for which
$Q\big{(}\text{OIP}(C)\big{)}=\Omega\left(\sqrt{\frac{n\log M}{\log(n/\log
M)+1}}\right)$.

##  II Classical to quantum query algorithm

A classical query algorithm can be modeled using a _decision tree_. A decision
tree is a directed acyclic graph which depicts the sequence of queried bits
along the algorithm. Any internal node of this tree is labeled by an index
$i\in[n]$, and output edges of this node are labeled using possible outcomes
of the query of $x_{i}$. Furthermore, every leaf of the decision tree is
labeled by an output value of the algorithm. Let $P_{x}$ be the unique path
from the root to a leaf, for which the labels of edges matches the bits of
$x$.

We consider _guessing algorithms_ along with decision trees. A guessing
algorithm is a classical algorithm that predicts the values of queries. Such
an algorithm can be represented by an edge-coloring of the decision tree. In
this coloring we color any edge associated to an output of the guessing
algorithm black and color the rest of edges red. We call such a coloring of
edges of a decision tree a _G-coloring_. More formally, a G-coloring is a
coloring of the edges of the decision tree using two colors black and red in
such a way that every node has exactly one outgoing edge with black color. See
Figure 1 for an example of a decision tree and its G-coloring.

Figure 1:  A decision tree for a function that outputs $1$ if the input string
contains at least two $0$s. Any leaf of the tree is labeled with its
associated output, and other vertices are labeled with the index that the
classical algorithm queries in that vertex. A G-coloring is represented by the
edge coloring of the decision tree using two colors black and red. This
guessing algorithm always guesses that the output of the queried bit is $1$
(black edges). Edges with label $0$ are colored red and indicate wrong guesses
of the guessing algorithm. The depth of this tree is $T=n$ which is the
classical query complexity of this algorithm for this function. The maximum
number of wrong guesses of the guessing algorithm is $G=2$ since after seeing
the first two mistakes (red edges) the output of the function will be
revealed.

To prove our main result, here we first give a bound on the quantum query
complexity of a problem based on a decision tree and a G-coloring on it. This
result is proven based on ideas from [2].

######  Theorem 2.

Assume that we have a decision tree $\mathcal{T}$ for a function
$f:D_{f}\to[m]$ with $D_{f}\subseteq\\{0,1\\}^{n}$ whose depth is $T$.
Furthermore, assume that for a G-coloring of the edges of $\mathcal{T}$, the
number of red edges in each path from the root to leaves of $\mathcal{T}$ is
at most $G$. Let $G_{x}$ be the number of red edges in $P_{x}$, and for $1\leq
g\leq G_{x}$, let $T_{g,x}$ be the number of black edges in $P_{x}$ after the
$g$-th red edge and before the next red one. Also let $T_{0,x}$ be the number
of black edges before the first red edge in $P_{x}$. Then there exists a
bounded error quantum query algorithm that computes the function $f$ with
query complexity

$O\left(\max_{x}\sum_{g=0}^{G_{x}}\sqrt{T_{g,x}}\right).$ (1)

Before getting into the proof of this theorem, let us make an intuition about
its statement. Think of the guessing algorithm as a reference that answers to
our queries. We look for its mistakes and truncate the classical algorithm
into $G+1$ parts, where $G$ is the number of mistakes that the guessing
algorithm makes. In each of these truncated pieces, the guessing algorithm
predicts all the necessary queries correctly, so in our quantum algorithm we
can follow the classical one without making any queries. To count the number
of queries, note that there exists a variant of Grover search algorithm that
finds the first marked element in a list of $n$ elements making
$O\left(\sqrt{j}\right)$ queries, where $j$ is the index of this marked
element [5]. Using this result we can find the $i$th mistake of the guessing
algorithm using $O\left(\sqrt{T_{i-1,x}}\right)$ queries. The total number of
queries will then be $\sum_{g=0}^{G_{x}}\sqrt{T_{g,x}}\log n$, where the extra
$\log n$ factor is needed for error reduction111The Grover search algorithm is
a bounded error algorithm meaning that the probability of getting a correct
answer is at least $\frac{2}{3}$. By successive calls of this algorithm on
different sets, the errors in different outputs aggregate and the probability
of getting the correct final answer becomes increasingly small. To reduce this
error we can repeat each Grover call $\log n$ times and get the majority vote.
This makes the total error bounded with the cost of a $\log n$ factor. In our
proof we design a non-binary span program for this problem that mimics the
behavior of the stated algorithm without getting this extra $\log n$ factor in
its complexity.

Figure 2:  A decision tree for the OR function. The output of this function is
the OR of the input bits, so the output is 1 if and only if there exists an
input index with value 1. The guessing algorithm always guesses that the
output of the queried bit is 0 (black edges). Red edges (wrong guesses of the
guessing algorithm) have label 1. The depth of this tree is $T=n$, which is
the classical query complexity of this algorithm for the OR function. The
maximum number of wrong guesses of the guessing algorithm is $G=1$, since
after seeing the first mistake (red edge) the output of the function will be
revealed. In the left, you can see all notations that are defined in the proof
of Theorem 2.

To present this span program first we need to develop some notations. Let
$V(\mathcal{T})$ be the vertex set of $\mathcal{T}$. Then for every internal
vertex $v\in V(\mathcal{T})$, its associated index is denoted by $J(v)$, i.e.,
$J(v)$ is the index $1\leq j\leq n$ that is queried by the classical algorithm
at node $v$. The two outgoing edges of $v$ are indexed by elements of
$\\{0,1\\}$ and connect $v$ to two other vertices. We denote these vertices by
$N(v,0)$ and $N(v,1)$. That is, $N(v,q)$, for $q\in\\{0,1\\}$, is the next
vertex that is reached from $v$ after following the outgoing edge with label
$q$. We also represent the G-coloring of edges of $\mathcal{T}$ by a function
$C(v,q)\in\\{{\rm{black}},{\rm{red}}\\}$ where $v$ is an internal vertex,
$q\in\\{0,1\\}$. $C(v,q)$ is the color of the outgoing edge of $v$ with label
$q$. Also let $b(v)$ be the number of black edges after the last red edge in
the path from the root to the vertex $v$. These notations are depicted in an
example in Figure 2.

###### Proof.

For every $x\in D_{f}$ there is an associated leaf of the tree $\mathcal{T}$
that is reached once we follow edges of the tree with labels $x_{j}$, starting
from the root. In order to find $f(x)$ it suffices to find this associated
leaf because this is what the classical query algorithm does; once we find the
leaf associated to $x$, we find the path that the classical query algorithm
would take and then find $f(x)$. Thus in order to compute $f$, we may compute
another function $\tilde{f}$ which given $x$ outputs its associated leaf of
$\mathcal{T}$, and to prove an upper bound on the quantum query complexity, it
suffices to design an algorithm for $\tilde{f}$.

We use the framework of non-binary span programs (NBSPs) for bounding the
quantum query complexity of $\tilde{f}$. For more details on NBSPs we refer to
[1] and here we only present the ingredients of an NBSP through the example of
our particular function $\tilde{f}$:

  * •

The first ingredient of an NBSP is a finite-dimensional vector space
$\mathcal{V}$ that is called the input space. Here, in our problem the input
space is determined by the orthonormal basis indexed by vertices of the
decision tree $\mathcal{T}$:

$\\{\left|v\right\rangle\,|\,v\in V(\mathcal{T})\\},$ (2)

  * •

An NBSP contains some target vectors
$|t_{0}\rangle,|t_{2}\rangle,\ldots,|t_{m-1}\rangle\in\mathcal{V}$, one for
any possible value of the function. Here, the output values of $\tilde{f}$ are
indexed by leaves $u$ of $\mathcal{T}$ and we let

$\left|t_{u}\right\rangle=\left|r\right\rangle-\left|u\right\rangle,$ (3)

where $r\in V(\mathcal{T})$ is the root of the tree.

  * •

The input vectors $I_{j,q}$ of an NBSP are some subsets
$I_{j,q}\subseteq\mathcal{V}$ for every $1\leq j\leq n$ and $q\in\\{0,1\\}$.
Here, the input sets are defined by

$I_{j,q}=\left\\{\sqrt{W_{C(v,q),b(v)}}\big{(}\left|v\right\rangle-\left|{N(v,q)}\right\rangle\big{)}\,\Big{|}\,\forall
v\in V(\mathcal{T})\text{ s.t. }J(v)=j\right\\},$ (4)

where $W_{{\rm{black}},b}$ and $W_{{\rm{red}},b}$ are positive real numbers to
be determined 222Note that the weight of every edge $(v,q)$ not only depends
on its color but also depends on $b(v)$, that is the number of black edges
after the last red edge in the path from root to $v$..

Having all these ingredients, we let $I\subseteq\mathcal{V}$ be

$I=\bigcup_{j=1}^{n}\bigcup_{q\in[\ell]}I_{j,q},$ (5)

and for every $x\in D_{f}$ we define the set of _available vectors_ $I(x)$ by

$I(x)=\bigcup_{j=1}^{n}I_{j,x_{j}}.$ (6)

Now we say that the above NBSP evaluates $\tilde{f}$ if for every $x$ the
target vector $\left|t_{\alpha}\right\rangle$ belongs to the span of available
vectors $I(x)$ if and only if $\alpha=\tilde{f}(x)$. Furthermore, there should
be _negative_ and _positive witnesses_ for this, that are explained below.

For every vertex $v$ of $\mathcal{T}$, let $P_{v}$ be the (unique) path from
the root $r$ to the vertex $v$. Then, for every $x\in D_{f}$ there exists a
path $P_{x}=P_{\tilde{f}(x)}$ from the root of the decision tree to the leaf
$\tilde{f}(x)$.

Thus, the target vector $\big{|}t_{\tilde{f}(x)}\big{\rangle}$ equals

$\displaystyle\big{|}t_{\tilde{f}(x)}\big{\rangle}=\left|r\right\rangle-\big{|}\tilde{f}(x)\big{\rangle}=\sum_{v\in
P_{x}}\frac{1}{\sqrt{W_{C\left(v,x_{J(v)}\right),b(v)}}}\left\\{\sqrt{W_{C\left(v,x_{J(v)}\right),b(v)}}\left(\left|v\right\rangle-\left|N(v,x_{J(v)})\right\rangle\right)\right\\},$
(7)

where the vectors in the braces are all available for $x$. This shows that
$\left|t_{\tilde{f}(x)}\right\rangle$ is the only target vector belonging to
the span of $I(x)$ We show that there exists a negative witness for it,
defined by

$\left|\bar{w}_{x}\right\rangle=\sum_{v\in
P_{x}}\left|v\right\rangle\in\mathcal{V}.$ (8)

We note that that $\left|\bar{w}_{x}\right\rangle$ is orthogonal to all
available vectors in $I(x)$. Moreover, it satisfies
$\left\langle\bar{w}_{x}\right|t_{\alpha}\rangle=\left\langle\bar{w}_{x}\right|r\rangle=1$
for every $\alpha\neq\tilde{f}(x)$. These two facts ensure that
$\left|t_{\alpha}\right\rangle$ for $\alpha\neq\tilde{f}(x)$ does not belong
to the span of $I(x)$.

Now based on the results of [1] the above NBSP gives a bound on the query
complexity of $\tilde{f}$. To evaluate this bound we need to estimate two
quantities called the _positive complexity_ and _negative complexity_ of the
NBSP. The positive complexity denoted by $\mathrm{wsize}^{+}(P,w,\bar{w})$ is
the maximum of the squared norm of the coefficient-vector in the expansion
(7). For any $x$ this norm is computed as

${\rm{wsize}}^{+}(x)=\sum_{g=0}^{G_{x}}\left(\frac{1}{W_{{\rm{red}},T_{g,x}}}+\sum_{b=0}^{T_{g,x}-1}\frac{1}{W_{{\rm{black}},b}}\right).$
(9)

The negative complexity for any $x$ denoted by ${\rm{wsize}}^{-}(x)$ is equal
to the sum of the squared inner product of $\left|\bar{w}_{x}\right\rangle$
with all vectors in $I$. To compute this quantity we need to compute the
overlap of $\left|\bar{w}_{x}\right\rangle$ with all vectors of the form

$\sqrt{W_{C(v,q),b(v)}}\Big{(}\left|v\right\rangle-\left|{N(v,q)}\right\rangle\Big{)}.$
(10)

We note that such a vector contributes in the negative complexity (is not
orthogonal to $\left|\bar{w}_{x}\right\rangle$) only if its corresponding edge
$\\{v,N(v,q)\\}$ leaves the path $P_{x}$, i.e., they have only the vertex $v$
in common. In this case, the contribution would be equal to $W_{C(v,q),b(v)}$,
which is the weight of that edge. Therefore, we have

${\rm{wsize}}^{-}(x)=\sum_{g=0}^{G_{x}}\left(W_{{\rm{black}},T_{g,x}}+\sum_{b=0}^{T_{g,x}-1}W_{{\rm{red}},b}\right).$
(11)

Now letting
$W_{{\rm{red}},b}=\frac{1}{W_{{\rm{black}},b}}=\sqrt{b+1}-\sqrt{b}$, both the
positive and negative witnesses are bounded by

$\sum_{g=0}^{G_{x}}\left(\frac{1}{\sqrt{T_{g,x}+1}-\sqrt{T_{g,x}}}+\sum_{b=0}^{T_{g,x}-1}\sqrt{b+1}-\sqrt{b}\right)=O\left(\sum_{g=0}^{G_{x}}\sqrt{T_{g,x}}\right),$
(12)

where we used the fact that

$\displaystyle\left(\sqrt{T_{g,x}+1}-\sqrt{T_{g,x}}\right)^{-1}=$
$\displaystyle\left(\sqrt{T_{g,x}}\Big{(}\sqrt{\frac{T_{g,x}+1}{T_{g,x}}}-1\Big{)}\right)^{-1}$
(13) $\displaystyle=$
$\displaystyle\;O\left(\frac{1}{\sqrt{T_{g,x}}}\Big{(}1+\frac{1}{2T_{g,x}}-1\Big{)}^{-1}\right)=O\left(\sqrt{T_{g,x}}\right).$
(14)

Putting these together the complexity of the NBSP is defined and bounded by

$\displaystyle\mathrm{wsize}=\sqrt{\max_{x}\mathrm{wsize}^{+}(x)\,\cdot\,\max_{x}\mathrm{wsize}^{-}(x)}=O\left(\max_{x}\sum_{g=0}^{G_{x}}\sqrt{T_{g,x}}\right).$
(15)

This is also a bound on the quantum query complexity of $\tilde{f}$, and that
of $f$. ∎

##  III Proof of the Main Result

Recall that in the oracle identification problem we are given a set
$C\subseteq\\{0,1\\}^{n}$, where $|C|=M$, and query access to a string
$x\in\\{0,1\\}^{n}$, with the promise that $x\in C$. We want to identify the
input $x$ using the least number of queries to the input oracle. We use
Theorem 2 to design a quantum query algorithm for this problem. To apply this
theorem we need a classical algorithm and a guessing algorithm, for which we
borrow ideas from [5]. The classical algorithm chooses the order in which we
query bits of $x$, and the guessing algorithm predicts the queries outputs
based on the set $C$. In the beginning, any element in $C$ is a possible
candidate for $x$. We need to choose an index of $x$ and predict a value for
it in such a way that an incorrect guess eliminates the biggest portion of
members of $C$ from our possible candidates. We continue removing impossible
candidates form $C$ until we are left with a single element which is equal to
$x$. Using this process we can find $x$ with minimum number of incorrect
guesses.

The following lemma indicates this special order of queries and their
predicted outcomes.

######  Lemma 3 ([4]).

For any set $C\in\\{0,1\\}^{n}$, there exists a string $s\in\\{0,1\\}^{n}$ and
a permutation $\pi:\\{1,\ldots,n\\}\to\\{1,\ldots,n\\}$, such that for any
$j\in\\{1,\ldots,n\\}$, we have $|C_{j}|\leq\frac{|C|}{\max\\{2,j\\}}$, where

$C_{j}=\big{\\{}c\in C\big{|}\;c_{\pi(j)}\neq
s_{\pi(j)},c_{\pi(i)}=s_{\pi(i)}\;\forall i<j\big{\\}}$ (16)

Now we have all ingredients to prove Theorem 1.

###### Proof of Theorem 1.

The classical algorithm for this problem is as follows. We first apply Lemma 3
to the set $C^{(1)}=C$ and obtain a string $s^{(1)}$ and a permutation
$\pi_{1}$. We then query the indices one by one using the order defined by the
permutation $\pi_{1}$ (the $i$-th query is the index $\pi_{1}(i)$). The
guessing algorithm predicts the value of this query to be
$s^{(1)}_{\pi_{1}(i)}$. We continue to query until we see a mismatch between a
queried value and the corresponding bit of $s^{(1)}$, that corresponds to a
wrong guess. Suppose that the first mismatch is in the $p_{1}$-th query, so we
have $x_{\pi_{1}(p_{1})}\neq s^{(1)}_{\pi_{1}(p_{1})}$. It is now guaranteed
that $x\in C^{(1)}_{p_{1}}$ and we can eliminate at least a fraction
$1-\frac{1}{\max\\{2,p_{1}\\}}$ of members of $C^{(1)}$ from possible
candidates for $x$. We then apply Lemma 3 once again to the set
$C^{(2)}=C^{(1)}_{p_{1}}$ and obtain a new string $s^{(2)}$ and a new
permutation $\pi_{2}$. We then continue to query using the order defined by
$\pi_{2}$ until we find a mismatch between a queried value and the string
$s^{(2)}$. This again eliminates some members of $C^{(2)}$. We continue this
process until we reach some set $C^{(r)}$ with $|C^{(r)}|=1$ and output the
only member of $C^{(r)}$ as the value of $x$.

In the above algorithm, the number of wrong guesses (red edges observed in the
decision tree) is $r$. Moreover, using Theorem 2, the quantum query complexity
of this problem is $\sum_{i=1}^{r-1}\sqrt{p_{i}}$. Indeed, to find an estimate
on the quantum query complexity, we need to compute the maximum of this
quantity over all possible choices of $r$ and $p_{i}$’s. To this end, note
that the total number of black edges in any path from the root to a leaf is at
most $n$ and

$\sum_{i=1}^{r-1}p_{i}\leq n.$ (17)

We also using Lemma 3 know that the $i$-th wrong guess reduces the size of
$C^{(i)}$ by a factor of $\max\\{2,p_{i}\\}$. Therefore, we have

$\displaystyle\prod_{i=1}^{r-1}\max\\{2,p_{i}\\}\leq M.$ (18)

Thus, using Theorem 2 the quantum query complexity of this problem is the
maximum value of

$\sum_{i=1}^{r-1}\sqrt{p_{i}},$ (19)

subject to the constraints (17) and (18). Kothari showed that the optimal
value of this optimization problem is bounded by $O\left(\sqrt{\frac{n\log
M}{\log(n/\log M)+1}}\right)$. The proof is by finding the optimal value of
another linear program that upperbounds the optimal value of this optimization
problem and is easier to solve (see Appendix B of the full version of [5] for
more details).

Putting these together we conclude that the quantum query complexity of the
oracle identification problem is $O\left(\sqrt{\frac{n\log M}{\log(n/\log
M)+1}}\right)$. ∎

As an example assume that $C=\\{0000,0001,0011,0111,1111\\}$. Figure 3 shows
how using Theorem 1 we can discover any given input $x\in C$. In the left we
see possible updates in the set $C$ along the algorithm based on queries. In
the right the associated decision tree is depicted. Note that similar colors
in the left picture and the decision tree determines the same stage of the
algorithm, for example in the root of the decision tree which we show using
color blue, the set $C$ is the leftmost set in the left picture which has blue
color too. We also color the index that we query in each step of the algorithm
by red color in $C$. Vertices with green color are the leaves of the decision
tree, which we label by an input $x$.

Figure 3:  An example for the decision tree and its coloring based on the
algorithm of Theorem 1. In this example we assume that the input set $C$ is
equal to $\\{0000,0001,0011,0111,1111\\}$. In the left we see the updates of
the set $C$ after each query, the order in which we query is determined using
Lemma 3. For any update of the set $C$, the queried index which is the most
informative bit based on lemma 3 is colored red. In the right the associated
decision tree is depicted. Note that any colored node of the decision tree is
associated to the set $C$ in the left with the same color.

To conclude, note that the classical algorithm that we used here is based on
Khotari’s [5]. This algorithm is based on repeated applications of the Grover
search algorithm to find the mismatches between $x$ and $s^{(i)}$ and then
uses properties of the $\gamma_{2}$-norm to eliminate the $\log$ factor that
has been added for error reduction. Using our approach this $\log$ factor does
not appear in the query complexity in the first place. This can be generalized
to any problem of this form, i.e., for any algorithm with repeated calls to
Grover search we can use the same idea to eliminate the $\log$ factor we get
for error reduction.

### Acknowledgements

I would like to thank Salman Beigi for his helpful discussion and suggestions.

## References

  * [1] Salman Beigi and Leila Taghavi.  Span Program for Non-binary Functions.  Quantum Information and Computation(QIC), 19(9&10):0760–0792, 2019.
  * [2] Salman Beigi and Leila Taghavi.  Quantum Speedup Based on Classical Decision Trees.  Quantum, 4:241, March 2020.
  * [3] Andrew M. Childs, Robin Kothari, Maris Ozols, and Martin Roetteler.  Easy and hard functions for the boolean hidden shift problem.  In Leibniz International Proceedings in Informatics, LIPIcs, volume 22, pages 50–79, 2013.
  * [4] Tibor Hegedűs.  Generalized teaching dimensions and the query complexity of learning.  In Proceedings of the eighth annual conference on Computational learning theory, pages 108–117, 1995.
  * [5] Robin Kothari.  An optimal quantum algorithm for the oracle identification problem.  In Leibniz International Proceedings in Informatics, LIPIcs, volume 25, pages 482–493, 2014.
  * [6] Cedric Yen-Yu Lin and Han-Hsuan Lin.  Upper bounds on quantum query complexity inspired by the Elitzur-Vaidman bomb tester.  Theory of Computing, 12:1–35, 2016.
  * [7] The Grover search algorithm is a bounded error algorithm meaning that the probability of getting a correct answer is at least $\frac{2}{3}$. By successive calls of this algorithm on different sets, the errors in different outputs aggregate and the probability of getting the correct final answer becomes increasingly small. To reduce this error we can repeat each Grover call $\mathop{log}\nolimits n$ times and get the majority vote. This makes the total error bounded with the cost of a $\mathop{log}\nolimits n$ factor.
  * [8] Note that the weight of every edge $(v,q)$ not only depends on its color but also depends on $b(v)$, that is the number of black edges after the last red edge in the path from root to $v$.
  * [9] Ben W Reichardt.  Span programs and quantum query complexity: The general adversary bound is nearly tight for every boolean function.  In Foundations of Computer Science, 2009. FOCS’09. 50th Annual IEEE Symposium on, pages 544–551. IEEE, 2009.


# Elastically-Constrained Meta-Learner for Federated Learning

Peng Lan**footnotemark: * Donglai Chen111Equal contribution Chong Xie  Keshu
Chen  Jinyuan He
Juntao Zhang  Yonghong Chen†††Contact Author&Yan Xu Linklogis Inc.
{lanpeng, chendonglai, xiechong, chenkeshu, hejinyuan, zhangjuntao,
chenyonghong<EMAIL_ADDRESS>
###### Abstract

Federated learning is an approach to collaboratively training machine learning
models for multiple parties that prohibit data sharing. One of the challenges
in federated learning is non-IID data between clients, as a single model can
not fit the data distribution for all clients. Meta-learning, such as Per-
FedAvg, is introduced to cope with the challenge. Meta-learning learns shared
initial parameters for all clients. Each client employs gradient descent to
adapt the initialization to local data distributions quickly to realize model
personalization. However, due to non-convex loss function and randomness of
sampling update, meta-learning approaches have unstable goals in local
adaptation for the same client. This fluctuation in different adaptation
directions hinders the convergence in meta-learning. To overcome this
challenge, we use the historical local adapted model to restrict the direction
of the inner loop and propose an elastic-constrained method. As a result, the
current round inner loop keeps historical goals and adapts to better
solutions. Experiments show our method boosts meta-learning convergence and
improves personalization without additional calculation and communication. Our
method achieved SOTA on all metrics in three public datasets.

##  1 Introduction

Federated learning (FL) is a privacy-preserving distributed machine learning
algorithm that trains models by exchanging model information of all parties
without exposure to their raw data Yang et al. (2019). Typical FL algorithms
like Google’s FedAvg McMahan et al. (2017), include a server that aggregates
and distributes parameters as well as clients that own data. In each round of
training, the server broadcasts the global model to clients. Clients use local
datasets to train their models. After local training, the clients gain diverse
models depending on the local distribution and capture latent information of
the data. The server collects and aggregates clients’ parameters to update the
global model. Then the next round begins. Through the parameter aggregation,
the message is passed across the clients iteratively over communication
rounds. After training, the final aggregated model is distributed to clients
for employment. In this way, information from clients is communicated to learn
a better model than isolated training without data breach in the whole
process.

Data heterogeneity is one of the challenges in FL. In real-world applications,
clients’ data are non-IID Zhao et al. (2018). FedAvg performs well in
identical and independent distributed (IID) data but its performance suffers
significantly on non-IID data, as one global model cannot fit different data
distributions simultaneously. One solution is to allow personalized model per
client for better local fitting, which is named personalized federated
learning (PFL) Kulkarni et al. (2020).

Meta-learning-based models have been proposed to solve PFL. Meta-learning aims
at training model initialization parameters (called meta-initialization) so
that they can adapt to any task quickly Vanschoren (2019); Hospedales et al.
(2020). In each training step, the meta-initialization $\varphi$ is updated on
a batch of different tasks (called meta-batch). The tasks in meta-batch
generate data from diverse distribution. The update for $\varphi$ can be
viewed as a double-loop optimisation: inner and outer loop. The inner loop
adapts meta-initialization $\varphi$ to each task via gradient descent on
labeled data, in order to get task-specific solution. The outer loop updates
meta-initialization based on the results of the outer loop.

The procedure of meta-learning can be naturally modified to adapt to FL. By
formulating each client as a task, the per-task operations can be executed on
each client locally, and the aggregation operation for a meta-batch can be
done on the server. FedMeta Chen et al. (2018) and Per-FedAvg Fallah et al.
(2020) introduce meta-learning and treat clients as different tasks. They
search for the common optimal initial parameters for all tasks. Such
parameters can quickly adapt to local data distribution by one or several
steps of gradient update to fulfill the personalization goal.

One limitation of meta-learning-based approaches is the fluctuation in meta-
learning’s training process. In the same client and different sampling loops,
goals are unstable. That is caused by non-convexity of loss function and
randomness of stochastic gradient descent (SGD). Specifically, in Figure 1(a),
in different sampling times, meta-learning’s inner loop may lean towards
different local optima, that is fluctuation in optimization goals.

Figure 1: The meta-learning local adaptation tracks for the same task and
different sampling time. The contour plots show levels of loss function at
different places. Due to changes in meta-learning parameter training two
sampling time and randomness of SGD, in (a) local adaptation may lean toward
different local optima. while in (b) the update keeps stable goal with elastic
constraint.

Whereas outer update relies on inner loop’s fast adaptation results. For
example in Reptile Nichol et al. (2018), the results of inner loop are used as
meta-initialization’s update direction. Therefore, deviation of inner update
would cause deviation of outer update and worsen the whole meta-learning
optimization results.

To address the problem, we propose a method of constraint based on past update
information. It aims to correct current inner loop optimization goal using
historical local adaptation point in order to deal with outer update
fluctuations. We discover that improper hard constraint, such as L2 would
hinder the algorithm searching for better solutions. The algorithm gets stuck
at suboptimal solutions and is slow to converge.

The purpose of the inner update is changed from fitting local data’s ground
truth to fitting the weighted average of historical models’ prediction
probability and ground truth. The goal takes historical models into account.
When the historical model has a better fitting ability, the model updates
towards a promising optimization direction and alleviates deviation caused by
non-convexity and randomness. On the other hand, when the historical model
fails to give useful information, the goal shifts to randomly smoothed ground
truth. It motivates the model to find a better fit and prevents suboptimal
solutions. Our method incorporates both exploitation of historical fitting
results and exploration of better solutions and improves meta-learning
performance under federated scenarios.

We did extensive experiments on three public datasets. Results show that our
algorithm effectively alleviates fluctuation in meta-learning, accelerates
convergence and achieves new SOTA.

We summarize our contributions as follows:

  * •

We find meta-learning-based federated learning has fluctuation in loss through
experiments and that using historical update information can alleviate the
problem.

  * •

Based on the first point, we propose a new implicit regularization method. It
uses historical information to constrain update direction and preserve the
ability to search for global optimum. It accelerate meta-learning convergence
and has better adaptation.

  * •

Extensive experiments show our method achieves SOTA performance in three
datasets and various statistical heterogeneity settings.

##  2 RELATED WORK

Our method is related to federated learning, meta-learning and regularization.
The related works are three-fold.

###  2.1 Federated Learning

Federated learning is a machine learning technique that preserve clients’
privacy. Google’s FedAvg McMahan et al. (2017) is a basic one. It takes
weighted average of clients’ local trained model to get global model. Then it
distributes to all clients as the initialize parameter for clients’ local
training. Those FL algorithms may not perform well if clients’ data are non-
IID.

The following researches improve FedAvg for non-IID data. FedProx Li et al.
(2018) adds constraint in the local training update. It uses the L2 distance
between local parameter and global parameter as regularization term and make
convergence process stable. SCAFFOLD Karimireddy et al. (2020) proposed
control variates to correct client-drift during local training and boost
convergence of FL. In FedDyn Acar et al. (2021), risk objective for each
device is dynamically updated, such that clients’ models move towards global
optimum.

The above methods use the same global model to adapt to different
distributions in different clients. While some researches propose local model
personalization solutions that can cope with the challenge more effectively.
The model personalization can be realized in different ways. MOCHA Smith et
al. (2017) and FedEM Marfoq et al. (2021) treat each client as a task and
realized multi-task learning under federated scenarios. FedPer Arivazhagan et
al. (2019), FedRep Collins et al. (2021) and LG-FedAvg Liang et al. (2020)
break down model to feature layers and classification layers. They use part of
the model as local private layer to realize personalization. The FedPer and
FedRep share feature layer and keep classification layer private, and the LG-
FedAvg does the opposite. FedMatch Chen et al. (2021) further devises private
patch for transformer structure to achieve personalized federated QA system.
And PartialFed-Adaptive Sun et al. (2021) automatically selects personalized
layers for each client to realize cross-domain personalization. Clustered FL
methods like CFL Sattler et al. (2019) and IFCA Ghosh et al. (2020) let
clients in one cluster share a model to realize cluster-level personalization.
Model mixture methods like APFL Deng et al. (2020) and Fed-MGL Hanzely and
Richtárik (2020) train local models and global models simultaneously, and
combine local models with global models to obtain final personalized model.
FedAMP Huang et al. (2020) introduces an attentive message passing mechanism,
which achieve personalized aggregate models by promoting parameter exchange
between clients with similar data distribution. Similarly, KT-pFL Zhang et al.
(2021) proposes a trainable knowledge coefficient matrix, which reinforces the
collaboration among similar clients and facilitate the communication using
local soft predictions. The pFedHN Shamsian et al. (2021) brings hyper network
into FL to generate personalized models. The pFedME Dinh et al. (2020)
introduces Moreau Envelopes to form a regularized loss, which leverages global
model to achieve better local training.

Meta-learning-based FL approaches, such as FedMeta Chen et al. (2018), Per-
Fedavg Fallah et al. (2020) are the most related ones to ours. They aim to
learn better initial parameter for each client. They realize personalization
in clients through one or several step gradient update to quickly adapt to
local data distribution. FedMeta employed empirical analysis to validate the
effectiveness of meta-learning under federated scenarios. Per-FedAvg
theoretically proved its convergence rate. GPAL Park et al. (2021) also
employs meta-learner for FL, but it handles the lack of data variability per
client and limited communication budget, rather than PFL.

###  2.2 Meta-Learning

In meta-learning, models are trained by different tasks and learn the ability
to quickly adapt to new tasks. One typical algorithm is LSTM Meta-Learner Chen
et al. (2016). It uses training gradients in different time as the input of
LSTM and treat model parameters as cell state, in order to exploit different
tasks’ commonality in gradient update and help optimization in new task
parameters. MAML Finn et al. (2017) aims to get an optimal initialization
parameters through different training tasks, so that the model can adapt to
new tasks by one or several steps of gradient descent.

Past researches improve the meta-learning from different perspectives. iMAML
Rajeswaran et al. (2019) introduces an implicit differentiation based method
to reduce the computational and memory burden in MAML, and Reptile Nichol et
al. (2018) proposes a first-order gradient-based meta-learner to further cut
the computation overhead. PLATIPUS Finn et al. (2018) and BMAML Yoon et al.
(2018) propose probabilistic or Bayes version of meta-learner, which generates
the model distribution depending on the observed data. LEO Rusu et al. (2018)
achieves improved performances on few-shot regimes by optimizing the latent
embedding of high-dimensional model parameter. MAML++ Antoniou et al. (2018)
discusses the network architectures sensitivity of the meta-learning and makes
various modiﬁcations to achieve better generalization, convergence and
computational cost. UNICORN-MAML Ye and Chao (2021) finds that meta learning
needs more update steps in the inner loop and is sensitive to label
permutation in meta test, and introduces a single-vector method for
classification to achieve new SOTA in few-shot learning. Researches like ANIL
Raghu et al. (2020) and BOIL Oh et al. (2020) explore the how meta-learner
works focusing on whether it achieve fast adaptation or feature reuse.

Meta-Learning is initially invented to solve few-shot learning and
reinforcement learning. But its quick adaptation method has been widely used
in other areas in the subsequent research. Shin et al. (2021) extends meta-
learning to many-shot application. MLDG Li et al. (2017) applies meta-learning
in semantic parsing model, in order to learns a domain-agnostic model that can
be applied to an unseen domain. FastDFKD Fang et al. (2022) uses meta-learning
to speed up knowledge distillation.

###  2.3 Regularization

Regularization is a common technique in machine learning to avoid overfitting
in model training by reducing the magnitude of model parameters. Traditional
regularization methods such as L1, L2 use the L1 and L2 norm as additional
loss term to limit model complexity.

EWC (Elastic Weight Consolidation, Kirkpatrick et al. (2017)) is one of the
regularization methods in continual learning to overcome catastrophic
forgetting. It calculates models parameter importance weights. New model has
larger constraint weight at high importance parameters. It helps keeping past
important knowledge in continual learning and alleviate catastrophic
forgetting.

Knowledge distillation (KD) aims to transfer the prediction ability from one
model (teacher) to another (student) Gou et al. (2021). Usually the teacher
model has high prediction power. During KD, the student model learns the soft
label from the teacher model, which may contain more information than the one-
hot label. KD adds an additional loss term in the objective function. It can
constrain student model’s representation and keep part of teacher model’s
knowledge and realize goal of generalization.

Knowledge distillation and label smoothing regularization (LSR) are related
Yuan et al. (2020). LSR replaces one-hot labels with soft labels Szegedy et
al. (2016). It can be viewed as KD with an uninformative teacher. KD can
improve student model’s performance in two scenarios. On one hand, the
uninformative teacher can improve the student by the effect of label smoothing
regularization in the KD process, on the other hand, the student can learn
from the teacher if the teacher has high performance.

##  3 Methodology

###  3.1 Problem Formulation

The objective of multi-party machine learning is to maximize the overall
performance of the models for all parties, which is formally

$\mathop{min}_{(\theta_{1},\ldots,\theta_{N})}\frac{1}{N}\sum_{i=1}^{N}l_{i}(\theta_{i}),$
(1)

where $N$ is the number of clients, $\theta_{i}:\mathop{R}^{d}\rightarrow Y$
is the learning model for _i_ -th client. The $l_{i}$ is expected loss over
local distribution. Namely $l_{i}:=E_{(x,y)\sim D_{i}}l(\theta_{i}(x),y)$,
where $x\in\mathop{R}^{d}$ and $y\in Y$ are the features and label for a data
point, $D_{i}$ is the data distribution for _i_ -th client and $l$ is the loss
function such as cross entropy.

To outperform isolated training, one straightforward option is centralized
training, in which the server collects data from all parties thus enriching
the overall training data. However, this is not practical, as participants
with privacy and security concerns usually prohibit any data breach. In FL,
the server communicates model parameters rather than data with clients so that
data are isolated and the performance is improved.

Standard FL algorithms like FedAvg aggregate the passed messages and get one
global model for all clients, formally $\theta=\theta_{1}=\ldots=\theta_{N}$.
In IID setting which assumes $D_{global}=D_{1}=\ldots=D_{N}$, one model that
fits one distribution $D_{global}$ is able to generalize on every client. The
non-IID setting brings significant performances degradation, as one model
$\theta$ is hard to fit $N$ diverse distributions $D_{1}\neq\ldots\neq D_{N}$.
Personalized federated learning (PFL) solves this more challenging
circumstance by diversifying models $\theta_{1}\neq\ldots\neq\theta_{N}$ to
fit each local distribution.

###  3.2 Meta-learner for Personalized Federated Learning

Typical meta learning methods like MAML and Reptile Nichol et al. (2018),
achieve learn-to-learn via training an initialization $\varphi$ (often named
meta-initialization) which can quickly adapt to any task. In every training
step, the meta-initialization $\varphi$ is updated on a batch of different
tasks (often called meta-batch). The tasks in meta-batch generate data from
diverse distribution ${D_{1}\neq D_{2}\neq\ldots}$. The update for $\varphi$
can be viewed as a double-loop optimisation: inner and outer loop. The inner
loop adapts meta-initialization $\varphi$ to each task via gradient descent on
labeled data, in order to get task-specific solution
${\theta_{1},\theta_{2},\ldots}$. The outer loop updates meta-initialization
based on ${\theta_{1},\theta_{2},\ldots}$ with different methods depending on
the variant of meta-learners.

MAML is a typical meta-learner. It divides the task data as support set and
query set for the two loops respectively. The motivation of this division is
using a separate data set to evaluate the generalization performance of the
local adaptation on a task. In the inner loop for $i$-th task, task-specific
solution $\theta_{i}$ is got via one gradient descent step from $\varphi$ on
support set, which is formally
$\theta_{i}\leftarrow\varphi-\varepsilon_{in}\nabla_{\varphi}l_{i}(\varphi)$.

In the outer loop, based on the loss of $\theta_{i}$ on query set, $\varphi$
is updated via another gradient descent,

$\varphi\leftarrow\varphi-\varepsilon_{out}\sum_{i}\nabla_{\varphi}l_{i}(\theta_{i}).$
(2)

The $\varepsilon_{in}$ and $\varepsilon_{out}$ are the learning rates for
inner and outer loops respectively, and the losses on support and query sets
for $i$-th task are both denoted as $l_{i}$ for notation simplicity. As
$\theta_{i}$ consists of the derivative with respect to $\varphi$, the outer
update of MAML involves the second-order derivative operation, which is
computationally intensive. Reptile is a ﬁrst-order gradient-based meta-learner
proposed by OpenAI Nichol et al. (2018). It removes the second-order
derivative operation in outer loop while achieves competitive performance as
MAML. The following demonstrates Reptile double-loop optimisation:
inner update:
initialize $\theta_{i}\leftarrow\varphi$, and do

$\theta_{i}\leftarrow\varepsilon_{in}\nabla_{\theta_{i}}l_{i}(\theta_{i})$ (3)

for $\tau$ steps,
outer update:

$\varphi\leftarrow\varepsilon_{out}\sum_{i}(\theta_{i}-\varphi).$ (4)

Inspired by Reptile, our method employs similar outer loop strategy for
computational overhead reduction.

We define the notations used in previous meta-learners for FL, like Per-FedAvg
and FedMeta. First of all, the server initializes the meta-parameter
$\varphi^{0}$ and starts up communication with clients. The communication
iterates $T$ rounds. At _t-th_ communication round, a clients subset $I^{t}$
is sampled for training. The server sends last round meta-initialization
$\varphi^{t-1}$ to all clients sampled. Each client $i$ does inner update from
$\varphi^{t-1}$ to get $\theta^{t-1}_{i}$, and then calculate the per-task
outer update component $\Delta\varphi_{i}^{t-1}$ based on $\theta^{t-1}_{i}$.
These per-task components are sent back, and the server completes the
aggregation operation of outer update:

$\varphi^{t}\leftarrow\varphi^{t-1}+\varepsilon_{out}\frac{1}{|I^{t}|}\sum_{i\in
I^{t}}\Delta\varphi_{i}^{t-1}.$ (5)

Then the next round starts. After training, each client fine-tunes from the
meta-initialization $\varphi^{T}$ on the local data to get task-specific
solution $\theta_{i}$ for final employment.

Algorithm 1 FedEC
---
Parameters:
number of clients $N$;
sampling rate $r$;
number of communication rounds $T$;
outer loop learning rate $\varepsilon_{out}$;
inner loop steps $\tau$;
inner loop learning rate $\varepsilon_{in}$;
regularization weight $\alpha$;
Overall Algorithm:
1: initialize $\varphi_{0}$ at random
2: for $t=1,2,\cdots,T$ do
3: sample clients set $I^{t}$ of size $rN$
4: for each client $i\in I^{t}$ in parallel do
5: $\theta^{t,\tau}_{i}\leftarrow\boldsymbol{InnerUpdate_{i}}(\varphi^{t-1})$
6: client $i$ sends updated parameter $\theta^{t,\tau}_{i}$ to server
8: end for
9: $\boldsymbol{OuterUpdate}$:
$\varphi^{t}\leftarrow\varphi^{t-1}+\varepsilon_{out}\frac{1}{rN}\sum_{i\in
I^{t}}(\theta^{t,\tau}_{i}-\varphi^{t-1})$
10: end for
$\boldsymbol{InnerUpdate_{i}}$:
input $\varphi^{t-1}$
11: $\theta^{t,0}_{i}\leftarrow\varphi^{t-1}$
12: for $s=1,2,\cdots,\tau$ do
13: if $\hat{\theta_{i}}$ initialized do
14: $\theta^{t,s}_{i}\leftarrow\theta^{t,s-1}_{i}-\varepsilon_{in}\nabla
l_{i}(\theta^{t,s-1})$
15: else
16:
$\theta^{t,s}_{i}\leftarrow\theta^{t,s-1}_{i}-\varepsilon_{in}\nabla(l_{i}(\theta^{t,s-1}_{i})+$
$\alpha c_{i}(\theta^{t,s-1}_{i},\hat{\theta_{i}}))$
17: end for
18: initialize or update local memory module
$\hat{\theta_{i}}\leftarrow\theta^{t,\tau}_{i}$
return $\theta^{t,\tau}_{i}$

###  3.3 Model Overview

As stated in the introduction, inner loop fluctuation worsen the performance
of meta-learner for FL. To cope with this problem, we proposed a elastically-
constrained meta-learner for Federated Learning (abbreviated as FedEC). FedEC
leverages the history optimisation information to facilitate efficient and
stable update in the inner loop. We also introduce first-order gradient-based
update strategy in the outer loop, which reduces the computation overhead. Our
method is summarized in Algorithm 1. The detailed procedures of FedEC are
described as follows.

For a communication round $t$ between $1$ and $T$, the server samples a subset
of clients $I^{t}$. The clients have heterogeneous distributions, so every
client is treated as a task and the $I^{t}$ corresponds to a meta-batch. The
server sends the latest meta-initialization $\varphi_{t-1}$ (for first
communication round $\varphi_{0}$ is randomly initialized) to all sampled
clients. Then the inner update is executed on clients in parallel following
the standard FL. However, in our method, inner update depends on not only
current meta initialization $\varphi_{t-1}$ but also the historical model
$\hat{\theta_{i}}$ by introducing the constraint. The $\hat{\theta_{i}}$
stores the last inner update result, and it is updated every time the $i$-th
client sampled. Thus, for most rounds the inner loss is changed from
$l_{i}(\theta)$ to $l_{i}(\theta)+\alpha c_{i}(\theta,\hat{\theta_{i}})$,
where $c_{i}(\theta,\hat{\theta_{i}})$ refers to the constraint and $\alpha$
refers to its weight. When the $i$-th client is sampled for the first time,
the constraint term is set to 0 as $\hat{\theta_{i}}$ is not initialized by
any training. Initialized from $\varphi_{t-1}$, the parameter trains on the
local data set for $\tau$ epochs. Formally the constrained inner loop loss
$\tilde{l}_{i}$ for $i$-th client is

$\tilde{l}_{i}(\theta)=\begin{cases}l_{i}(\theta)+\alpha
c_{i}(\theta,\hat{\theta_{i}})&\hat{\theta_{i}}\quad initialized\\\
l_{i}(\theta)&else.\\\ \end{cases}$ (6)

The training can use any gradient methods such SGD, RMSprop or Adam. The
result of $i$-th client at $t$-th communication round after $\tau$ epochs of
local adaptation is denoted as $\theta^{t,\tau}_{i}$, which is tested as
employed model at current communication round. The $\theta^{t,\tau}_{i}$ is
then copied and sent to server for outer update, and also stored into
$\hat{\theta_{i}}$ to constrain next round inner loop. The inner loop is
formulated in the $InnerUpdate_{i}$ part of Algorithm 1.

FedEC adopts a first-order approximation strategy for outer update.
Concretely, given all local adaptations $\theta^{t,\tau}_{i}$ for $i$ from $1$
to $N$, FedEC approximates the per-task component of outer update by
$\theta^{t,\tau}_{i}-\varphi^{t-1}$. We denote $U_{i}^{\tau}$ as local update
on data from $i$-th client for $\tau$ epochs, and $\theta^{t,\tau}_{i}$ can be
written as $\theta^{t,\tau}_{i}=U_{i}^{\tau}(\varphi^{t-1})$.

So the outer update component for $i$-th client is

$\Delta\varphi_{i}^{t-1}=U_{i}^{\tau}(\varphi^{t-1})-\varphi^{t-1},$ (7)

which is the direction from meta-initialization to the local adaptation. The
outer update takes the average for components of all clients, formally

$\varphi^{t}\leftarrow\varphi^{t-1}+\varepsilon_{out}\frac{1}{|I^{t}|}\sum_{i\in
I^{t}}\Delta\varphi_{i}^{t-1}.$ (8)

This averaging aggregation acts like a voting process, and the averaged
direction drives meta-initialization to a point that is good for most local
adaptation. Moreover, FedEC significantly reduces the communication overhead
and simplifies the procedure, as the second-order derivative and data-split
operations are removed.

The details of the elastic constraint is illustrated below. The constraint
term is

$c_{i}(\theta,\hat{\theta_{i}})=KL(\hat{\theta_{i}}(x),\theta(x)),$ (9)

where $D_{KL}$ is the Kullback-Leibler divergence and $\hat{\theta_{i}}$ is
the reserved model. Considering the local empirical loss

$l_{i}(\theta)=CE(y,\theta(x)),$ (10)

where $CE$ is the cross entropy, $\theta(x)$ is the predicted probabilities of
training model and $y$ the one-hot embedding of the ground truth. The
constrained loss can be written as

$l_{i}(\theta)+\alpha c_{i}(\theta,\hat{\theta_{i}})=CE(y,\theta(x))+\alpha
KL(\hat{\theta_{i}}(x),\theta(x))$ (11)

We introduce a generated distribution $q(x)$ as

$q(x)=\frac{1}{1+\alpha}y+\frac{\alpha}{1+\alpha}\hat{\theta_{i}}(x).$ (12)

The $q(x)$ can be seen as a distribution generated from weighted average of
one-hot embedding and $\hat{\theta_{i}}$ predictions. Substituting $q(x)$ into
the equation (11), (11) equals to

$(1+\alpha)CE(q(x),\theta(x))+\alpha\hat{\theta_{i}}(x)\log(\hat{\theta_{i}}(x)).$
(13)

where the second term $\alpha\hat{\theta_{i}}(x)\log(\hat{\theta_{i}}(x))$ is
a constant. The constrained loss changes the optimization target by combining
the one-hot embedding and the distribution of previous prediction. Compared
with the unconstrained loss $l_{i}(\theta)$, this combination encourages
$\theta(x)$ not only to predict higher probability corresponding to true
label, but also to imitate the prediction distribution for historical model.
There exists two circumstances for historical model $\hat{\theta_{i}}$ in the
training. One is that when historical model does not fit the local
distribution properly, which often happens in the early communication. Thus,
the prediction of $\hat{\theta_{i}}$ can be seen as a random distribution.
This plays a role of regularization by smoothing target of training. The
training target is one-hot embedding of ground truth in unconstrained loss,
and the model learns to predict more extreme distribution, thus prone to over-
fitting. While the constrained target reduces the predicted probability gap
between classes, so the generalization of model is kept. The other is that
$\hat{\theta_{i}}$ gains fitting ability after certain rounds of learning.
Thus the introducing of historical prediction facilitates exploitation of
previous knowledge, thus the update can avoid local optimal and converge
quickly to equivalent performance of historical adaptation. While the ground
truth part keep model exploring better fitting for training data. So the model
would not be stuck at a historical solution and the continual exploration for
new solution is kept.

##  4 EXPERIMENT

We do extensive experiments to evaluate FedEC’s performance by comparing with
different benchmarks, as well as L2 regularization under different experiment
settings. In the case study, visualizations show the effectiveness of FedEC to
deal with model fluctuation.

###  4.1 Dataset

To evaluate the performance of our method, we use three public datasets to
build benchmark. CIFAR10 is a dataset of 60000 images with sizes of 32*32 from
10 classes. Classes include car, plane, bird, etc. CIFAR100 is similar to
CIFAR10 with 60000 images. Its difference is that this dataset further
classifies images into 20 super classes and 100 subclasses. FEMNIST Caldas et
al. (2018) is built by partitioning the data in Extended MNIST. Data summary
statistics are shown in Table 1.

DataSet | total | train | test | classes
---|---|---|---|---
CIFAR10 | 60000 | 50000 | 10000 | 10
CIFAR100 | 60000 | 50000 | 10000 | 100
FEMNIST | 34674 | 31123 | 3551 | 10

Table 1: Details of three Datasets

###  4.2 Experiment Setting

####  4.2.1 Model structure

We evaluate on different model structures. For CIFAR10 and CIFAR100, the main
model structure are two convolutional layers and three fully connected layers.
For FEMNIST, the structure is three fully connected layers. All models use
ReLU as the activate function. For evaluation metric, we use the average _acc_
values of the local test set of all clients when parameter aggregation has
been completed for each communication round.

####  4.2.2 Data splitting method

We split the data into training set and test set in a non-IID way. For CIFAR10
and CIFAR100, we set classes at each client to control the level of non-IID.
In this scenario, each client has the same number of classes and different
combination of classes. Number of images of each class at a client equals to
the total number of images at that class divided by the number of clients. For
instance, assume we set 100 clients in total and each client contains two
distinct classes, and CIFAR10 is chosen to conduct the experiment. The data
splitting method is described as follows: initially, separate each class in
CIFAR10 into 20 pieces (where 20 is computed via $2\times 100\div 10$); then,
randomly allocate each client with two classes (duplication is permitted), the
total number of allocation for each class is identical; ultimately, specific
data in the corresponding class is assigned to each client. For each
communication round, we randomly selects 10% of clients to participate in
federated training. Additionally, the allocated data to each client is
different. By changing the number of classes per client, the level of
statistical heterogeneity is tuned. Specifically, lower number classes per
client, which means fewer shared classes between clients, corresponds to a
more heterogeneous case. We also experiment on various number of clients, the
detailed statistical result of experiments on different datasets is in Table
2. The settings are consistent with Collins et al. (2021). To examine the
strength of the model-agnostic methods, we implement various common networks
on different datasets.

####  4.2.3 Hardware resource

We use one Nvidia A100 card to test our model. A100 has 40G video memory. CPU
is Intel(R) Xeon(R) Gold 6248R @ 3.0Ghz.

Model | CIFAR10 | CIFAR100 | FEMNIST |
---|---|---|---|---
| $\mathbf{(100,2)}$ | $\mathbf{(100,5)}$ | $\mathbf{(1000,2)}$ | $\mathbf{(100,5)}$ | $\mathbf{(100,20)}$ | $\mathbf{(150,3)}$ |
Local Only | 89.79 | 70.68 | 78.30 | 75.29 | 41.29 | 40.12 |
FedAvg | 42.65 | 51.78 | 44.31 | 23.94 | 31.97 | 31.26 |
SCAFFOLD | 37.72 | 47.33 | 33.79 | 20.32 | 22.52 | 17.65 |
pFedME | 60.23 | 70.01 | 62.33 | 23.64 | 34.66 | 47.66 |
LG-FedAvg | 84.14 | 63.02 | 77.48 | 72.44 | 38.76 | 41.38 |
L2GD | 81.04 | 59.98 | 71.96 | 72.13 | 42.84 | 46.56 |
APFL | 83.77 | 72.29 | 82.39 | 78.20 | 55.44 | 52.46 |
FedPer | 87.13 | 73.84 | 81.73 | 76.00 | 55.68 | 59.07 |
Per-FedAvg | 82.27 | 67.2 | 67.36 | 72.05 | 52.49 | 71.51 |
FedRep | 87.70 | 75.68 | 83.27 | 79.15 | 56.10 | 59.21 |
FedEC-wo | 91.03 | 82.26 | 90.87 | 80.29 | 58.07 | 53.27 |
FedEC-l2 | 87.38 | 71.15 | 87.13 | 74.90 | 43.09 | 72.03 |
FedEC | 92.35 | 83.78 | 92.11 | 81.75 | 59.1 | 77.3 |

Table 2: Average accuracy of local test over the ﬁnal 10 rounds out of T = 100
total rounds, where $(x,y)$ denotes $x$ clients and $y$ classes per client

###  4.3 Benchmark

We compare FedEC with other baseline methods thoroughly. Isolated training is
also added to compare, which is denoted as local only. In Table 2, we show the
average acc at the last ten rounds for each method under six groups of
parameters. FedEC outperforms the current best model in almost all criteria.
We find that traditional FL methods such as FedAvg, SCAFFOLD Karimireddy et
al. (2020) fall behind the other personalized FL methods. LG-FedAvg Liang et
al. (2020), L2GD Hanzely and Richtárik (2021) perform worse than local only.
APFL Deng et al. (2020), FedPer Arivazhagan et al. (2019) and FedRep Collins
et al. (2021) improve from previous models, but they are lower than FedEC. As
the data heterogeneity increases, such as the more number of classes in local
clients or more clients, the gap between FedEC and other methods are more
significant. This shows that our model can effectively tackle non-IID problem
and outperform other methods including traditional meta-learning-based
methods.

###  4.4 Ablation study for the constraint

We conduct ablation study and compare results of L2 regularization. The
results are in Table 2. We compare our method, FedEC with traditional meta-
learning-based FL models (called FedEC-wo).

Methods with elastic regularization improve the acc for about 1-2 percents on
CIFAR10 and CIFAR100 and 24 percents on FEMINST, but L2 regularization method
has lower metrics in three datasets. The more significant gap happens on
FEMINST is reasonable, as FEMINST, with more complicated images and smaller
smaller data size, corresponds to a more challenging setting. The possible
reason for failure of L2 is that it poses a more rigid regularization compared
to our method. The goal of regularization is to let current model adapt faster
to equivalent optimum as historical model. The L2 imposes hard constraint on
the parameter by only allowing parameter searching within a hypersphere with
historical model as the center. This restricts the expressiveness of the model
severely and worsen model performance. While elastic constraint allows any
parameter searching direction as long as current model tracks the historical
prediction distribution, which allows the model to search for better
solutions.

Figure 2: The t-SNE visualization for embedding changes over training rounds:
(a) and (c) corresponds to FedEC-wo over 30 rounds and (b) and (d) corresponds
to FedEC-wo over 100 rounds.

###  4.5 Case Study

We compare the change of local data embedding of FedEC-wo and FedEC in
different communication rounds. We record multiple clients’ average local
feature representation after update at each communication round, and display
their changes in consecutive multiple rounds using t-SNE Van der Maaten and
Hinton (2008). As shown in Figure 2 (a) and (b) represent changes of
embeddings of FedEC-wo and FedEC in the first 30 rounds. (c) and (d) show
their pattern change in the first 100 rounds. Each row shows the change of one
client in different rounds and different methods. Each dot represents the
reduced dimension feature representation in one round.

From the comparison of the pair of (a) and (b) and the pair of (c) and (d), in
the same iteration round, FedEC-wo has large change in the feature, sparse
distribution of the reduced dimension feature representation and cannot
effective converge due to model update fluctuation. FedEC can effectively
solve this problem. Its features has small changes and has centralized
distribution.

We compare (a) and (c) and find that FedEC-wo has multiple clusters. Its
models jump between different clusters in the update process. This phenomenon
validates our hypothesis in the introduction. Local models have multiple
convergence points. Due to the randomness of sampling, the initialization of
model parameters are much different from last round sampling local update. The
update in different round converges to different optima. The inconsistency of
optimization direction cause fluctuation in the local model update. We can see
from (b) and (d), FedEC can effectively solve this problem. As the
communication goes on, FedEC has gradual change apart from isolated points.
There is only one cluster. FedEC effective alleviate model fluctuation.

##  5 CONCLUSION

In this work, we propose FedEC to deal with optimization fluctuation in
traditional meta-learning-based FL models. FedEC incorporates Reptile in FL
and learns from past models. It is immune to the discrepancy of goals in the
same client in different sampling rounds. We compare our method with baseline
methods in three public datasets under different heterogeneity conditions. The
results show that our method effectively alleviates model fluctuation and
improve personalization ability for client models. In the future, we will
explore data heterogeneity in FL and expand FedEC to other personalized FL
methods.

## References

  * Acar et al. [2021] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama.  Federated learning based on dynamic regularization.  arXiv preprint arXiv:2111.04263, 2021.
  * Antoniou et al. [2018] Antreas Antoniou, Harrison Edwards, and Amos Storkey.  How to train your maml.  Learning, 2018.
  * Arivazhagan et al. [2019] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary.  Federated learning with personalization layers.  arXiv preprint arXiv:1912.00818, 2019.
  * Caldas et al. [2018] Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H. Brendan McMahan, Virginia Smith, and Ameet Talwalkar.  Leaf: A benchmark for federated settings.  arXiv: Learning, 2018.
  * Chen et al. [2016] Yutian Chen, Matthew W. Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Timothy P. Lillicrap, Matthew Botvinick, and Nando de Freitas.  Learning to learn without gradient descent by gradient descent.  international conference on machine learning, 2016.
  * Chen et al. [2018] Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiuqiang He.  Federated meta-learning with fast convergence and efficient communication.  arXiv preprint arXiv:1802.07876, 2018.
  * Chen et al. [2021] Jiangui Chen, Ruqing Zhang, Jiafeng Guo, Yixing Fan, and Xueqi Cheng.  Fedmatch: Federated learning over heterogeneous question answering data.  arXiv: Information Retrieval, 2021.
  * Collins et al. [2021] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai.  Exploiting shared representations for personalized federated learning.  international conference on machine learning, 2021.
  * Deng et al. [2020] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi.  Adaptive personalized federated learning.  arXiv preprint arXiv:2003.13461, 2020.
  * Dinh et al. [2020] Canh T. Dinh, Nguyen H. Tran, and Josh Nguyen.  Personalized federated learning with moreau envelopes.  neural information processing systems, 2020.
  * Fallah et al. [2020] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar.  Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach.  Advances in Neural Information Processing Systems, 33:3557–3568, 2020.
  * Fang et al. [2022] Gongfan Fang, Kanya Mo, Xinchao Wang, Jie Song, Shitao Bei, Haofei Zhang, and Mingli Song.  Up to 100x faster data-free knowledge distillation.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6597–6604, 2022.
  * Finn et al. [2017] Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-agnostic meta-learning for fast adaptation of deep networks.  international conference on machine learning, 2017.
  * Finn et al. [2018] Chelsea Finn, Kelvin Xu, and Sergey Levine.  Probabilistic model-agnostic meta-learning.  neural information processing systems, 2018.
  * Ghosh et al. [2020] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran.  An efficient framework for clustered federated learning.  neural information processing systems, 2020.
  * Gou et al. [2021] Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao.  Knowledge distillation: A survey.  International Journal of Computer Vision, 129(6):1789–1819, 2021\.
  * Hanzely and Richtárik [2020] Filip Hanzely and Peter Richtárik.  Federated learning of a mixture of global and local models.  arXiv preprint arXiv:2002.05516, 2020.
  * Hanzely and Richtárik [2021] Filip Hanzely and Peter Richtárik.  Federated learning of a mixture of global and local models.  arXiv: Learning, 2021.
  * Hospedales et al. [2020] Timothy M. Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey.  Meta-learning in neural networks: A survey.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020\.
  * Huang et al. [2020] Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang.  Personalized cross-silo federated learning on non-iid data.  national conference on artificial intelligence, 2020.
  * Karimireddy et al. [2020] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh.  Scaffold: Stochastic controlled averaging for federated learning.  international conference on machine learning, 2020.
  * Kirkpatrick et al. [2017] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al.  Overcoming catastrophic forgetting in neural networks.  Proceedings of the national academy of sciences, 114(13):3521–3526, 2017.
  * Kulkarni et al. [2020] Viraj Kulkarni, Milind Kulkarni, and Aniruddha Pant.  Survey of personalization techniques for federated learning.  In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4), pages 794–797. IEEE, 2020.
  * Li et al. [2017] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales.  Learning to generalize: Meta-learning for domain generalization.  national conference on artificial intelligence, 2017.
  * Li et al. [2018] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith.  Federated optimization in heterogeneous networks.  Proceedings of Machine Learning and Systems, 2018.
  * Liang et al. [2020] Paul Pu Liang, Terrance Liu, Liu Ziyin, Ruslan Salakhutdinov, and Louis-Philippe Morency.  Think locally, act globally: Federated learning with local and global representations.  arXiv: Learning, 2020.
  * Marfoq et al. [2021] Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal.  Federated multi-task learning under a mixture of distributions.  neural information processing systems, 2021.
  * McMahan et al. [2017] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In Artificial intelligence and statistics, pages 1273–1282. PMLR, 2017.
  * Nichol et al. [2018] Alex Nichol, Joshua Achiam, and John Schulman.  On first-order meta-learning algorithms.  arXiv: Learning, 2018.
  * Oh et al. [2020] Jaehoon Oh, Hyungjun Yoo, ChangHwan Kim, and Se-Young Yun.  Boil: Towards representation change for few-shot learning.  Learning, 2020.
  * Park et al. [2021] Younghyun Park, Dong-Jun Han, Doyeon Kim, Jun Seo, and Jaekyun Moon.  Few-round learning for federated learning.  neural information processing systems, 2021.
  * Raghu et al. [2020] Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals.  Rapid learning or feature reuse? towards understanding the effectiveness of maml.  Learning, 2020.
  * Rajeswaran et al. [2019] Aravind Rajeswaran, Chelsea Finn, Sham M. Kakade, and Sergey Levine.  Meta-learning with implicit gradients.  neural information processing systems, 2019.
  * Rusu et al. [2018] Andrei Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osindero, and Raia Hadsell.  Meta-learning with latent embedding optimization.  Learning, 2018.
  * Sattler et al. [2019] Felix Sattler, Klaus-Robert Müller, and Wojciech Samek.  Clustered federated learning: Model-agnostic distributed multi-task optimization under privacy constraints.  IEEE Transactions on Neural Networks, 2019.
  * Shamsian et al. [2021] Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik.  Personalized federated learning using hypernetworks.  international conference on machine learning, 2021.
  * Shin et al. [2021] Jaewoong Shin, Hae Beom Lee, Boqing Gong, and Sung Ju Hwang.  Large-scale meta-learning with continual trajectory shifting.  international conference on machine learning, 2021.
  * Smith et al. [2017] Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar.  Federated multi-task learning.  Advances in neural information processing systems, 30, 2017.
  * Sun et al. [2021] Benyuan Sun, Hongxing Huo, Yi Yang, and Bo Bai.  Partialfed: Cross-domain personalized federated learning via partial initialization.  neural information processing systems, 2021.
  * Szegedy et al. [2016] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna.  Rethinking the inception architecture for computer vision.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2818–2826, 2016.
  * Van der Maaten and Hinton [2008] Laurens Van der Maaten and Geoffrey Hinton.  Visualizing data using t-sne.  Journal of machine learning research, 9(11), 2008.
  * Vanschoren [2019] Joaquin Vanschoren.  Meta-learning.  In Automated Machine Learning, pages 35–61. Springer, Cham, 2019\.
  * Yang et al. [2019] Qiang Yang, Yang Liu, Yong Cheng, Yan Kang, Tianjian Chen, and Han Yu.  Federated learning.  Synthesis Lectures on Artificial Intelligence and Machine Learning, 13(3):1–207, 2019.
  * Ye and Chao [2021] Han-Jia Ye and Wei-Lun Chao.  How to train your maml to excel in few-shot classification.  arXiv: Learning, 2021.
  * Yoon et al. [2018] Jaesik Yoon, Taesup Kim, Ousmane Amadou Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn.  Bayesian model-agnostic meta-learning.  neural information processing systems, 2018.
  * Yuan et al. [2020] Li Yuan, Francis EH Tay, Guilin Li, Tao Wang, and Jiashi Feng.  Revisiting knowledge distillation via label smoothing regularization.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3903–3911, 2020.
  * Zhang et al. [2021] Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang, Wenchao Xu, and Feijie Wu.  Parameterized knowledge transfer for personalized federated learning.  neural information processing systems, 2021.
  * Zhao et al. [2018] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra.  Federated learning with non-iid data.  arXiv preprint arXiv:1806.00582, 2018.