# Transformer Mechanisms Mimic Frontostriatal Gating Operations When Trained
on Human Working Memory Tasks

Aaron Traylor,1 Jack Merullo,1 Michael J. Frank,2 & Ellie Pavlick1
1 Department of Computer Science
2 Department of Cognitive, Linguistic & Psychological Sciences
Brown University, Providence, Rhode Island, USA
{aaron_traylor, jack_merullo, michael_frank<EMAIL_ADDRESS>
###### Abstract

Models based on the Transformer neural network architecture have seen success
on a wide variety of tasks that appear to require complex “cognitive
branching”– or the ability to maintain pursuit of one goal while accomplishing
others. In cognitive neuroscience, success on such tasks is thought to rely on
sophisticated frontostriatal mechanisms for selective gating, which enable
role-addressable updating– and later readout– of information to and from
distinct “addresses” of memory, in the form of clusters of neurons. However,
Transformer models have no such mechanisms intentionally built-in. It is thus
an open question how Transformers solve such tasks, and whether the mechanisms
that emerge to help them to do so bear any resemblance to the gating
mechanisms in the human brain. In this work, we analyze the mechanisms that
emerge within a vanilla attention-only Transformer trained on a simple
sequence modeling task inspired by a task explicitly designed to study working
memory gating in computational cognitive neuroscience. We find that, as a
result of training, the self-attention mechanism within the Transformer
specializes in a way that mirrors the input and output gating mechanisms which
were explicitly incorporated into earlier, more biologically-inspired
architectures. These results suggest opportunities for future research on
computational similarities between modern AI architectures and models of the
human brain.

Keywords: transformers; neural networks; working memory; computational
neuroscience; gating; computational cognitive science; mechanistic
interpretability

## Introduction

Computational models based on the Transformer architecture Vaswani . (2017)
have seen success on a wide variety of tasks that appear to require complex
“cognitive branching”: the ability to maintain pursuit of one over-arching
goal while performing other subtasks along the way. For example, Transformer-
based large language models (LLMs) have demonstrated impressive abilities in
not just language Brown . (2020), but planning Huang . (2022), navigation Du .
(2021), and problem solving Lewkowycz . (2022).

In humans, performance on such tasks is known to depend on a neural mechanism
for gating Frank Badre (2012); Badre Frank (2012); Chatham . (2014); Rac-
Lubashevsky Kessler (2016); Rac-Lubashevsky Frank (2021), which controls
whether new information is maintained in working memory or not, the address in
memory where it is stored, and the address from which stored information is
recalled in response to a task. Typical Transformer models have no specialized
architecture for working memory, in spite of their ability to succeed at tasks
which appear to require it. Although some Transformer models have additional
built-in structure for memory Dai . (2019); Burtsev . (2020), recurrence Dai .
(2019), or hierarchy YS. Wang . (2019), vanilla Transformer models without any
such inductive biases remain the dominant architecture for the modern AI
systems which are especially visible in their success on complex tasks Brown .
(2020); Touvron . (2023). This raises the question: in solving such tasks,
does a mechanism for selective input and output gating emerge within the
vanilla Transformer?

Transformers are good candidates for learning gating behavior because of
inductive biases within the self-attention mechanism, i.e., the Transformer’s
defining architectural component. Within self-attention, numerous “attention
heads” construct contextual representations for each item in their input
sequence through a learned weighted combination of the previous items in the
sequence. Attention heads could in principle learn gating behavior by marking
sequence elements with a key (input gating) and reading out those values later
by querying for those keys (output gating). Moreover, the attention mechanism
in Transformers is decomposed in a way that enables it to readily
differentiate “reading” and “writing” operations. This behavior is analogous
functionally to the neurobiological roles of corticostriatal circuits in
humans and other animals, in which isolated clusters of prefrontal neurons
represent distinct “addresses” in memory that can be updated or read out from
via selective gating actions triggered by basal ganglia and thalamus O’Reilly
Frank (2006); Frank Badre (2012); Calderon . (2022). In computational
neuroscience models of this process, the prefrontal clusters (or “stripes”)
can also serve as latent abstract “roles” that condition how to interpret
content within them, affording functions such as variable binding,
indirection, and hierarchical generalization to new situations O’Reilly Frank
(2006); Collins Frank (2013); Kriete . (2013); Bhandari Badre (2018). Thus,
Transformers may use their attention heads to learn a gating strategy that
mimics certain functions of the brain.

In this work, we train vanilla Transformer models with self-attention on a
working memory task paradigm that was specifically designed to evaluate models
of selective gating and working memory in computational neuroscience O’Reilly
Frank (2006); Rac-Lubashevsky Frank (2021). We use recent techniques from
mechanistic interpretability Olah (2022); Nanda Bloom (2022) to expose the
mechanism that the Transformer uses in order to perform the task. We find
that, as a result of training, the self-attention mechanism specializes in a
way that resembles existing models of input-output gating. Specifically, we
find that the trained model uses the key vectors within the attention
mechanism to control input gating, i.e., determining which elements in the
input to consider vs. ignore, as well as controlling how the information is
stored– in other words, assigning it a role such that the model can access it
later. The model uses the query vectors to control output gating, i.e.,
determining which information is accessed in order to complete a task. Our
findings highlight the importance of considering the emergent mechanisms that
result from training in addition to the innate architectural mechanisms when
drawing comparisons between AI systems and human cognitive processes, and
opens the door for future analysis and work which can enable more principled
studies of the similarities and differences between human vs. machine
cognition.

## Background

### Gating Mechanisms

Working memory in human brains is known to make use of a gating mechanism,
which processes and stores information roughly analogously to gates being
opened and closed Rac-Lubashevsky Kessler (2016); Rac-Lubashevsky Frank
(2021); Bhandari Badre (2018). The input gate controls which information is
stored or not stored in memory, and if stored, into which “address”. The
output gate controls which content within working memory is accessed in order
to produce a response or to make subsequent gating operations. Both input and
output gating policies are also dependent on the learned task-dependent
context (i.e. role) of the information to be stored and accessed in working
memory.

In cognitive neuroscience, a variety of tasks are used to study the capacity
of working memory. In this work, we focus on a variant of the “reference-back
2” task Rac-Lubashevsky Frank (2021), a human paradigm meant to mimic a task
designed to showcase the need for selective gating of independent contents of
information in frontostriatal neural networks O’Reilly Frank (2006). In the
reference-back paradigm, symbols such as letters or numbers are viewed one at
a time, and the subject must determine whether the current symbol is the same
or different as that stored in memory for a given role (letter or number).
They also are given a cue to indicate whether to update the current symbol as
the “reference” to be compared on subsequent trials of the same role, or if
instead they should continue to maintain the previous reference. Thus this
task requires selective updating and accessing of information in a role-
addressable manner.

### Transformers

Transformers are powerful language models which create contextualized
representations of sequences of words, where they learn to to predict the next
token one at a time using an “attention” mechanism Bahdanau . (2014) to scan
the previously seen tokens for relevant information. These models are able to
learn and represent any arbitrary sequence modelling task.

For a given prediction, Transformer attention generates three separate vectors
at each position in the sequence: a query, key, and value ($q$, $k$, $v$). The
query vector scans the previous context (including the current token) for
relevant keys, and calculates how much the current prediction should “attend”
to those positions. Then, the value vectors at those positions are multiplied
by the corresponding weights, summed up, and added to the next representation:
for token $i$ at layer $j$, the contextual representation is
$\sum_{k}q_{i}^{j}\cdot k_{k}^{j}*v_{k}^{j}$. Thus, the next token prediction
includes information from earlier in the sequence by combining the value
vectors from previous tokens. In other words, Transformer attention can be
viewed as a read/write mechanism: for a given token, the queries and keys
dictate which tokens to read from, the values are the content that is read
proportional to the attention calculated by the keys and queries, and the
summed content is written to a new representation at the given token. As we
shall see below, the comparison to role-addressable input and output gating
operations is evident, whereby the key vectors form addresses analogous to the
PFC “stripes”, with key construction determining input gating, and the query
vectors control which of them are accessed, with query construction
determining output gating.

#### Limitations as Cognitive Models:

Transformers have properties which make them obviously bad models of human
sequence processing. In particular, because Transformers can attend to any
part of the sequence when creating a representation, they are not limited by
memory representation constraints. Transformers could thus solve tasks that
would push the limits of human working memory, but it should be noted they
accordingly require large amounts of training data. The question addressed in
this work is orthogonal to these limitations. That is, we focus specifically
on if and how Transformers can learn to implement an efficient gating
mechanism to solve tasks with human working memory demands.

Figure 1: Graphical diagram of the path-patching process. Attention heads are
represented as circles (layer,head index), and contextual representations of
each token (as well as the next token prediction) are represented as
rectangles.

### Mechanistic Interpretability

We use a set of recently introduced analysis tools Elhage . (2021) which
enable us to uncover specific mechanisms defined in terms of model weights
within the Transformer. Specifically, we use path-patching K. Wang . (2022);
Goldowsky-Dill . (2023), a generalization of causal mediation analysis Pearl
(2001) that allows us to determine which components of a neural model (e.g.,
attention heads) work together in order to produce observed behavior on a
task. The discovered components are referred to as a circuit Räuker . (2023).

Path-patching involves making a incisive edit to the representations of a
trained model and observing how the model’s behavior is affected (see Fig. 1).
Path-patching typically requires a minimal pair of examples: the “clean”
example and the “corrupted” example, in which one token from the clean example
is changed, as well as the correct label. Given representations from the model
for both the clean and the corrupted examples (the blue and orange components
in the figure), we can chose a specific component anywhere in the model
(referred to as the “sender”), and insert the corrupted representation at that
component into the clean representation. We then use the model to recalculate
the representations up until another component of the model (the “receiver”),
thus “patching” the path. In the figure, we send from layer 0, head 0 to layer
1, both heads 0 and 1. All clean representations that are not along this path
are not modified and are unaffected by the patch. The model then recomputes
all representations after the receiver (the “patched” representations), and
arrives at a new prediction. If the model output matches the corrupt
prediction rather than the clean one, that prediction is causally dependent on
the path from sender to receiver. See K. Wang . (2022) and Goldowsky-Dill .
(2023) for a more comprehensive review of path-patching methods.

Figure 2: Above: example of textual reference-back task as model input. Below:
step-by-step task process; models do not view task-internal grey words.
“Update Instruction” executes after “Answer” despite appearing earlier
sequentially.

## Task

We create a modified text-based version of the reference-back 2 task Rac-
Lubashevsky Frank (2021) designed to tax selective WM gating O’Reilly Frank
(2006).The textual reference-back task requires making same/different
judgments between incoming symbols assigned to a particular “register” in
memory, with respect to those seen previously and linked to those same
registers. Like the original tasks, the textual reference-back task is
sequential, and requires the maintenance and independent updating of two
memory registers, each containing one of $S$ arbitrary symbols at a time. The
contents of the registers are updated over the course of the task. At the
beginning of each sequence, each register is initialized individually to one
$s\in S$ (the pool of symbols is shared between registers, which was shown to
more substantively tax gating mechanisms in O’Reilly Frank (2006)). Each
sequence is composed of $L$ tuples, each containing register address
$\text{Reg}_{i}$, symbol $\text{Sym}_{i}$, same/different label
$\text{Ans}_{i}$, and update instruction $\text{Ins}_{i}$. For a tuple $i\in
L$, the answer $\text{Ans}_{i}$ is a binary value that is either same if
symbol $\text{Sym}_{i}$ is currently stored in the register with address
$\text{Reg}_{i}$, or different otherwise. The update instruction
$\text{Ins}_{i}$ also takes one of two values, evenly distributed: if ignore,
then there is no effect further on in the sequence. If store, then from that
point on in the sequence, $\text{Sym}_{i}$ is stored in the register with
address $\text{Reg}_{i}$ until otherwise updated. An example is shown in Fig.
2.

We implement each reference-back task example in our data as a single
sequence, and measure models’ ability to predict same versus different for
each $\text{Ans}_{i}$. Each sequence has 10 same/different answers, and we
generate 100,000 train, 1,000 dev, and 1,000 held-out test sequences.

The class balance of same to different answer labels in the train/test
datasets is roughly 1:2, making a “maximum class” heuristic solution 0.66
accuracy, 0.33 precision, and 0.5 recall. We test several other heuristics,
the strongest of which is predicting same if another tuple including Store and
the target register and target symbol exists in the sequence, which scores
0.80 accuracy, 0.82 precision, and 0.85 recall.

## Model

We use vanilla Transformers in order to facilitate interpretability, as done
in prior work that analyzes emergent mechanisms Elhage . (2022). Our models
contain two decoder-only layers, each with only two heads (four in total), and
no multilayer perceptrons or layer normalization, followed by a linear
“unembed” layer to project the output of the last decoder into the space of
the entire vocabulary at each timestep 111In practice, only ‘same’ and
‘different’ are ever predicted.. Our network uses absolute positional
embeddings Vaswani . (2017). The vocabulary contains all possible tokens,
represented individually with embedding size $E$. Models are trained to
predict the next token with the language modelling objective, meaning if the
model is predicting $\text{Ans}_{c}$, it will have access to all
($\text{Ins}_{i}$, $\text{Reg}_{i}$, $\text{Sym}_{i}$, $\text{Ans}_{i}$)
tuples where $i<c$, as well as $\text{Ins}_{c}$, $\text{Reg}_{c}$, and
$\text{Sym}_{c}$. However, the models only receive loss at positions where a
same/different token must be predicted. Furthermore, each layer gets a causal
attention mask– when constructing each token representation, it cannot look
ahead at tokens further down the sequence.

The models are trained over 60 epochs of the 100k training data points,
learning from 6 million examples in total. Models are evaluated on their
accuracy (whether the correct $\text{Ans}_{i}$ is predicted for each tuple
$i$), measured in precision and recall, as well as the same versus different
token logit difference.

## Experiments

First, we select and analyze a single Transformer model which succeeds on the
task, and upon investigation of its weights find that it learns a mechanism
for input/output gating. Second, when we conduct a search over more trained
models, we find that model performance correlates with markers of learning a
gating policy, analogous to findings in the frontostriatal neural networks
Frank Badre (2012).

Figure 3: Model behavior when predicting same/different (token 15) is shown.
We measure attention visualized as a shade of purple, with deeper shade
corresponding to higher attention to that token. We create “corrupted” minimal
pairs in which changing a token (light blue) either changes the correct label
at index 15 (examples b, c, e) or does not (d, f). We make small path-patching
edits with the minimal pair to targeted network components (layer 1 keys for
b, c, d, f; queries for e,f). In other words, we replace specific components
(denoted with red text) with their corresponding representation from the
“corrupted” sequence, but hold all other representations constant, and run the
model and get a new same/different prediction. In all test examples, making
the small patch successfully results in the model’s prediction changing to
align with the “corrupted” example.

We first establish that a Transformer model is able to succeed on the
reference-back task. We perform a small hyperparameter search and select a
model that reaches 100% accuracy on the held-out test data for further
analysis. We determine the circuit that the model uses to solve the textual
reference-back task through an array of path-patching experiments with a
simplistic minimal pairs paradigm. Our “sender” within path-patching is always
both attention heads at layer 0, and our “receiver” is always both attention
heads at layer 1.

At layer 0, the model learns to condense the task-critical information from
each tuple into one embedding, at the position for
$\text{Sym}_{i}$222Redundantly, the model does the same at the position for
$\text{Ans}_{i}$. Through additional experimentation, we determine that this
is a quirk of Transformer learning, and does not impact our analysis.. At this
layer, the model pays 85.8% of total attention to the task-critical
information to that tuple, and just 14.2% of attention to other tuples.

At layer 1, the attention heads learn to attend to the $\text{Sym}_{i}$ key
vector representing the tuple where information was last stored in the target
register. The heads pay 70.2% of total attention to this tuple (the “stored”
tuple), and only 29.8% of attention to all other tokens. This behavior is tied
to the target register matching the register in the stored tuple, which is
analogous to gating of the relevant role-addressable PFC stripe. We focus our
analysis on the Layer 1 representations which exhibit this learned gating
policy, shown in Fig. 3.

### Input Gating through Key Vectors

Input gates in working memory control what incoming information is remembered
and “role-addressable”– i.e., stored in memory in such a way that it is able
to be freely accessed later when it is needed for task completion. In a
Transformer, the key vectors serve the role of addresses (analogous to PFC
“stripes”), which are retrieved based on their match to a query vector from
the current or later timesteps during self-attention. Thus the input gating in
Transformers is controlled through key construction; the composition of the
output of the Layer 0 attention controls which content is stored in the key
vector at Layer 1 for later use. At a later timestep, a query vector will
address the information in the key vector.

In our model analysis, we find that key composition at the $\text{Sym}_{i}$
position (positions 2, 6, 10, and 14 in Fig. 3) roughly represents each tuple.
A query’s ability to address this position depends on whether the represented
tuple contains a Store or an Ignore. Key vectors representing an Ignore tuple
receive very little attention (0.4% of layer 1 attention averaged over test
set), whereas those representing a Store tuple receive the bulk of the
attention (86.8%). We determine this effect causally with path-patching (Fig.
1). First, we create clean sequences sampled from our test set, and then
corrupt these sequences by switching a Store within tuple $i$ to an Ignore. We
then path-patch only the key vectors of $i$. We expect, if the key controls
input gating, that patching these key vectors should “block” attention to all
of tuple $i$. An example attention pattern is in Fig. 3, examples a and b. We
find that the model’s attention shifts away from the tuple accordingly in 100%
of patched instances. The presence of an Ignore or a Store within a tuple
controls whether the key construction acts as an open input gate or a closed
input gate.

Key construction also depends on the role of the represented content; within
our task, that means whether $\text{Reg}_{i}$ is Reg0 or Reg1. When making a
same/different prediction, key vectors representing a tuple that matches the
target register receive most of the model’s attention (92.5% of total
attention), while those that do not match are not attended to (3.3% of total
attention). Similarly to input gating, we use path-patching to determine that
key construction encodes roles. This time, given a target tuple $i$ with a
target register, we corrupt the register of the stored tuple, changing it from
Reg0 to Reg1 or vice versa; to predict the answer for $i$, the model must
attend to an earlier tuple with the target register. An example is Fig. 3, row
c. The model’s attention shifts away accordingly across every example in the
test set. Note that the stored tuple must be modified; if the same corruption
is made earlier (as in row d), attention does not shift. This behavior shows
that the gating within self-attention is role-addressable; the registers
within the task function as roles, and are embedded within the key vectors as
part of the representation.

### Output Gating through Query Vectors

Within working memory, output gates control which addressable information is
accessed in order to complete a task. Given that key vectors serve the role of
addresses, query vectors in turn control which key vectors are accessed,
through the final Q*K dot product in attention. Query construction thus
performs the role of output gating within Transformers.

The query composition controls which addressable $\text{Sym}_{i}$
representations are attended to based on the identity of the target register;
changing Reg1 to Reg0 controls which role-addressable content is accessed. We
determine this through a final set of path-patching experiments, an example of
which is shown in Fig. 3, row e. Rather than editing the register in the
stored tuple as done in row c, we corrupt the target register itself; this
means that the query must now find representations corresponding to the other
tuple. In row e, the model finds $Sym_{5}$, and predicts different; and in row
f, we patch in $Sym_{4}$ at index 2, and the model predicts same. Upon
inspection, the query is constructed to correspond with key vectors that
represent tuples which also contain the target register. When we edit the
target register in minimal pair experiments, we observe that the attention
shifts from the original stored tuple (74.1% of attention) to the stored tuple
that matches the edited register, and successfully makes an updated
same/different comparison to the symbol in the edited register in every
instance.

Editing aspects of the target tuple other than target register has minimal
effect on the query construction behavior. No edits to the query cause the
model to attend to a Ignore tuple, further evidencing of output gating
behavior– only content that has been made “addressable” can be accessed for a
response. Furthermore, we find that the target instruction and symbol do not
factor into the query composition– changing them through path-patching to the
query does not affect attention. This is notable because the model could
employ other strategies for determining which tuples are eligible to be the
stored tuple; e.g. attending to all symbols to match if any of them are the
same as the target symbol.

### Successful Task Performance is Related to Discovering Gating Policies

To further identify how readily Transformer models learn a gating policy, and
how useful such a policy is to succeed on the task, we train new models with
the same hyperparameters across many different random seeds, and measure their
performance on the target task as well as on markers of the gating policy. We
train 20 new models on the same textual reference-back data, each with a
different random initialization, and measure both training loss and test set
accuracy. 5 of the models succeed 100% of the time, and the other 15 models
succeed between 94%-99.99% of the time, with a mean of 97.72% and a standard
deviation of 2.03. The models are trained on the same amount of data (6
million examples).

We observed in the prior sections that the trained Transformer model uses its
key composition to control input gating and its query composition to control
output gating and role addressability. To identify whether the new models
learn to gate similarly, we evaluate the key and query composition of all 20
new models by making minimal pair path-patching edits for every test example,
where the answer changes from same to different or vice versa, similarly to
Figure 3.

To evaluate the ability to open and close input gates, we corrupt the stored
tuple’s Store to Ignore, and path-patch only to the stored tuple’s key
vectors. To evaluate the role addressability of the content during output
gating, we corrupt the target register (changing Reg0 to Reg1 or vice versa),
and path-patch only to the query vector for making the same/different
judgment. This is a more challenging task than patching to all of the keys or
queries respectively; a model will only succeed on these subtasks if it
implements a gating policy with the same markers that the model analyzed in
earlier sections does. A model makes the patched prediction successfully if
its prediction matches the corrupted sequence and not the original sequence–
i.e. the targeted path-patch was sufficient to change its same/different
prediction.

Figure 4: Model performance over training on patching subtasks. Each row
corresponds to an individual model’s training loss (solid line) and subtask
accuracy (dashed line) over time; blue and orange lines respectively
correspond to models which reach 100% test accuracy and to those that do not.

We visualize the 5 runs that reach 100% accuracy on the test data in Figure 4,
as well as 5 randomly selected runs that do not reach 100% accuracy on the
test data, and plot their training loss versus their accuracy (between 0 and
100) on the two patched subtasks over the course of training.

Two trends become apparent from the data: first, models which score a perfect
test accuracy appear to succeed at the subtasks more readily than models which
do not. Of the former, 3 of 5 models reach 100% accuracy on both subtasks
readily, plateauing less than halfway through training duration. However,
examples from the latter category of models do not reach such immediate
success at the patched subtasks (including the 10 not pictured in this graph);
in fact, many categorically fail, scoring as low as 49% accuracy. These
results do not indicate that this class of models’ representations are useless
for the task– they all score between 94% and 99.99%, well above heuristic
performance. Failing to succeed at the targeted patching subtasks reveals that
these models may implement some other strategy that is not gating, and may be
brittle or heuristic in some way. We take these results as evidence that
learning a robust policy for gating correlates with model performance at the
textual reference-back task.

The second emergent trend is that many models across both classes have a sharp
decline in training loss, which correlates with a similarly steep increase in
accuracy on both subtasks. Sudden jumps in performance is a noted phenomenon
that has been observed in other Transformer models in cases of e.g., grokking
Power . (2022). We interpret this phase transition as suddenly learning a
gating mechanism. Models that do not exhibit phase transitions to the same
degree take longer to fit the task, and do not reach high subtask accuracy.
The cause of phase transitions and what is learned during this process is left
for future work.

## Summary and Discussion

In this work, we investigate Transformer models for emergence of a learned
gating mechanism; a network component performing role-addressable gating,
similar to that in working memory of humans. We observe that the model readily
learns a gating policy, and upon training more models find that task
performance is correlated with gating ability. Our results show how competence
at cognitive branching tasks can emerge in Transformers, and suggest that
integrating Transformer components may improve existing computational
neuroscience models of working memory.

The Transformer models which were trained on the textual reference-back task
are capable of making use of key composition for input gating and query
composition for output gating. We find that making precise corruptions to
specific architectural elements of the network causes the model’s prediction
to change from same to different or vice versa, indicating that those
components are causally responsible for the gating mechanism. The
architectural biases of attention within the vanilla Transformer model lend
themselves well to representing role-addressable content, as the learnable
nature of keys, queries, and values allows the model to learn to create
internal representations in a manner which allows it to represent roles and
addresses, mimicking the variable binding and input / output gating mechanisms
in biological neural networks O’Reilly Frank (2006); Frank Badre (2012);
Collins Frank (2013); Kriete . (2013).

In this work, we focused on characterizing the mechanism that the Transformer
model learns as a result of training on cognitive branching tasks, and did not
evaluate the robustness of the mechanism. The textual reference-back task is
similar in nature to the FFLM task Liu . (2023), comprised of 1 register, 2
symbols (0 or 1), and long sequences. Liu . found that similar Transformer
models were able to succeed on FFLM data, but struggled to generalize outside
of their training distribution. We leave experiments characterizing
generalization of Transformer mechanism behavior on data with the reference-
back paradigm– e.g. more than 2 registers, distinct sets of symbols, novel
symbols introduced at test time– for future work.

When we trained more models on the task, we found that the models which
perform best on the task correlate with the markers of gating we observed in
our circuit analysis, and that the learning trajectory shows a steep decrease
in training loss and a steep rise in patched subtask accuracy simultaneously,
suggesting that the model has learned a component of the gating policy at that
time. Both findings are analogous to those of frank2012mechanisms, in which
they find that networks which learned a hierarchical gating policy performed
better at a hierarchical learning task, and humans that learn this policy also
show a sharp decrease in loss when they discover it.

Ultimately, finding connections between emergent behavior of Transformer
models and human working memory serves to benefit both computational cognitive
neuroscience and artificial intelligence. Although Transformer models
themselves are limited in their biological plausibility, in this setting they
learned behavior mimicking the functionality of working memory, and their
application within computational models of the brain should be further
explored. From the perspective of artificial intelligence, understanding the
strengths and limitations of Transformer models on cognitive branching tasks
may inform model analysis across the many diverse settings in which these
models are applied.

## References

  * Badre Frank (2012) badre2012mechanismsBadre, D. Frank, MJ.  2012\.  Mechanisms of hierarchical reinforcement learning in cortico–striatal circuits 2: Evidence from fMRI Mechanisms of hierarchical reinforcement learning in cortico–striatal circuits 2: Evidence from fmri.  Cerebral cortex223527–536.
  * Bahdanau . (2014) bahdanau2014neuralBahdanau, D., Cho, K. Bengio, Y.  2014\.  Neural machine translation by jointly learning to align and translate Neural machine translation by jointly learning to align and translate.  arXiv preprint arXiv:1409.0473.
  * Bhandari Badre (2018) bhandari2018learningBhandari, A. Badre, D.  2018\.  Learning and transfer of working memory gating policies Learning and transfer of working memory gating policies.  Cognition17289–100.  https://www.sciencedirect.com/science/article/pii/S0010027717303037 https://doi.org/10.1016/j.cognition.2017.12.001
  * Brown . (2020) brown2020languageBrown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, JD., Dhariwal, P.others  2020\.  Language models are few-shot learners Language models are few-shot learners.  Advances in neural information processing systems331877–1901.
  * Burtsev . (2020) burtsev2020memoryBurtsev, MS., Kuratov, Y., Peganov, A. Sapunov, GV.  2020\.  Memory transformer Memory transformer.  arXiv preprint arXiv:2006.11527.
  * Calderon . (2022) calderon2022thunderstruckCalderon, CB., Verguts, T. Frank, MJ.  2022\.  Thunderstruck: The ACDC model of flexible sequences and rhythms in recurrent neural circuits Thunderstruck: The acdc model of flexible sequences and rhythms in recurrent neural circuits.  PLoS Computational Biology182e1009854.
  * Chatham . (2014) chatham2014corticostriatalChatham, CH., Frank, MJ. Badre, D.  2014\.  Corticostriatal output gating during selection from working memory Corticostriatal output gating during selection from working memory.  Neuron814930–942.
  * Collins Frank (2013) collins2013cognitiveCollins, AG. Frank, MJ.  2013\.  Cognitive control over learning: creating, clustering, and generalizing task-set structure. Cognitive control over learning: creating, clustering, and generalizing task-set structure.  Psychological review1201190.
  * Dai . (2019) dai2019transformerDai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, QV. Salakhutdinov, R.  2019\.  Transformer-XL: Attentive language models beyond a fixed-length context Transformer-xl: Attentive language models beyond a fixed-length context.  arXiv preprint arXiv:1901.02860.
  * Du . (2021) du2021vtnetDu, H., Yu, X. Zheng, L.  2021\.  VTNet: Visual transformer network for object goal navigation Vtnet: Visual transformer network for object goal navigation.  arXiv preprint arXiv:2105.09447.
  * Elhage . (2022) elhage2022superpositionElhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S.Olah, C.  2022\.  Toy Models of Superposition Toy models of superposition.  Transformer Circuits Thread.  https://transformer-circuits.pub/2022/toy_model/index.html
  * Elhage . (2021) elhage2021mathematicalElhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B.others  2021\.  A mathematical framework for transformer circuits A mathematical framework for transformer circuits.  Transformer Circuits Thread1.
  * Frank Badre (2012) frank2012mechanismsFrank, MJ. Badre, D.  2012\.  Mechanisms of hierarchical reinforcement learning in corticostriatal circuits 1: computational analysis Mechanisms of hierarchical reinforcement learning in corticostriatal circuits 1: computational analysis.  Cerebral cortex223509–526.
  * Goldowsky-Dill . (2023) goldowsky2023localizingGoldowsky-Dill, N., MacLeod, C., Sato, L. Arora, A.  2023\.  Localizing model behavior with path patching Localizing model behavior with path patching.  arXiv preprint arXiv:2304.05969.
  * Huang . (2022) huang2022languageHuang, W., Abbeel, P., Pathak, D. Mordatch, I.  2022\.  Language models as zero-shot planners: Extracting actionable knowledge for embodied agents Language models as zero-shot planners: Extracting actionable knowledge for embodied agents.  International Conference on Machine Learning International conference on machine learning ( 9118–9147).
  * Kriete . (2013) kriete2013indirectionKriete, T., Noelle, DC., Cohen, JD. O’Reilly, RC.  2013\.  Indirection and symbol-like processing in the prefrontal cortex and basal ganglia Indirection and symbol-like processing in the prefrontal cortex and basal ganglia.  Proceedings of the National Academy of Sciences1104116390–16395.
  * Lewkowycz . (2022) lewkowycz2022solvingLewkowycz, A., Andreassen, A., Dohan, D., Dyer, E., Michalewski, H., Ramasesh, V.others  2022\.  Solving quantitative reasoning problems with language models Solving quantitative reasoning problems with language models.  Advances in Neural Information Processing Systems353843–3857.
  * Liu . (2023) liu2023exposingLiu, B., Ash, JT., Goel, S., Krishnamurthy, A. Zhang, C.  2023\.  Exposing Attention Glitches with Flip-Flop Language Modeling Exposing attention glitches with flip-flop language modeling.  arXiv preprint arXiv:2306.00946.
  * Nanda Bloom (2022) nanda2022transformerlensNanda, N. Bloom, J.  2022\.  TransformerLens. Transformerlens.  https://github.com/neelnanda-io/TransformerLens.
  * Olah (2022) olahblogOlah, C.  2022\.  Mechanistic interpretability, variables, and the importance of interpretable bases. Mechanistic interpretability, variables, and the importance of interpretable bases.  https://www.transformer-circuits.pub/2022/mech-interp-essay.
  * O’Reilly Frank (2006) o2006makingO’Reilly, RC. Frank, MJ.  2006\.  Making working memory work: a computational model of learning in the prefrontal cortex and basal ganglia Making working memory work: a computational model of learning in the prefrontal cortex and basal ganglia.  Neural computation182283–328.
  * Pearl (2001) pearl2001Pearl, J.  2001\.  Direct and indirect effects Direct and indirect effects.  Proceedings of the Seventeenth Conference on Uncertainty in Artificial Intelligence Proceedings of the seventeenth conference on uncertainty in artificial intelligence ( 411–420).  San Francisco, CA, USAMorgan Kaufmann Publishers Inc.
  * Power . (2022) power2022grokkingPower, A., Burda, Y., Edwards, H., Babuschkin, I. Misra, V.  2022\.  Grokking: Generalization beyond overfitting on small algorithmic datasets Grokking: Generalization beyond overfitting on small algorithmic datasets.  arXiv preprint arXiv:2201.02177.
  * Rac-Lubashevsky Frank (2021) rac2021analogousRac-Lubashevsky, R. Frank, MJ.  2021\.  Analogous computations in working memory input, output and motor gating: Electrophysiological and computational modeling evidence Analogous computations in working memory input, output and motor gating: Electrophysiological and computational modeling evidence.  PLoS computational biology176e1008971.
  * Rac-Lubashevsky Kessler (2016) rac2016dissociatingRac-Lubashevsky, R. Kessler, Y.  2016\.  Dissociating working memory updating and automatic updating: The reference-back paradigm. Dissociating working memory updating and automatic updating: The reference-back paradigm.  Journal of Experimental Psychology: Learning, Memory, and Cognition426951.
  * Räuker . (2023) rauker2023towardRäuker, T., Ho, A., Casper, S. Hadfield-Menell, D.  2023\.  Toward transparent ai: A survey on interpreting the inner structures of deep neural networks Toward transparent ai: A survey on interpreting the inner structures of deep neural networks.  2023 IEEE Conference on Secure and Trustworthy Machine Learning (SaTML) 2023 ieee conference on secure and trustworthy machine learning (satml) ( 464–483).
  * Touvron . (2023) touvron2023llamaTouvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, MA., Lacroix, T.others  2023\.  Llama: Open and efficient foundation language models Llama: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971.
  * Vaswani . (2017) vaswani2017attentionVaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, AN.Polosukhin, I.  2017\.  Attention is all you need Attention is all you need.  Advances in neural information processing systems30.
  * K. Wang . (2022) wang2022interpretabilityWang, K., Variengien, A., Conmy, A., Shlegeris, B. Steinhardt, J.  2022\.  Interpretability in the wild: a circuit for indirect object identification in gpt-2 small Interpretability in the wild: a circuit for indirect object identification in gpt-2 small.  arXiv preprint arXiv:2211.00593.
  * YS. Wang . (2019) wang2019treeWang, YS., Lee, HY. Chen, YN.  2019\.  Tree transformer: Integrating tree structures into self-attention Tree transformer: Integrating tree structures into self-attention.  arXiv preprint arXiv:1909.06639.


# Flashback: Understanding and Mitigating
Forgetting in Federated Learning

Mohammed Aljahdali  Ahmed M. Abdelmoniem  Marco Canini  Samuel Horváth

###### Abstract

In Federated Learning (FL), forgetting, or the loss of knowledge across
rounds, hampers algorithm convergence, particularly in the presence of severe
data heterogeneity among clients. This study explores the nuances of this
issue, emphasizing the critical role of forgetting in FL’s inefficient
learning within heterogeneous data contexts. Knowledge loss occurs in both
client-local updates and server-side aggregation steps; addressing one without
the other fails to mitigate forgetting. We introduce a metric to measure
forgetting granularly, ensuring distinct recognition amid new knowledge
acquisition. Leveraging these insights, we propose Flashback, an FL algorithm
with a dynamic distillation approach that is used to regularize the local
models, and effectively aggregate their knowledge. Across different
benchmarks, Flashback outperforms other methods, mitigates forgetting, and
achieves faster round-to-target-accuracy, by converging in 6 to 16 rounds.

Machine Learning, ICML

##  1 Introduction

Federated Learning (FL) is a distributed learning paradigm that allows
training over decentralized private data. These datasets belong to different
clients that participate in training a global model. Federated Averaging
(FedAvg) (McMahan et al., 2017) is a prominent training algorithm that uses a
centralized server to orchestrate the process. At every round, the server
samples a proportion of the available clients. Starting form the current
version of the global model, each sampled client performs $E$ epochs of local
training using their private data and sends its updated model to the server.
Then, the server aggregates the models by averaging them to obtain the new
global model. This process is typically repeated for many rounds until a
desired model performance is obtained.

A main challenge in FL is the heterogeneity in distribution between the
private datasets, which are unbalanced and non-IID (Kairouz et al., 2019).
Data heterogeneity causes local model updates to drift – the local optima
might not be consistent with the global optima – and can lead to slow
convergence of the global model – where more rounds of communication and local
computation are needed – or worse, when the desired performance may not be
reached. Addressing data heterogeneity in FL has been the focus of several
prior studies. For instance, FedProx (Li et al., 2020) proposes a proximal
term to limit the distance between the global model and the local model
updates, mitigating the drift in the local updates. MOON (Li et al., 2021b)
mitigates the local drift using a contrastive loss to minimize the distance
between the feature representation of the global model and the local model
updates while maximizing the distance between the current model updates and
the previous model updates. FedDF (Lin et al., 2020) addresses heterogeneity
in local models by using ensemble distillation during the aggregation step
(instead of averaging the model updates).

However, we experimentally observe that under severe data heterogeneity, these
proposals provide little or even no advantage over FedAvg. Figure 1(a)
illustrates the test accuracy of FedAvg and other baselines while training a
DNN over the CIFAR10 dataset (Krizhevsky, 2009) (details in § 5).

This motivates us to better understand how data heterogeneity poses a
challenge for FL and devise a new approach of handling non-IID datasets. We
investigate the evolution of the global model accuracy broken down by its per-
class accuracy. Figure 1(b) shows as a heatmap the per-class accuracy for
FedAvg; each rectangle represents the accuracy of the global model on a class
at a round. Other baseline methods show similar results. Our key observation
is that there is a significant presence of _forgetting_ : i.e., cases where
some knowledge obtained by the global model at round $t$ is dropped at round
$t+1$, causing the accuracy to decline (as shown by the prominent number of
light-shaded rectangles at the right side of darker ones in the figure).

(a) Global model accuracy of FedAvg and baselines.

(b) Per-class accuracy of FedAvg’s global model.

Figure 1: Performance of FedAvg and other baselines over training rounds with
CIFAR10.

A similar phenomenon is known as _catastrophic forgetting_ in Continual
Learning (CL) literature (Parisi et al., 2019). CL addresses the challenge of
sequentially training a model on a series of tasks, denoted as
$\\{T_{1},T_{2},\ldots,T_{n}\\}$, without revisiting data from prior tasks.
Formally, given a model with parameters $\theta$ and task-specific loss
functions $L_{t}(\theta)$ for each task $T_{t}$, the objective in CL is to
update $\theta$ such that performance on the current task is optimized without
significantly degrading the model’s performance on previously learned tasks.
This is non-trivial, as naïve sequential training often leads to catastrophic
forgetting, where knowledge from prior tasks is overridden when learning a new
task. An inherent assumption in this paradigm is that once the model
transitions from task $T_{i}$ to task $T_{i+1}$, data from $T_{i}$ becomes
inaccessible, amplifying the importance of knowledge retention strategies (De
Lange et al., 2021).

While the premises and assumptions of FL differ from those of traditional
machine learning and continual learning, forgetting remains an issue. This can
be viewed as a side effect of data heterogeneity, a commonality FL shares with
CL. In FL, the global model evolves based on a fluctuating data distribution.
Specifically, in each communication round, a diverse set of sampled clients,
each with distinct data distributions, contribute with a model update.
Furthermore, these model updates need to be aggregated to obtain a global
model. This situation presents dual levels of data heterogeneity. Firstly, at
the intra-round level, heterogeneity arises from the participation of clients
with varied data distributions within the same round. This diversity can
inadvertently lead to “forgetting” specific data patterns or insights from
certain clients. Secondly, at the inter-round level, the participating clients
generally change from one round to the next. As a result, the global model may
“forget” or dilute insights gained from clients in previous rounds.

To remedy this issue, we propose Flashback, a FL algorithm that employs a
dynamic distillation approach to mitigate the effects of data heterogeneity.
Flashback’s dynamic distillation ensures that the local models learn new
knowledge while retaining knowledge from the global model during the client
updates by adaptively adjusting the distillation loss. Moreover, during the
server update, Flashback uses a very small public dataset as a medium to
integrate the knowledge from the local models to the global model using the
same dynamic distillation. Flashback performs these adaptations by estimating
the knowledge in each model using label counts as a proxy of the model
knowledge. Overall, Flashback results in a more stable and faster convergence
compared to existing methods.

Our contributions are the following:
$\bullet$ We investigate the forgetting problem in FL. We show that under
severe data heterogeneity, FL sufferers from forgetting. We dissect how and
where forgetting happens (§ 3).
$\bullet$ We propose a new metric for measuring forgetting over the
communication rounds(§ 3).
$\bullet$ We introduce _Flashback_ , a FL algorithm that employs a dynamic
distillation during the local updates and the server update (§ 4). By
addressing the forgetting issue, Flashback not only mitigates its detrimental
effects but also converges to the desired accuracy faster than existing
methods (§ 5).

##  2 Background

We consider a standard cross-device FL setup in which there are $N$ clients.
Each client $i$ has a unique dataset $D_{i}=\\{(x_{j},y_{j})\\}_{j=1}^{n_{i}}$
where $x_{j}$ represents the input features and $y_{j}$ is the ground-truth
label for $j$-th data point and $n_{i}$ represent the size of the local
dataset of client $i$. The goal is to train a single global model that
minimizes the objective:

$\min_{w\in\mathbb{R}^{d}}\sum^{N}_{i=1}\frac{|D_{i}|}{|\cup_{i\in[N]}D_{i}|}\left\\{L_{i}(w)=\frac{1}{|D_{i}|}\sum_{j=1}^{|D_{i}|}l(w;(x_{j},y_{j}))\right\\},$

where $L_{i}(w)$ represents the local loss for client $i$, and
$l(w;(x_{j},y_{j}))=\mathcal{L}_{\text{CE}}(F_{w}(x),y)$ is the cross-entropy
loss for a single data point, where $F_{w}$ denotes the model parameterized by
learnable weights $w$.

FedAvg provides a structured approach to efficiently address this distributed
problem. At each communication round $t$, the server randomly selects $K$
clients from the total available $N$ clients. These clients (denoted with
${\mathbb{S}}_{t}$) receive the previous global model, $w_{t-1}$. Then, they
update this model based on their local data using their local loss function
$L_{i}$. After updating, each client sends their modified model $w_{k,t}$ back
to the server that updates the global model using a weighted average of local
models, i.e.,
$w_{t}=\sum_{k\in{\mathbb{S}}_{t}}\frac{|D_{k}|w_{k,t}}{|\cup_{k\in[K]}D_{k}|}.$
To accommodate the intrinsic heterogeneity in client data, various FL
algorithms introduce modifications either at the local update level or during
the global aggregation. The nuances of these variations are further explored
in § 6.

Figure 2: Local (client) & Global Forgetting in some of the baselines using
CIFAR10. The first row represents the global model per-class test accuracy at
round $t-1$; then, the rows in the middle are all the clients that
participated in round $t$, and finally, in the last row, the global model at
the end of round $t$. Local forgetting happens when clients at round $t$ lose
the knowledge that the global model had at round $t-1$. The global forgetting
happens when the global model at round $t$ loses the knowledge that in the
clients’ models at round $t$.

Among these, Knowledge Distillation (KD) is a training method wherein a
smaller model, referred to as the student, is trained to reproduce the
behavior of a more complex model or ensemble called the teacher. Let
$F_{w_{s}}$ denote the student model with weights $w_{s}$ and $F_{w_{t}}$
represent the teacher model with weights $w_{t}$. For a given input $x$, the
student aims to minimize the following distillation loss:

$\displaystyle\mathcal{L}_{\text{KD}}((x,y);w_{s},w_{t})=\mathcal{L}_{\text{CE}}(F_{w_{s}}(x),y)(1-\alpha)$
(1) $\displaystyle+\mathcal{L}_{\text{KL}}(F_{w_{t}}(x),F_{w_{s}}(x))\alpha$

Here, $\mathcal{L}_{\text{CE}}$ is the standard cross-entropy loss with true
label $y$, and $\mathcal{L}_{\text{KL}}$ represents the Kullback-Leibler (KL)
divergence between the teacher’s and the student’s output probabilities. It is
defined as
$\mathcal{L}_{\text{KL}}({\bm{p}},{\bm{q}})=\sum_{c=1}^{C}{p}^{c}\log\left(\frac{{p}^{c}}{{q}^{c}}\right)$,
where $C$ is the number of classes, ${\mathbf{p}}$ is the target output
probability vector, and ${\mathbf{q}}$ is the predicted output probability
vector. The hyperparameter $\alpha\in[0,1]$ balances the importance between
the learning from the true labels and the teacher’s outputs.

While distillation originally emerged as a method for model compression
(Hinton et al., 2015; Buciluǎ et al., 2006; Schmidhuber, 1991), its utility
extends to FL. In the federated context, distillation can combat challenges
like data heterogeneity (Lin et al., 2020; Lee et al., 2021) and communication
efficiency (Jeong et al., 2018). Specifically, the global model can act as a
guiding teacher during local updates, directing the training process for each
client. Additionally, distillation techniques streamline the aggregation step,
assisting in the incorporation of varied knowledge from diverse clients to
update the global model. Furthermore, employing distillation for aggregation
mitigates model heterogeneity, allowing for the use of different model
architectures. Perhaps most interestingly, by using distillation, FL systems
can potentially bypass the traditional method of transmitting weight updates.
This is accomplished by sending soft labels that encapsulate the essence of
local updates–communication becomes more efficient, reducing bandwidth usage.

##  3 Forgetting in FL

We now investigate where forgetting happens and devise a metric to quantify
this phenomenon. Recall that in FL, the models are updated in two distinct
phases: 1) during local training – when each client $k$ starts from global
model $w_{t-1}$ and locally trains $w_{k,t}$ – and 2) during the aggregation
step – when the server combines the client models to update the new global
model $w_{t}$.

Intuitively, forgetting in FL is when knowledge contained in the global model
will be lost after the completion of communication round $w_{t-1}\rightarrow
w_{t}$. We observe that forgetting may occur in the two phases of FL. We refer
to the former case as local forgetting, where some knowledge in the global
model will be lost during the local training $w_{t-1}\rightarrow w_{k,t}$.
This is due to optimizing for the clients’ local objectives, which depend on
their datasets. Local forgetting is akin to the form of forgetting seen in CL,
where tasks change over time (as with clients in FL) and, consequently, the
data distribution. We refer to the latter case as aggregation forgetting,
where some knowledge contained in the clients’ model updates will be lost
during aggregation $\sum\\{w_{k,t}\mid k\in{\mathbb{S}}_{t}\\}\rightarrow
w_{t}$. This might be due to the coordinate-wise aggregation of weights as
opposed to matched averaging in the parameter space of DNNs (Wang et al.,
2020a).

We illustrate forgetting in Fig. 2 based on actual experiments with several
baseline methods. The figure shows the per-class accuracy of $w_{t-1}$, all
local models $w_{k,t}$, and the new global model $w_{t}$. The local forgetting
is evident in the drop in accuracy (lighter shade of blue) of the local models
$w_{k,t}$ compared to the global model $w_{t-1}$. The aggregation forgetting
is evident in the drop in accuracy of the global model $w_{t}$ compared to the
local models $w_{k,t}$. The figure also previews a result of our method,
Flashback, which shows a significant mitigation of forgetting. Local and
aggregation forgetting lead to the main forgetting problem in FL, which we
term as round forgetting, affecting $w_{t-1}\rightarrow w_{t}$.

In CL, forgetting is often quantified using Backward Transfer (BwT) (Chaudhry
et al., 2018). Lee et al. (2021) adapted this metric for FL as:

$\textstyle\mathcal{F}=\frac{1}{C}\sum_{c=1}^{C}\operatorname*{arg\,max}_{t\in{1,T-1}}(A^{c}_{t}-A^{c}_{T}),$
(2)

where $C$ is the number of classes and $A^{c}_{t}$ is the global model
accuracy on class $c$ at round $t$.

However, $\mathcal{F}$ is a coarse-grain score that evaluates forgetting in
aggregate across all rounds. We seek a finer-grain metric that measures
forgetting round-by-round. Furthermore, we wish to account for knowledge
replacement scenarios, such as when a decline in accuracy for one class might
be accompanied by an increase in another, essentially masking the negative
impact of forgetting in aggregate measures. Thus, for our evaluation results
(§ 5), we propose to measure round forgetting by focusing only on drops in
accuracy using the following metric:

$\textstyle\mathcal{F}_{t}=-\frac{1}{C}\sum_{c=1}^{C}\min(0,(A^{c}_{t}-A^{c}_{t-1}))$

where $t>1$ is the round at which forgetting is measured.

Our metric accounts for the pitfalls of the previously proposed forgetting
metric. It discounts knowledge replacement scenarios that can happen between
rounds by only focusing on the negative changes in accuracy. Furthermore, it
provides a granular view of forgetting, because it measures forgetting within
the rounds (whereas Eq. 2 measures how much the global model has forgotten by
the end of training).

###  3.1 Forgetting in Federated Learning and Continual Learning

Forgetting is a prominent problem in CL, where tasks change over time, and
consequently, the data distribution, places models at risk of overriding
previously learned knowledge. Looking at FL with the same perspective, we have
the intra-round and inter-round data distribution heterogeneity. Intra-round
clients with different data distributions participate, by updating the current
global model weights $w_{t-1}$. The server obtains a set of $\\{w_{k,t}\mid
k\in S_{t}\\}$ where $S_{t}$ is the set of clients participating in round $t$.
The goal is after the aggregation step, a new global model $w_{t}$ is obtained
containing all the knowledge that was present in $\\{w_{k,t}\mid k\in
S_{t}\\}$. Inter-round the global model $w_{t}$ has learned knowledge that
over the prior rounds $1\ldots t$. At each round, different sets of clients
participate. The goal is to carry this knowledge to the next round $t+1$ even
though the new set $S_{t+1}$ most likely is different with respect to the data
distribution than the previous set $S_{t}$. Presenting distinct forgetting
challenges in FL compared to CL.

##  4 Forgetting-Robust Federated Learning

Our key idea to mitigate forgetting is to leverage a dynamic form of knowledge
distillation, which is fine-tuned in response to the evolving knowledge
captured by the different models in the training process. During local
training, distillation ensures that each local model learns from the client’s
local dataset while retaining knowledge from the current global model. On the
server side, after the clients’ updates, Flashback begins by aggregating the
locally updated models—much in the vein of FedAvg. Then, Flashback distills
the knowledge of the freshly updated global model using our dynamic
distillation, learning from both its immediate predecessor—the global model
obtained at the previous round—and the ensemble of locally updated models,
which all play the role of teachers. The Flashback algorithm is detailed in
Algorithm 1. The remainder of this section discusses our distillation approach
in detail.

Algorithm 1 Flashback algorithm.

0: Initial global model $w_{0}$, number of rounds $T$, fraction of clients
$R$, minibatch size $B$, number of local epochs $E$, number of server epochs
$E_{s}$, learning rate $\eta$

0: global model $w_{T}$

1: $\bm{\pi}=\bm{0}\in\mathbb{R}^{C}$ // Global model’s label count vector

2: for $t=1$ to $T$ do

3: ${\mathbb{S}}_{t}\leftarrow$ Randomly select $\lceil R\cdot N\rceil$
clients

4: for each client $k\in{\mathbb{S}}_{t}$ do

5: $w_{k,t}\leftarrow w_{t-1}$ // Initialize local model with current global
model

6: $B_{k}\leftarrow$ Split local dataset into batches of size $B$

7: Compute $\bm{\alpha}$ with $\bm{\nu}$ as the local label count and a single
teacher $\bm{\mu}\leftarrow\bm{\pi}$

8: for $e=1$ to $E$ do

9: for each batch $b\in B_{k}$ do

10: Update $w_{k,t}$ using dKD loss $\mathcal{L}_{\text{dKD}}$

11: end for

12: end for

13: end for

14: $m_{t}\leftarrow\sum_{k\in{\mathbb{S}}_{t}}n_{k}$ // Total data points in
this round

15: $w_{t}\leftarrow\sum_{k\in{\mathbb{S}}_{t}}\frac{n_{k}}{m_{t}}w_{k,t}$ //
Average to obtain the new global model

16: $B_{s}\leftarrow$ Split the public dataset into batches of size $B$

17: ${\mathbb{T}}\leftarrow\\{w_{k,t}\mid
k\in{\mathbb{S}}_{t}\\}\cup\\{w_{t-1}\\}$

18: Compute $\bm{\alpha}$ with $\bm{\nu}\leftarrow\bm{\pi}$ and $\bm{\mu}_{i}$
as the label count $\forall w_{i}\in{\mathbb{T}}$

19: for $e=1$ to $E_{s}$ do

20: for each batch $b\in B_{s}$ do

21: Update $w_{t}$ using dKD loss $\mathcal{L}_{\text{dKD}}$

22: end for

23: end for

24: $r_{k}\leftarrow$ (Increment $r_{k}$ for every client
$k\in{\mathbb{S}}_{t}$)

25: Update participation count for each client $k\in{\mathbb{S}}_{t}$

26: for each client $k\in{\mathbb{S}}_{t}$ do

27: if $\gamma r_{k}\leq 1$ then

28: $\bm{\pi}\leftarrow\bm{\pi}+\gamma\bm{\mu}_{k}$

29: end if

30: end for

31: end for

###  4.1 Dynamic Distillation

As established in § 3, a client’s local model can forget and override model
knowledge with what is present in its private data. Moreover, even the global
model can be imperfect for two reasons:  i) As we established before, the
global model is susceptible to forgetting in the aggregation step.  ii)
Assuming no forgetting in the aggregation step, the knowledge contained in the
clients who participated so far might not represent all the available
knowledge, especially in the early rounds.  Overall, both local models and the
global model can be imperfect. Therefore, the logits of all the different
classes cannot be treated equally.

We propose to use the label count as an approximation of the knowledge within
a model. Here, the label count refers to the occurrences of each class in the
training data that the model saw during training.

In machine learning, a model’s knowledge is fundamentally tied to the data it
has been exposed to. If certain classes have higher representation (or label
counts) in the training data, it’s intuitive that the model would have more
opportunities to learn the distinguishing features of such classes.
Conversely, underrepresented classes might not offer the model sufficient
exposure to learn their nuances effectively. Our experimental results suggest
that per-class performance on the test set correlates highly with the label
counts in the training data. In scenarios where certain classes were more
abundant, the model demonstrated higher proficiency in predicting those
classes on the test set. As an example, Fig. 5 in the appendix, illustrates
for a randomly-chosen client that the client’s model performance on the test
set well reflects the label count distribution of its private data. From this
and many similar observations, we conclude that the label count can be
indicative of a model’s knowledge.

In standard knowledge distillation (Eq. 1), all logits are treated equally
since it is assumed that the teacher model has been trained on a balanced
dataset. Owing to the heterogeneity of data distribution in local datasets,
this assumption does not hold in FL. As a result, we cannot directly treat the
current global model nor the local model updates as equally reliable teachers
across all classes. Instead, we propose weighting the logits using the label
count as an approximation of the per-class knowledge within a model.

We now revisit the distillation loss in Eq. 1 and transform the scalar
$\alpha$ to a matrix form that is automatically tuned according to the label
count of both the student and the teachers and used directly within the KL
divergence loss. Namely, the dynamic $\bm{\alpha}$ parameter (defined below)
will change during the training as the label counts change. Flashback
maintains the global model counts over the rounds; this mechanism is detailed
in the next section.

We consider a single student model $F_{w_{s}}$ with weights $w_{s}$ and a set
${\mathbb{T}}$ of $K$ teacher models; the $i$-th teacher model is denoted as
$F_{w_{i}}$ with weights $w_{i}$. Let $\bm{\nu}\in\mathbb{R}^{C}$ be the
relative label count vector of the student model, where $\nu^{c}$ is the
relative occurrences of class $c$ in the dataset. Similarly, let
$\bm{\mu}_{i}\in\mathbb{R}^{C}$ be the relative label count vector of the
$i$-th teacher model.

The dynamic $\bm{\alpha}\in[0,1]^{K\times C}$ is defined as
$[\bm{\alpha}_{1}^{\intercal},\ldots,\bm{\alpha}_{K}^{\intercal}]$, with
$\alpha^{c}_{i}=\frac{\mu^{c}_{i}}{\nu^{c}+\sum_{k}\mu^{c}_{k}}$.

Then, we embed $\bm{\alpha}$ directly in the KL divergence loss
($\mathcal{L}_{\text{KL}}$ in Eq. 1) as follows:

$\mathcal{L}_{\text{dKL}}({\bm{p}},{\bm{q}};\bm{\alpha}_{i})=\sum_{c=1}^{C}\alpha_{i}^{c}\cdot{p}^{c}\log\left(\frac{{p}^{c}}{{q}^{c}}\right)$

Similar to standard distillation, to account for the student model knowledge
with respect to the ground-truth class $y$, we define
$\alpha^{c}_{s}=\frac{\nu^{c}}{\nu^{c}+\sum_{k}\mu^{c}_{k}}$. Thus,
$\alpha^{c}_{s}+\sum_{k=1}^{K}\alpha^{c}_{k}=1$ for all classes $c\in[C]$.

Finally, the dynamic knowledge distillation loss ($\mathcal{L}_{\text{dKD}}$)
is:

$\displaystyle\mathcal{L}_{\text{dKD}}((x,y);w_{s},{\mathbb{T}},\bm{\alpha})=\alpha^{y}_{s}\mathcal{L}_{\text{CE}}(F_{w_{s}}(x),y)$
(3)
$\displaystyle\;+\sum_{{w_{i}}\in{\mathbb{T}}}\mathcal{L}_{\text{dKL}}(F_{w_{i}}(x),F_{w_{s}}(x);\bm{\alpha}_{i})$

The dynamic $\bm{\alpha}$ will weigh the divergence between the logits of
different classes, making the student model focus more on learning from the
teacher’s strengths while being cautious of its weaknesses. This is of great
importance in FL because of the data heterogeneity problem. For instance, in
the initial training rounds, the global model may not encounter certain
classes. Our distillation approach assigns a zero weight to the divergence of
these classes, shielding the client model from adopting unreliable knowledge
from the global model. Similarly, if a client possesses significantly larger
data for a specific class compared to what the global model has encountered,
the weight assigned to that class’s divergence will be small. This implies
that the client model’s remains more grounded in classes where it has more
comprehensive data.

A property of our distillation is that it will ignore the global model as a
teacher in the first communication round. Since the initial global model label
counts are all zeros, Eq. 3 reduces to just the cross-entropy:
$\mathcal{L}_{\text{dKD}}=\mathcal{L}_{\text{CE}}$.

(a) CIFAR10.

(b) CINIC10.

(c) FEMNIST.

Figure 3: Round-to-accuracy performance of Flashback and other baselines over
training rounds.

(a) CIFAR10.

(b) CINIC10.

Figure 4: Transition of local models loss to the global model loss over the
rounds.

###  4.2 Estimating the Global Model Knowledge

Note that to apply the dynamic distillation loss Eq. 3, we require to obtain
the student’s and teachers’ label count vectors. While the label count of
local models can be easily obtained (from the class frequency of local
datasets), the label count of the global model is not readily available. We
construct $\bm{\pi}$, the global model’s relative label count, as follows. Let
$r_{k}$ denote the number of rounds that client $k$ has participated in. For
every client $k$ that participates at round $t$, Flashback adds a fraction
$\gamma\in(0,1]$ of client $k$ label count ($\bm{\mu}_{k}$) to $\bm{\pi}$,
unless $\gamma r_{k}>1$, in which case $\bm{\pi}$ is not updated based on
$k$’s label count. The latter case means that client $k$ has participated
enough times that its label count is fully accounted for in $\bm{\pi}$.

Intuitively, the parameter $\gamma$ indicates the rate at which we rely on the
global model. When $\gamma$ is set to 1, it implies complete trust in the
global model’s ability to incorporate the clients’ knowledge after just one
round of participation. However, expecting such immediate and full
assimilation is unrealistic, so we typically set $\gamma<1$. The gradual
build-up of the global label count plays a vital role in maintaining a
balanced distillation in Eq. 3 during local updates. This progressive approach
mirrors our growing trust in the global model’s capabilities. It prevents the
risk of assigning excessively high weights too soon, which could otherwise
hurt the learning process.

###  4.3 Label Count Motivation

Fig. 5 shows, on the left, the per-class accuracy of a randomly-chosen example
client from a FedAvg training experiment. On the right, the figure shows the
corresponding label count at the client. This example suggests that label
count can be representative of the model performance.

Figure 5: (left) Per-class accuracy of a client model on all the rounds where
it participated. (right) Data distribution of that client.

##  5 Experiments & Results

We outline and analyze our experimental findings to investigate whether
mitigating forgetting successfully addresses the issues of slow and unstable
convergence observed in the initial problem laid out in § 1. The experimental
results stem from three settings: CIFAR10 and CINIC10, where heterogeneous
data partitions are created using Dirichlet distribution with $\beta=0.1$ and
FEMNIST with 3,432 clients, following the natural heterogeneity of the
dataset. Furthermore, we do an ablation study on the different components of
the algorithm. We use the same neural network architecture that is used in Lee
et al. (2021); McMahan et al. (2017), which is a 2-layer Convolutional Neural
Network (CNN). Summaries of the datasets, partitions, and more details on the
experimental setup, as well as additional results, are reported in Appendix A.

We compare Flashback against several baseline methods, namely:  1) FedAvg
(McMahan et al., 2017),  2) FedDF (Lin et al., 2020),  3) FedNTD (Lee et al.,
2021),  4) FedProx (Li et al., 2020),  5) FedReg (Xu et al., 2022),  6) MOON
(Li et al., 2021b).  It is noteworthy that both FedNTD and FedReg target
forgetting in FL (discussed further in § 6). Flashback server distillation is
performed until early stopping is triggered on the public validation set
(details in Appendix A). Moreover, Flashback only introduces one additional
hyperparameters $\gamma$, which represents how fast trust is built in the
global model. We analyze the effect of $\gamma$ later in the section. We start
by evaluating Flashback performance by showing round-to-accuracy, round
forgetting, and the local-global loss over the rounds.

Improved round-to-accuracy. We evaluate the learning efficiency of Flashback
and other baselines by showing the accuracy over rounds in Fig. 3. Flashback
consistently shows a faster convergence to a high accuracy. Furthermore, we
show the number of rounds it takes to reach a target accuracy and fractions of
that target accuracy in Table 1; we include the result of a central training
on the public dataset. Flashback shows a much faster convergence than other
baselines. This indicates that addressing forgetting on the clients’ local
update and at the aggregation step does provide training stability and indeed
a faster convergence.

Table 1: Number of rounds to reach accuracy $A_{x}=A\cdot x$ where $A$ is the target accuracy and $x$ is a fraction of it. | CIFAR10, $A=48.2$% | CINIC10, $A=43.5$% | FEMNIST, $A=69.5$%
---|---|---|---
| $A_{0.5}$ | $A_{0.75}$ | $A_{0.95}$ | $A_{0.5}$ | $A_{0.75}$ | $A_{0.95}$ | $A_{0.5}$ | $A_{0.75}$ | $A_{0.95}$
FedAvg | 12 | 82 | - | 13 | - | - | 49 | 75 | 138
FedDF | 7 | 40 | 112 | 2 | 30 | - | - | - | -
FedNTD | 12 | 41 | 112 | 13 | 46 | - | - | - | -
FedProx | 35 | 93 | - | 13 | - | - | 142 | - | -
FedReg | 35 | 108 | - | 16 | - | - | - | - | -
MOON | 82 | - | - | 124 | - | - | - | - | -
Flashback | 2 | 4 | 10 | 4 | 5 | 6 | 3 | 5 | 16

Less round forgetting. We show the empirical cumulative distribution function
(ECDF) of the round forgetting in Fig. 6. We see that Flashback successfully
reduces round forgetting. Also, FedNTD has less round forgetting than the
remaining baselines. In the appendix, we show the round forgetting over the
rounds (Fig. 12).

(a) CIFAR10.

(b) CINIC10.

Figure 6: Distribution of round forgetting of Flashback compared to other
baselines.

Minimizing local models divergence. To further understand the effect of
Flashback on the training behavior, we show the transition of the mean loss of
the local models to the loss of the global model over the rounds in Fig. 4.
This gives us an insight into the effect of the regularization made by our
dynamic distillation. We see that the mean loss of the local models of the
other baselines always spikes, signifying a divergence of these models from
the global training objective, while Flashback has a much more stable loss.
This shows that Eq. 3 regularizes the local models well so that they do not
diverge too much from the global training objective.

In the rest of the section, we delve deeper into Flashback to understand its
behavior and validate its performance gains.

Dissecting the distillation. To show the importance of performing dynamic
distillation during clients’ updates and at the server’s aggregation step, we
conduct an experiment where we run Flashback with local distillation only and
with server distillation only (c.f. Fig. 7). We observe that doing dynamic
distillation at either side of the algorithm – client update and aggregation
step – doesn’t address forgetting or gets a similar performance to Flashback.

Figure 7: Performing distillation at only one side of the algorithm (client &
server) on CIFAR10.

To validate the _importance of dynamic distillation at client- and server
sides_ towards Flashback’s performance gains, we create a baseline where we
replace the local distillation loss with not-true distillation (NTD) loss Lee
et al. (2021). From Fig. 8, we observe that this baseline doesn’t perform as
well as Flashback, and performs similarly to FedNTD. This indicates that
performing the distillation at the server needs well-regularized local models
(teachers), which is further supported by our previous experiment contrasting
Flashback to single-side distillation (c.f. Fig. 7).

Figure 8: Using NTD loss instead of Flashback’s dynamic distillation at the
local update on CIFAR10.

Effect of public dataset size. Flashback requires the availability of a public
labeled dataset. This may be a limiting assumption in some cases. To study
this limitation, we explore a few scenarios for the size of the public dataset
in Fig. 9:  1) 9000 samples ($15$% of CIFAR10),  2) 1125 samples ($1.88$% of
CIFAR10),  3) 1283 samples that have unbalanced class distribution ($2.14$% of
CIFAR10),  4) 450 samples ($0.75$% of CIFAR10).  For all of these scenarios,
we train a model centrally on the public dataset. We find that Flashback can
benefit from a large balanced public dataset. Most importantly, Flashback can
work well with a small public dataset (1125 samples is the default in all
experiments). Furthermore, even if the public dataset has a class imbalance
Flashback still performs relatively well. In all of the cases, Flashback
always outperforms central training on the public dataset. Overall, Flashback
requires the availability of a public dataset, however, it does not require a
huge amount of data or hard requirements for the class distribution to be very
balanced.

Figure 9: Flashback using different public datasets, and the results of
central training on public datasets of similar sizes.

Training on the public dataset. We also experiment to answer the following
question: _does the performance improvement of Flashback come from the fact
that we train the global model on a public labeled dataset?_ To answer this
question we create a naive baseline, where we extend FedAvg to fine-tune the
global model after the aggregation step at every communication round. From
Fig. 10, we see that FedAvg with fine-tuning quickly reaches a stale model and
eventually collapses. We believe this collapse happens due to the local models
diverging too much after the local update such that aggregating those models
fails. This is evident by the spike of the local model’s loss.

Figure 10: Comparing Flashback to FedAvg with fine-tuning: (left) test
accuracy over 3 runs on CIFAR10; (right) the transition of the local models’
loss to the global model loss over the rounds.

The importance of  $\gamma$. As mentioned, Flashback has a single
hyperparameter $\gamma$, which dictates how fast models will trust the global
model as a competent teacher. We explore the effect of this hyperparameter in
Fig. 11. We find that setting this parameter to a larger value leads the
learning process to get to a stale solution quickly. This is intuitive since
large $\gamma$ leads the global model label count to grow faster, therefore,
this dominates the loss term in Eq. 3 during both the client- and the server
update. Smaller $\gamma$ gives the best performance because it gives the local
models time to learn from their own private data by having a small weight to
the distillation term in Eq. 3. However, too small value for $\gamma$ such as
$0.001$ slows the training process, since during the local training the
distillation term in Eq. 3 will be very small in the early rounds.

Figure 11: Varying the trust of the global model parameter $\gamma$.

###  5.1 Additional Results

In Fig. 12, we show the round forgetting score computed over the rounds. We
see that the baselines have very flaunting round forgetting score.

In Fig. 13, we show the transition of the average loss of the local models to
the global model loss on the test set. Interestingly, we can see that even
though performing local distillation only doesn’t have the same performance as
Flashback, it does mitigate the local forgetting. That is, we do not see a
spike in the loss after the clients perform their local update.

In Fig. 14, we see the ECDF of the round forgetting. The Flashback variant
with NTD even shows worse forgetting than Flashback and FedNTD. Further
showing that just performing distillation at both sides doesn’t address the
forgetting problem.

(a) CIFAR10.

(b) CINIC10.

Figure 12: The round forgetting of Flashback and other baselines over training
rounds.

Figure 13: Performing distillation at only one side of the algorithm (client &
server) on CIFAR10.

Figure 14: Using Not-True Distillation instead of Flashback’s dynamic
distillation at the local update on CIFAR10.

##  6 Related Work

Federated learning. FL is commonly viewed as a ML paradigm wherein a server
distributes the training process on a set of decentralized participants that
train a shared global model using local datasets that are never shared
(Konečnỳ et al., 2015; Shokri & Shmatikov, 2015; Konečnỳ et al., 2016;
Konečnỳ, 2017; Li et al., 2020; McMahan et al., 2017; Kairouz et al., 2019).
FL has been used to enhance prediction quality for virtual keyboards among
other applications (Bonawitz et al., 2019; Yang et al., 2018). A number of FL
frameworks have facilitated research in this area (Caldas et al., 2019;
OpenMined, 2020; tensorflow.org, 2020; Abdelmoniem et al., 2023).

Heterogeneity in FL. A key challenge in FL systems is uncertainties stemming
from learner, system, and data heterogeneity. The non-IID distributions of
Learners’ data can significantly slow down convergence (McMahan et al., 2017;
Kairouz et al., 2019) and several algorithms are proposed as means of
mitigation (Fourati et al., 2023; Wang et al., 2020b; Karimireddy et al.,
2020; Li et al., 2020, 2021a).

In this section, we review works that targeted the issue of forgetting in FL.
Forgetting in FL is an under-studied area that poses significant challenges,
leading to slow model convergence and loss of crucial knowledge acquired
during the learning process (Chaudhry et al., 2018; Dupuy et al., 2023). There
have been some notable attempts to mitigate the impact of forgetting on the
learning process (Lee et al., 2021; Xu et al., 2022).

FedReg (Xu et al., 2022) addresses the issue of slow convergence in FL,
asserting it to be a result of forgetting at the local update phase. They
demonstrate this by comparing the loss of the global model $w_{t-1}$ on
specific client data points with the averaged loss of updated clients’ models
$\\{w_{t,k}\mid k\in{\mathbb{S}}_{t}\\}$ on the same data points, highlighting
a significant increase in the average loss, indicative of forgetting. However,
in our work, we propose a systematic way of measuring forgetting using a
metric designed to capture it. Furthermore, we show that forgetting doesn’t
only occur in the local update, but it also happens at the aggregation step
(§§ 3 and 2). FedReg proposes to generate fake data that carries the
previously attained knowledge. During the local update, Fast Gradient Sign
Method (Goodfellow et al., 2014) is used to generate these data using the
global model $w_{t-1}$ and the client data. Then, the loss of the generated
data is used to regularize the local update. While FedReg employs
regularization using synthetic data during local updates, our work, Flashback,
leverages dynamic distillation to ensure knowledge retention at both local
updates and aggregation steps.

FedNTD (Lee et al., 2021) makes a connection between CL and FL, suggesting
that forgetting happens in FL as well. Similarly to FedReg, their analysis
shows that forgetting happens at the local update, where global knowledge that
lies outside of the local distribution of the client is susceptible to
forgetting. To address this, they propose to use a new variant of distillation
Eq. 1 named Not-True Distillation (NTD), that masks the ground-truth class
logits in the KL divergence as
$\mathcal{L}_{\text{KL}}({\bm{p}},{\bm{q}})=\sum_{i=c,c\neq
y}^{C}{p}^{c}\log(\frac{{p}^{c}}{{q}^{c}})$, where $y$ is the ground-truth
class. NTD is used at the local update, while all the other steps in the
algorithm remain the same as FedAvg. FedNTD aims to preserve global knowledge
during the local update.

Both FedReg and FedNTD diagnose the issue of forgetting primarily within the
realm of local updates, asserting that this stage risks losing valuable global
knowledge. Consequently, both works present innovative solutions specifically
tailored to counteract this local update forgetting. However, their
perspective overlooks a pivotal aspect of the forgetting problem: the
occurrence of forgetting during the aggregation step. As we delve into in § 3,
this oversight in recognizing and addressing forgetting during aggregation has
repercussions on the later local updates. In contrast, Flashback takes a
holistic approach, targeting forgetting comprehensively across both the local
updates and the aggregation phase, leading to faster convergence.

##  7 Conclusion

We explored the phenomenon of forgetting in FL. Our investigation revealed
that forgetting occurs during both local and global update phases of FL
algorithms. We presented Flashback, a novel FL algorithm explicitly designed
to counteract forgetting by employing dynamic knowledge distillation. Our
approach leverages data label counts as a proxy for knowledge, ensuring a more
targeted and effective forgetting mitigation. Our empirical results showed
Flashback’s efficacy in mitigating round forgetting, thereby supporting the
hypothesis that the observed slow and unstable convergence in FL algorithms is
closely linked to forgetting. This result underlines the importance of
addressing forgetting, paving the way for the advancement of more robust and
efficient FL algorithms.

## References

  * Abdelmoniem et al. (2023) Abdelmoniem, A. M., Sahu, A. N., Canini, M., and Fahmy, S. A.  REFL: Resource-Efficient Federated Learning.  In _ACM EuroSys_ , 2023.
  * Bonawitz et al. (2019) Bonawitz, K., Eichner, H., Grieskamp, W., Huba, D., Ingerman, A., Ivanov, V., Kiddon, C., Konečný, J., Mazzocchi, S., McMahan, B., Van Overveldt, T., Petrou, D., Ramage, D., and Roselander, J.  Towards Federated Learning at Scale: System Design.  In _MLSys_ , 2019.
  * Buciluǎ et al. (2006) Buciluǎ, C., Caruana, R., and Niculescu-Mizil, A.  Model compression.  In _KDD_ , 2006.
  * Caldas et al. (2019) Caldas, S., Duddu, S. M. K., Wu, P., Li, T., Konečný, J., McMahan, H. B., Smith, V., and Talwalkar, A.  LEAF: A Benchmark for Federated Settings.  In _Workshop on Federated Learning for Data Privacy and Confidentiality_ , 2019.
  * Chaudhry et al. (2018) Chaudhry, A., Dokania, P. K., Ajanthan, T., and Torr, P. H.  Riemannian walk for incremental learning: Understanding forgetting and intransigence.  In _Proceedings of the European conference on computer vision (ECCV)_ , 2018.
  * Darlow et al. (2018) Darlow, L. N., Crowley, E. J., Antoniou, A., and Storkey, A. J.  Cinic-10 is not imagenet or cifar-10.  _arXiv preprint arXiv:1810.03505_ , 2018.
  * De Lange et al. (2021) De Lange, M., Aljundi, R., Masana, M., Parisot, S., Jia, X., Leonardis, A., Slabaugh, G., and Tuytelaars, T.  A continual learning survey: Defying forgetting in classification tasks.  _IEEE transactions on pattern analysis and machine intelligence_ , 2021.
  * Dupuy et al. (2023) Dupuy, C., Majmudar, J., Wang, J., Roosta, T. G., Gupta, R., Chung, C., Ding, J., and Avestimehr, S.  Quantifying catastrophic forgetting in continual federated learning.  In _IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2023.
  * Fourati et al. (2023) Fourati, F., Kharrat, S., Aggarwal, V., Alouini, M.-S., and Canini, M.  FilFL: Accelerating Federated Learning via Client Filtering, 2023.  URL https://arxiv.org/abs/2302.06599.
  * Goodfellow et al. (2014) Goodfellow, I. J., Shlens, J., and Szegedy, C.  Explaining and harnessing adversarial examples.  _arXiv preprint arXiv:1412.6572_ , 2014.
  * Hinton et al. (2015) Hinton, G., Vinyals, O., and Dean, J.  Distilling the knowledge in a neural network, 2015.
  * Jeong et al. (2018) Jeong, E., Oh, S., Kim, H., Park, J., Bennis, M., and Kim, S.-L.  Communication-Efficient On-Device machine learning: Federated distillation and augmentation under Non-IID private data.  2018\.
  * Kairouz et al. (2019) Kairouz, P., Brendan McMahan, H., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z., Cormode, G., Cummings, R., D’Oliveira, R. G. L., Eichner, H., El Rouayheb, S., Evans, D., Gardner, J., Garrett, Z., Gascón, A., Ghazi, B., Gibbons, P. B., Gruteser, M., Harchaoui, Z., He, C., He, L., Huo, Z., Hutchinson, B., Hsu, J., Jaggi, M., Javidi, T., Joshi, G., Khodak, M., Konečný, J., Korolova, A., Koushanfar, F., Koyejo, S., Lepoint, T., Liu, Y., Mittal, P., Mohri, M., Nock, R., Özgür, A., Pagh, R., Raykova, M., Qi, H., Ramage, D., Raskar, R., Song, D., Song, W., Stich, S. U., Sun, Z., Suresh, A. T., Tramèr, F., Vepakomma, P., Wang, J., Xiong, L., Xu, Z., Yang, Q., Yu, F. X., Yu, H., and Zhao, S.  Advances and open problems in federated learning.  _arXiv 1912.04977_ , 2019.
  * Karimireddy et al. (2020) Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., and Suresh, A. T.  Scaffold: Stochastic controlled averaging for federated learning.  In _International conference on machine learning_ , pp. 5132–5143, 2020.
  * Konečnỳ (2017) Konečnỳ, J.  Stochastic, distributed and federated optimization for machine learning.  _arXiv preprint arXiv:1707.01155_ , 2017.
  * Konečnỳ et al. (2015) Konečnỳ, J., McMahan, B., and Ramage, D.  Federated optimization: Distributed optimization beyond the datacenter.  _arXiv preprint arXiv:1511.03575_ , 2015.
  * Konečnỳ et al. (2016) Konečnỳ, J., McMahan, H. B., Yu, F. X., Richtárik, P., Suresh, A. T., and Bacon, D.  Federated learning: Strategies for improving communication efficiency.  _arXiv preprint arXiv:1610.05492_ , 2016.
  * Krizhevsky (2009) Krizhevsky, A.  Learning multiple layers of features from tiny images.  Technical report, University of Toronto, 2009.
  * Lee et al. (2021) Lee, G., Jeong, M., Shin, Y., Bae, S., and Yun, S.-Y.  Preservation of the global knowledge by Not-True distillation in federated learning.  2021\.
  * Li et al. (2021a) Li, L., Duan, M., Liu, D., Zhang, Y., Ren, A., Chen, X., Tan, Y., and Wang, C.  FedSAE: A Novel Self-Adaptive Federated Learning Framework in Heterogeneous Systems.  In _IJCNN_ , 2021a.
  * Li et al. (2021b) Li, Q., He, B., and Song, D.  Model-contrastive federated learning.  In _Conference on computer vision and pattern recognition_ , 2021b.
  * Li et al. (2020) Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V.  Federated optimization in heterogeneous networks.  _Proceedings of Machine learning and systems_ , 2020.
  * Lin et al. (2020) Lin, T., Kong, L., Stich, S. U., and Jaggi, M.  Ensemble distillation for robust model fusion in federated learning.  _Advances in Neural Information Processing Systems_ , 2020.
  * McMahan et al. (2017) McMahan, B., Moore, E., Ramage, D., Hampson, S., and Arcas, B. A. y.  Communication-Efficient Learning of Deep Networks from Decentralized Data.  In Singh, A. and Zhu, J. (eds.), _AISTATS_ , 2017.
  * OpenMined (2020) OpenMined.  Syft + grid provides secure and private deep learning in python, 2020.  URL https://github.com/OpenMined/PySyft.
  * Parisi et al. (2019) Parisi, G. I., Kemker, R., Part, J. L., Kanan, C., and Wermter, S.  Continual lifelong learning with neural networks: A review.  _Neural networks_ , 2019.
  * Schmidhuber (1991) Schmidhuber, J.  _Neural sequence chunkers_.  Inst. für Informatik, 1991.
  * Shokri & Shmatikov (2015) Shokri, R. and Shmatikov, V.  Privacy-preserving deep learning.  In _Proceedings of the 22nd ACM SIGSAC conference on computer and communications security_ , pp. 1310–1321, 2015.
  * tensorflow.org (2020) tensorflow.org.  Tensorflow federated: Machine learning on decentralized data, 2020.  URL https://www.tensorflow.org/federated.
  * Wang et al. (2020a) Wang, H., Yurochkin, M., Sun, Y., Papailiopoulos, D., and Khazaeni, Y.  Federated learning with matched averaging.  _arXiv preprint arXiv:2002.06440_ , 2020a.
  * Wang et al. (2020b) Wang, J., Liu, Q., Liang, H., Joshi, G., and Poor, H. V.  Tackling the objective inconsistency problem in heterogeneous federated optimization.  In _NeurIPS_ , 2020b.
  * Xu et al. (2022) Xu, C., Hong, Z., Huang, M., and Jiang, T.  Acceleration of federated learning with alleviated forgetting in local training.  _arXiv preprint arXiv:2203.02645_ , 2022.
  * Yang et al. (2018) Yang, T., Andrew, G., Eichner, H., Sun, H., Li, W., Kong, N., Ramage, D., and Beaufays, F.  Applied Federated Learning: Improving Google Keyboard Query Suggestions, 2018.

##  Appendix A Experiments Details

###  A.1 Datasets

In this section, we provide an overview of the datasets used, the data split,
and the specific experimental setups. For each dataset, we perform two sets of
experiments to analyze the effects of data heterogeneity on the algorithms’
performance. The datasets used are CIFAR10, CINIC10, and FEMNIST.

CIFAR10 (Krizhevsky, 2009). A famous vision dataset that includes 50k training
images and 10k testing images. We emulate a realistic, heterogeneous data
distribution by using a Dirichlet distribution with $\beta=0.1$. A $\beta$
value of $0.1$ is chosen to simulate a more heterogeneous, and challenging
data distribution. A $2.5$% random sample of the training set creates a public
dataset, further divided into training and validation sets. This yields a very
small public training dataset with size of $1.88$%. The remaining $97.5$% is
distributed among 100 clients, with each client’s data being split into
training (90%) and validation (10%) subsets.

CINIC10 (Darlow et al., 2018). A drop-in replacement of CIFAR10, this dataset
comprises 90k training, 90k validation, and 90k test images. We merge the
training and validation sets and adopt a similar approach as with CIFAR10,
taking out $2.5$% of the data to be the public dataset, similar to the CIFAR10
case this $2.5$% is further divided into training set and validation set. Then
employing Dirichlet distribution with $\beta$ value of 0.1 to split the
$97.5$% remaining data into 200 clients, with each client’s data further
divided into training (90%) and validation (10%) sets.

(a) CIFAR10, Dir$(\beta=0.1)$

(b) CINIC10, Dir$(\beta=0.1)$

(c) FEMNIST with 3432 clients

Figure 15: Clients data distribution. The x-axis is the clients and the y-axis
is the labels.

FEMNIST (Caldas et al., 2019). This federated learning dataset is based on
extended MNIST with natural heterogeneity, where each writer is considered a
client. From the 3597 total writers, those with less than 50 samples are
excluded. We randomly selected 150 writers to form a public dataset. The
remaining 3432 writers’ data is divided into train (approx. 70%), validation
(approx. 15%), and test (approx. 15%) sets. The collective test sets from all
writers form the overall test set. At every round, 32 clients are randomly
selected for participation.

For CIFAR10 and CINIC10, we chose $\beta$ values of 0.1 and client
participation value of 10. While for FEMNIST we have 3432 clients (writers)
with client participation vale of 32 In all cases, the training data
distribution among clients is illustrated in Fig. 15.

###  A.2 Baselines & Hyperparameters

We evaluate the following algorithms as baselines:  1. FedAvg (McMahan et al.,
2017);  2. FedDF (Lin et al., 2020);  3. FedNTD (Lee et al., 2021);  4.
FedProx (Li et al., 2020);  5. FedReg (Xu et al., 2022); and  6. MOON (Li et
al., 2021b) . Both FedNTD and FedReg target forgetting in FL (discussed in §
6). We use the same neural network architecture that is used in Lee et al.
(2021); McMahan et al. (2017), which is a 2-layer CNN. Note that for MOON (Li
et al., 2021b) we add an additional layers to the model as described in their
source code for the projection head. Moreover, for the optimizer, learning
rate, and model we follow Lee et al. (2021); McMahan et al. (2017), and when a
baseline has different hyperparameters we use their proposed values. For
example, in FedDF the number of local epochs is set to 40, while in the other
baselines and Flashback, it is set to 5 epochs. As for Flashback
hyperparameters, during the server distillation, we train until early stopping
gets triggered using the validation set; we set the label count fraction
$\gamma=0.025$ for CIFAR10, i.e., we add $2.5$% of the client label count each
time it participates, while we set $\gamma=0.1$ for CINIC10 and FEMNIST. As
for distillation-specific hyperparameters, we have one fewer hyperparameter
since $\bm{\alpha}$ is computed automatically, and for temperature, we use the
standard $T=3$.