11institutetext: Gordon Center for Medical Imaging, Department of Radiology,
Massachusetts General Hospital and Harvard Medical School, Boston, MA,
0211422institutetext: Department of Radiation Oncology, Massachusetts General
Hospital and Harvard Medical School, Boston, MA, 02114

# Incremental Learning for Heterogeneous Structure Segmentation in Brain Tumor
MRI

Xiaofeng Liu 11 Helen A. Shih 22 Fangxu Xing 11 Emiliano Santarnecchi 11
Georges El Fakhri 11 Jonghye Woo 11

###### Abstract

Deep learning (DL) models for segmenting various anatomical structures have
achieved great success via a static DL model that is trained in a single
source domain. Yet, the static DL model is likely to perform poorly in a
continually evolving environment, requiring appropriate model updates. In an
incremental learning setting, we would expect that well-trained static models
are updated, following continually evolving target domain data—e.g.,
additional lesions or structures of interest—collected from different sites,
without catastrophic forgetting. This, however, poses challenges, due to
distribution shifts, additional structures not seen during the initial model
training, and the absence of training data in a source domain. To address
these challenges, in this work, we seek to progressively evolve an “off-the-
shelf” trained segmentation model to diverse datasets with additional
anatomical categories in a unified manner. Specifically, we first propose a
divergence-aware dual-flow module with balanced rigidity and plasticity
branches to decouple old and new tasks, which is guided by continuous batch
renormalization. Then, a complementary pseudo-label training scheme with self-
entropy regularized momentum MixUp decay is developed for adaptive network
optimization. We evaluated our framework on a brain tumor segmentation task
with continually changing target domains—i.e., new MRI scanners/modalities
with incremental structures. Our framework was able to well retain the
discriminability of previously learned structures, hence enabling the
realistic life-long segmentation model extension along with the widespread
accumulation of big medical data.

##  1 Introduction

Accurate segmentation of a variety of anatomical structures is a crucial
prerequisite for subsequent diagnosis or treatment [28]. While recent advances
in data-driven deep learning (DL) have achieved superior segmentation
performance [29], the segmentation task is often constrained by the
availability of costly pixel-wise labeled training datasets. In addition, even
if static DL models are trained with extraordinarily large amounts of training
datasets in a supervised learning manner [29], there exists a need for a
segmentor to update a trained model with new data alongside incremental
anatomical structures [24].

In real-world scenarios, clinical databases are often sequentially constructed
from various clinical sites with varying imaging protocols [19, 20, 21, 23].
As well, labeled anatomical structures are incrementally increased with
additional lesions or new structures of interest, depending on study goals or
clinical needs [27, 18]. Furthermore, access to previously used data for
training can be restricted, due to data privacy protocols [18, 17]. Therefore,
efficiently utilizing heterogeneous structure-incremental (HSI) learning is
highly desired for clinical practice to develop a DL model that can be
generalized well for different types of input data and varying structures
involved. Straightforwardly fine-tuning DL models with either new structures
[30] or heterogeneous data [17] in the absence of the data used for the
initial model training, unfortunately, can easily overwrite previously learned
knowledge, i.e., catastrophic forgetting [30, 17, 14].

At present, satisfactory methods applied in the realistic HSI setting are
largely unavailable. $First$, recent structure-incremental works cannot deal
with domain shift. Early attempts [27] simply used exemplar data in the
previous stage. [5, 33, 30, 18] combined a trained model prediction and a new
class mask as a pseudo-label. However, predictions from the old model under a
domain shift are likely to be unreliable [38]. The widely used pooled feature
statistics consistency [5, 30] is also not applicable for heterogeneous data,
since the statistics are domain-specific [2]. In addition, a few works [13,
25, 34] proposed to increase the capacity of networks to avoid directly
overwriting parameters that are entangled with old and new knowledge. However,
the solutions cannot be domain adaptive. $Second$, from the perspective of
continuous domain adaptation with the consistent class label, old exemplars
have been used for the application of prostate MRI segmentation [32]. While Li
et al. [17] further proposed to recover the missing old stage data with an
additional generative model, hallucinating realistic data, given only the
trained model itself, is a highly challenging task [31] and may lead to
sensitive information leakage [35]. $Third$, while, for natural image
classification, Kundu et al. [16] updated the model for class-incremental
unsupervised domain adaption, its class prototype is not applicable for
segmentation.

In this work, we propose a unified HSI segmentor evolving framework with a
divergence-aware decoupled dual-flow (D3F) module, which is adaptively
optimized via HSI pseudo-label distillation using a momentum MixUp decay (MMD)
scheme. To explicitly avoid the overwriting of previously learned parameters,
our D3F follows a “divide-and-conquer” strategy to balance the old and new
tasks with a fixed rigidity branch and a compensated learnable plasticity
branch, which is guided by our novel divergence-aware continuous batch
renormalization (cBRN). The complementary knowledge can be flexibly integrated
with the model re-parameterization [4]. Our additional parameters are constant
in training, and 0 in testing. Then, the flexible D3F module is trained
following the knowledge distillation with novel HSI pseudo-labels.
Specifically, inspired by the self-knowledge distillation [15] and self-
training [38] that utilize the previous prediction for better generalization,
we adaptively construct the HSI pseudo-label with an MMD scheme to smoothly
adjust the contribution of potential noisy old model predictions on
heterogeneous data and progressively learned new model predictions along with
the training. In addition, unsupervised self-entropy minimization is added to
further enhance performance.

Our main contributions can be summarized as follow:

$\bullet$ To our knowledge, this is the first attempt at realistic HSI
segmentation with both incremental structures of interest and diverse domains.

$\bullet$ We propose a divergence-aware decoupled dual-flow module guided by
our novel continuous batch renormalization (cBRN) for alleviating the
catastrophic forgetting under domain shift scenarios.

$\bullet$ The adaptively constructed HSI pseudo-label with self-training is
developed for efficient HSI knowledge distillation.

We evaluated our framework on anatomical structure segmentation tasks from
different types of MRI data collected from multiple sites. Our HSI scheme
demonstrated superior performance in segmenting all structures with diverse
data distributions, surpassing conventional class-incremental methods without
considering data shift, by a large margin.

##  2 Methodology

For the segmentation model under incremental structures of interest and domain
shift scenarios, we are given an off-the-shelf segmentor
$f_{\theta^{0}}:\mathcal{X}^{0}\rightarrow\mathcal{Y}^{0}$ parameterized with
$\theta^{0}$, which has been trained with the data
$\\{x_{n}^{0},y_{n}^{0}\\}_{n=1}^{N^{0}}$ in an initial source domain
$\mathcal{D}^{0}=\\{\mathcal{X}^{0},\mathcal{Y}^{0}\\}$, where
$x_{n}^{0}\in\mathbb{R}^{H\times W}$ and $y_{n}^{0}\in\mathbb{R}^{H\times W}$
are the paired image slice and its segmentation mask with the height of $H$
and width of $W$, respectively. There are $T$ consecutive evolving stages with
heterogeneous target domains
$\mathcal{D}^{t}=\\{\mathcal{X}^{t},\mathcal{S}^{t}\\}_{t=1}^{T}$, each with
the paired slice set $\\{x_{n}^{t}\\}_{n=1}^{N^{t}}\in\mathcal{X}^{t}$ and the
current stage label set $\\{s_{n}^{t}\\}_{n=1}^{N^{t}}\in\mathcal{S}^{t}$,
where $x_{n}^{t},s_{n}^{t}\in\mathbb{R}^{H\times W}$. Due to heterogeneous
domain shifts, $\mathcal{X}^{t}$ from different sites or modalities follows
diverse distributions across all $T$ stages. Due to incremental anatomical
structures, the overall label space, across the previous $t$ stages,
$\mathcal{Y}^{t}$ is expanded from $\mathcal{Y}^{t-1}$ with the additional
annotated structures $\mathcal{S}^{t}$ in stage $t$, i.e.,
$\mathcal{Y}^{t}=\mathcal{Y}^{t-1}\cup\mathcal{S}^{t}=\mathcal{Y}^{0}\cup\mathcal{S}^{1}\cdots\cup\mathcal{S}^{t}$.
We are targeting to learn
$f_{\theta^{T}}:\\{\mathcal{X}^{t}\\}_{t=1}^{T}\rightarrow\mathcal{Y}^{T}$
that performs well on all $\\{\mathcal{X}^{t}\\}_{t=1}^{T}$ for delineating
all of the structures $\mathcal{Y}^{T}$ seen in $T$ stages.

Figure 1: Illustration of one layer in our proposed divergence-aware decoupled
dual-flow module guided with cBRN for our cross-MR-modality HSI task, i.e.,
subject-independent (CoreT with T1) $\rightarrow$ (EnhT with T2) $\rightarrow$
(ED with FLAIR). Notably, we do not require the dual-flow or cBRN, for the
initial segmentor.

###  2.1 cBRN guided divergence-aware decoupled dual-flow

To alleviate the forgetting through parameter overwriting, caused by both new
structures and data shift, we propose a D3F module for flexible decoupling and
integration of old and new knowledge.

Specifically, we duplicate the convolution in each layer initialized with the
previous model $f_{\theta^{t-1}}$ to form two branches as in [13, 25, 34]. The
first $rigidity$ branch $f_{\theta^{t}}^{r}$ is fixed at the stage $t$ to keep
the old knowledge we have learned. In contrast, the extended $plasticity$
branch $f_{\theta^{t}}^{p}$ is expected to be adaptively updated to learn the
new task in $\mathcal{D}^{t}$. At the end of current training stage $t$, we
can flexibly integrate the convolutions in two branches, i.e.,
$\\{W_{t}^{r},b_{t}^{r}\\}$ and $\\{W_{t}^{p},b_{t}^{p}\\}$ to
$\\{W_{t+1}^{r}=\frac{W_{t}^{r}+W_{t}^{p}}{2},b_{t+1}^{r}=\frac{b_{t}^{r}+b_{t}^{p}}{2}\\}$
with the model re-parameterization [4]. In fact, the dual-flow model can be
regarded as an implicit ensemble scheme [9] to integrate multiple sub-modules
with a different focus. In addition, as demonstrated in [6], the fixed modules
will regularize the learnable modules to act as the fixed one. Thus, the
plasticity modules can also be implicitly encouraged to keep the previous
knowledge along with its HSI learning.

However, under the domain shift, it can be sub-optimal to directly average the
parameters, since $f_{\theta^{t}}^{r}$ may not perform well to predict
$\mathcal{Y}^{t-1}$ on $\mathcal{X}^{t}$. It has been demonstrated that batch
statistics adaptation plays an important role in domain generalizable model
training [22]. Therefore, we propose a continual batch renormalization (cBRN)
to mitigate the feature statistics divergence between each training batch at a
specific stage and the life-long global data distribution.

Of note, as a default block in the modern convolutional neural networks (CNN)
[8, 37], batch normalization (BN) [11] normalizes the input feature of each
CNN channel $z\in\mathbb{R}^{H_{c}\times W_{c}}$ with its batch-wise
statistics, e.g., mean $\mu_{B}$ and standard deviation $\sigma_{B}$, and
learnable scaling and shifting factors $\\{\gamma,\beta\\}$ as
$\tilde{z}_{i}=\frac{z_{i}-\mu_{B}}{\sigma_{B}}\cdot\gamma+\beta,$ where $i$
indexes the spatial position in $\mathbb{R}^{H_{c}\times W_{c}}$. BN assumes
that the same mini-batch training and testing distribution [10], which does
not hold in HSI. Simply enforcing the same statistics across domains as [5,
33, 30] can weaken the model expressiveness [36].

The recent BRN [10] proposes to rectify the data shift between each batch and
the dataset by using the moving average ${\mu}$ and $\sigma$ along with the
training:

$\displaystyle\mu=(1-\eta)\cdot\mu+\eta\cdot\mu_{B},~{}~{}\sigma=(1-\eta)\cdot\sigma+\eta\cdot\sigma_{B},$
(1)

where $\eta\in[0,1]$ is applied to balance the global statistics and the
current batch. In addition, $\gamma=\frac{\sigma_{B}}{\sigma}$ and
$\beta=\frac{\mu_{B}-\mu}{\sigma}$ are used in both training and testing.
Therefore, BRN renormalizes $\tilde{z}_{i}=\frac{z_{i}-\mu}{\sigma}$ to
highlight the dependency on the global statistics $\\{\mu,\sigma\\}$ in
training for a more generalizable model, while limited to the static learning.

In this work, we further explore the potential of BRN in the continuously
evolving HSI task to be general for all of domains involved. Specifically, we
extend BRN to cBRN across multiple consecutive stages by updating
$\\{\mu_{c},\sigma_{c}\\}$ along with all stages of training, which is
transferred as shown in Fig. 1. The conventional BN also inherits
$\\{\mu,\sigma\\}$ for testing, while not being used in training [11]. At the
stage $t$, $\mu_{c}$ and $\sigma_{c}$ are succeeded from $t-1$ stage, and are
updated with the current batch-wise $\\{\mu_{B}^{r},\sigma_{B}^{r}\\}$ and
$\\{\mu_{B}^{p},\sigma_{B}^{p}\\}$ in rigidity and plasticity branches:

$\displaystyle\mu_{c}=(1-\eta)\cdot\mu_{c}+\eta\cdot\frac{1}{2}\\{\mu_{B}^{r}+\mu_{B}^{p}\\},~{}~{}\sigma_{c}=(1-\eta)\cdot\sigma_{c}+\eta\cdot\frac{1}{2}\\{\sigma_{B}^{r}+\sigma_{B}^{p}\\}.$
(2)

For testing, the two branches in final model $f_{\theta^{T}}$ can be merged
for the lightweight implementation:

$\displaystyle\tilde{z}=\frac{W_{T}^{r}z+b_{T}^{r}+\mu_{c}}{2\sigma_{c}}+\frac{W_{T}^{p}z+b_{T}^{p}+\mu_{c}}{2\sigma_{c}}=\frac{W_{T}^{r}+W_{T}^{p}}{2\sigma_{c}}z+\frac{b_{T}^{r}+b_{T}^{p}-2\mu_{c}}{2\sigma_{c}}=\hat{W}z+\hat{b}.$
(3)

Therefore, $f_{\theta}^{T}$ does not introduce additional parameters for
deployment.

Figure 2: Illustration of the proposed HSI pseudo-label distillation with MMD

###  2.2 HSI pseudo-label distillation with momentum MixUp decay

The training of our developed $f_{\theta^{t}}$ with D3F is supervised with the
previous model $f_{\theta^{t-1}}$ and current stage data
$\\{x_{n}^{t},s_{n}^{t}\\}_{n=1}^{N^{t}}$. In conventional class incremental
learning, the knowledge distillation [31] is widely used to construct the
combined label ${y}_{n}^{t}\in\mathbb{R}^{H\times W}$ by adding $s_{n}^{t}$
and the prediction of $f_{\theta^{t-1}}(x_{n}^{t})$. Then, $f_{\theta^{t}}$
can be optimized by the training pairs of
$\\{x_{n}^{t},{y}_{n}^{t}\\}_{n=1}^{N^{t}}$. However, with heterogeneous data
in different stages, $f_{\theta^{t-1}}(x_{n}^{t})$ can be highly unreliable.
Simply using it as ground truth cannot guide the correct knowledge transfer.

In this work, we construct a complementary pseudo-label
$\hat{y}_{n}^{t}\in\mathbb{R}^{H\times W}$ with a MixUp decay scheme to
adaptively exploit the knowledge in the old segmentor for the progressively
learned new segmentor. In the initial training epochs, $f_{\theta^{t-1}}$
could be a more reliable supervision signal, while we would expect
$f_{\theta^{t}}$ can learn to perform better on predicting
$\mathcal{Y}^{t-1}$. Of note, even with the rigidity branch, the integrated
network can be largely distracted by the plasticity branch in the initial
epochs. Therefore, we propose to dynamically adjust their importance in
constructing pseudo-label along with the training progress. Specifically, we
MixUp the predictions of $f_{\theta^{t-1}}$ and $f_{\theta^{t}}$ w.r.t.
$\mathcal{Y}^{t-1}$, i.e., $f_{\theta^{t}}(\cdot)[:t-1]$, and control their
pixel-wise proportion for the pseudo-label $\hat{y}_{n}^{t}$ with MMD:

$\displaystyle\hat{y}_{n:i}^{t}=\\{\lambda
f_{\theta^{t-1}}(x_{n:i}^{t})+(1-\lambda)f_{\theta^{t}}(x_{n:i}^{t})[:t-1]\\}\cup
s_{n:i}^{t},~{}\lambda=\lambda^{0}\text{exp}(-I),$ (4)

where $i$ indexes each pixel, and $\lambda$ is the adaptation momentum factor
with the exponential decay of iteration $I$. $\lambda^{0}$ is the initial
weight of $f_{\theta^{t-1}}(x_{n:i}^{t})$, which is empirically set to 1 to
constrain $\lambda\in(0,1]$. Therefore, the weight of old model prediction can
be smoothly decreased along with the training, and
$f_{\theta^{t}}(x_{n:i}^{t})$ gradually represents the target data for the old
classes in $[:t-1]$. Of note, we have ground-truth of new structure
$s_{n:i}^{t}$ under HSI scenarios [5, 33, 30, 18]. We calculate the cross-
entropy loss $\mathcal{L}_{CE}$ with the pseudo-label $\hat{y}_{n:i}^{t}$ as
self-training [15, 38].

In addition to the old knowledge inherited in $f_{\theta^{t-1}}$, we propose
to explore unsupervised learning protocols to stabilize the initial training.
We adopt the widely used self-entropy (SE) minimization [7] as a simple add-on
training objective. Specifically, we have the slice-level segmentation SE,
which is the averaged entropy of the pixel-wise softmax prediction as
$\mathcal{L}_{SE}=\mathbb{E}_{i}\\{-{f_{\theta^{t}}(x_{n:i}^{t})\text{log}f_{\theta^{t}}(x_{n:i}^{t})}\\}.$
In training, the overall optimization loss is formulated as follows:

$\displaystyle\mathcal{L}=\mathcal{L}_{CE}(\hat{y}_{n:i}^{t},f_{\theta^{t}}(x_{n:i}^{t}))+\alpha\mathcal{L}_{SE}(f_{\theta^{t}}(x_{n:i}^{t})),~{}~{}\alpha=\frac{I_{max}-I}{I_{max}}{\alpha^{0}},$
(5)

where $\alpha$ is used to balance our HSI distillation and SE minimization
terms, and $I_{max}$ is the scheduled iteration. Of note, strictly minimizing
the SE can result in a trivial solution of always predicting a one-hot
distribution [7], and a linear decreasing of $\alpha$ is usually applied,
where $\lambda^{0}$ and $\alpha^{0}$ are reset in each stage.

Table 1: Numerical comparisons and ablation studies of the cross-subset brain
tumor HSI segmentation task

Method | Data shift | Dice similarity coefficient (DSC) [%] $\uparrow$ | Hausdorff distance (HD)[mm] $\downarrow$
---|---|---|---
consideration |  Mean |  CoreT |  EnhT |  ED |  Mean |  CoreT |  EnhT |  ED
PLOP [5] | $\times$ | 59.83$\pm 0.131$ | 45.50 | 57.39 | 76.59 | 19.2$\pm$0.14 | 22.0 | 19.8 | 15.9
MargExcIL [18] | $\times$ | 60.49$\pm$0.127 | 48.37 | 56.28 | 76.81 | 18.9$\pm$0.11 | 21.4 | 19.8 | 15.5
UCD [30] | $\times$ | 61.84$\pm$0.129 | 49.23 | 58.81 | 77.48 | 19.0$\pm$0.15 | 21.8 | 19.4 | 15.7
HSI-MMD | $\surd$ | 66.87$\pm$0.126 | 59.42 | 61.26 | 79.93 | 16.8$\pm$0.13 | 18.5 | 17.8 | 14.2
HSI-D3F | $\surd$ | 67.18$\pm$0.118 | 60.18 | 63.09 | 78.26 | 16.7$\pm$0.14 | 18.0 | 17.5 | 14.5
HSI-cBRN | $\surd$ | 68.07$\pm$0.121 | 61.52 | 63.45 | 79.25 | 16.3$\pm$0.14 | 17.8 | 17.3 | 13.8
HSI | $\surd$ | 69.44$\pm$0.119 | 63.79 | 64.71 | 79.81 | 15.7$\pm$0.12 | 16.7 | 16.9 | 13.6
Joint Static |  $\surd$(upper bound) | 73.98$\pm$0.117 | 71.14 | 68.35 | 82.46 | 15.0$\pm$0.13 | 15.7 | 16.2 | 13.2

##  3 Experiments and Results

We carried out two evaluation settings using the BraTS2018 database [1],
including cross-subset (relatively small domain shift) and cross-modality
(relatively large domain shift) tasks. The BraTS2018 database is a continually
evolving database [1] with a total of 285 glioblastoma or low-grade gliomas
subjects, comprising three consecutive subsets, i.e., 30 subjects from
BraTS2013 [26], 167 subjects from TCIA [3], and 88 subjects from CBICA [1].
Notably, these three subsets were collected from different clinical sites,
vendors, or populations [1]. Each subject has T1, T1ce, T2, and FLAIR MRI
volumes with voxel-wise labels for the tumor core (CoreT), the enhancing tumor
(EnhT), and the edema (ED).

We incrementally learned CoreT, EnhT, and ED structures throughout three
consecutive stages, each following different data distributions. We used
subject-independent 7/1/2 split for training, validation, and testing. For a
fair comparison, we adopted the ResNet-based 2D nnU-Net backbone with BN as in
[12] for all of the methods and all stages used in this work.

Figure 3: Segmentation examples in $t=1$ and $t=2$ in the cross-subset brain
tumor HSI segmentation task.

###  3.1 Cross-subset structure incremental evolving

In our cross-subset setting, three structures were sequentially learned across
three stages: (CoreT with BraTS2013) $\rightarrow$ (EnhT with TCIA)
$\rightarrow$ (ED with CBICA). Of note, we used a CoreT segmentator trained
with BraTS2013 as our off-the-shelf segmentor in $t=0$. Testing involved all
subsets and anatomical structures. We compared our framework with the three
typical structure-incremental (SI-only) segmentation methods, e.g., PLOP [5],
MargExcIL [18], and UCD [30], which cannot address the heterogeneous data
across stages. As tabulated in Table 1, PLOP [5] with additional feature
statistic constraints has lower performance than MargExcIL [18], since the
feature statistic consistency was not held in HSI scenarios. Of note, the
domain-incremental methods [17, 32] cannot handle the changing output space.
Our proposed HSI framework outperformed SI-only methods [5, 18, 30] with
respect to both DSC and HD, by a large margin. For the anatomical structure
CoreT learned in $t=0$, the difference between our HSI and these SI-only
methods was larger than 10% DSC, which indicates the data shift related
forgetting lead to a more severe performance drop in the early stages. We set
$\eta=0.01$ and $\ alpha^{0}=10$ according to the sensitivity study in the
supplementary material.

For the ablation study, we denote HSI-D3F as our HSI without the D3F module,
simply fine-tuning the model parameters. HSI-cBRN used dual-flow to avoid
direct overwriting, while the model was not guided by cBRN for more
generalized prediction on heterogeneous data. As shown in Table 1, both the
dual-flow and cBRN improve the performance. Notably, the dual-flow model with
flexible re-parameterization was able to alleviate the overwriting, while our
cBRN was developed to deal with heterogeneous data. In addition, HSI-MMD
indicates our HSI without the momentum MixUp decay in pseudo-label
construction, i.e., simply regarding the prediction of
$f_{\theta^{t-1}}(x^{t})$ is ground truth for $\mathcal{Y}^{t-1}$. However,
$f_{\theta^{t-1}}(x^{t})$ can be quite noisy, due to the low quantification
performance of early stage structures, which can be aggravated in the case of
the long-term evolving scenario. Of note, the pseudo-label construction is
necessary as in [5, 18, 30]. We also provide the qualitative comparison with
SI-only methods and ablation studies in Fig. 3.

Table 2: Numerical comparisons and ablation studies of the cross-modality
brain tumor HSI segmentation task

Method | Data shift | Dice similarity coefficient (DSC) [%] $\uparrow$ | Hausdorff distance (HD)[mm] $\downarrow$
---|---|---|---
consideration |  Mean |  CoreT |  EnhT |  ED |  Mean |  CoreT |  EnhT |  ED
PLOP [5] | $\times$ | 39.58$\pm$0.231 | 13.84 | 38.93 | 65.98 | 30.7$\pm$0.26 | 48.1 | 25.4 | 18.7
MargExcIL [18] | $\times$ | 42.84$\pm$0.189 | 19.56 | 41.56 | 67.40 | 29.1$\pm$0.28 | 46.7 | 22.1 | 18.6
UCD [30] | $\times$ | 44.67$\pm$0.214 | 21.39 | 45.28 | 67.35 | 29.4$\pm$0.32 | 46.2 | 23.6 | 18.4
HSI-MMD | $\surd$ | 59.81$\pm$0.207 | 51.63 | 53.82 | 73.97 | 19.4$\pm$0.26 | 21.6 | 20.5 | 16.2
HSI-D3F | $\surd$ | 60.81$\pm$0.195 | 53.87 | 55.42 | 73.15 | 19.2$\pm$0.21 | 21.4 | 19.9 | 16.2
HSI-cBRN | $\surd$ | 61.87$\pm$0.180 | 54.90 | 56.62 | 74.08 | 18.5$\pm$0.25 | 20.1 | 19.5 | 16.0
HSI | $\surd$ | 64.15$\pm$0.205 | 58.11 | 59.51 | 74.83 | 17.7$\pm$0.29 | 18.9 | 18.6 | 15.8
Joint Static |  $\surd$ (upper bound) | 70.64$\pm$0.184 | 67.48 | 65.75 | 78.68 | 16.7$\pm$0.26 | 17.2 | 17.8 | 15.1

###  3.2 Cross-modality structure incremental evolving

In our cross-modality setting, three structures were sequentially learned
across three stages: (CoreT with T1) $\rightarrow$ (EnhT with T2)
$\rightarrow$ (ED with T2 FLAIR). Of note, we used the CoreT segmentator
trained with T1 modality as our off-the-shelf segmentor in $t=0$. Testing
involved all MRI modalities and all structures. With the hyperparameter
validation, we empirically set $\eta=0.01$ and $\alpha^{0}=10$.

In Table 2, we provide quantitative evaluation results. We can see that our
HSI framework outperformed SI-only methods [5, 18, 30] consistently. The
improvement can be even larger, compared with the cross-subset task, since we
have much more diverse input data in the cross-modality setting. Catastrophic
forgetting can be severe, when we use SI-only method for predicting early
stage structures, e.g., CoreT. We also provide the ablation study with respect
to D3F, cBRN, and MMD in Table 2. The inferior performance of HSI-D3F/cBRN/MMD
demonstrates the effectiveness of these modules for mitigating domain shifts.

##  4 Conclusion

This work proposed an HSI framework under a clinically meaningful scenario, in
which clinical databases are sequentially constructed from different
sites/imaging protocols with new labels. To alleviate the catastrophic
forgetting alongside continuously varying structures and data shifts, our HSI
resorted to a D3F module for learning and integrating old and new knowledge
nimbly. In doing so, we were able to achieve divergence awareness with our
cBRN-guided model adaptation for all the data involved. Our framework was
optimized with a self-entropy regularized HSI pseudo-label distillation scheme
with MMD to efficiently utilize the previous model in different types of MRI
data. Our framework demonstrated superior segmentation performance in learning
new anatomical structures from cross-subset/modality MRI data. It was
experimentally shown that a large improvement in learning anatomic structures
was observed.

## Acknowledgements

This work is supported by NIH R01DC018511, R01DE027989, and P41EB022544. The
authors would like to thank Dr. Jonghyun Choi for his valuable insights and
helpful discussions.

## References

  * [1] Bakas, S., Reyes, M., Jakab, A., Bauer, S., Rempfler, M., Crimi, A., Shinohara, R.T., Berger, C., Ha, S.M., Rozycki, M., et al.: Identifying the best machine learning algorithms for brain tumor segmentation, progression assessment, and overall survival prediction in the brats challenge. arXiv:1811.02629 (2018)
  * [2] Chang, W.G., You, T., Seo, S., Kwak, S., Han, B.: Domain-specific batch normalization for unsupervised domain adaptation. In: CVPR. pp. 7354–7362 (2019)
  * [3] Clark, K., Vendt, B., Smith, K., Freymann, J., Kirby, J., Koppel, P., Moore, S., Phillips, S., Maffitt, D., Pringle, M., et al.: The cancer imaging archive (tcia): maintaining and operating a public information repository. Journal of digital imaging 26(6), 1045–1057 (2013)
  * [4] Ding, X., Zhang, X., Ma, N., Han, J., Ding, G., Sun, J.: Repvgg: Making vgg-style convnets great again. In: CVPR. pp. 13733–13742 (2021)
  * [5] Douillard, A., Chen, Y., Dapogny, A., Cord, M.: Plop: Learning without forgetting for continual semantic segmentation. In: CVPR. pp. 4040–4050 (2021)
  * [6] Fu, S., Li, Z., Liu, Z., Yang, X.: Interactive knowledge distillation for image classification. Neurocomputing 449, 411–421 (2021)
  * [7] Grandvalet, Y., Bengio, Y.: Semi-supervised learning by entropy minimization. In: NeurIPS (2005)
  * [8] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: CVPR (2016)
  * [9] Huang, G., Sun, Y., Liu, Z., Sedra, D., Weinberger, K.Q.: Deep networks with stochastic depth. In: European Conference on Computer Vision. pp. 646–661 (2016)
  * [10] Ioffe, S.: Batch renormalization: Towards reducing minibatch dependence in batch-normalized models. NeurIPS 30 (2017)
  * [11] Ioffe, S., Szegedy, C.: Batch normalization: Accelerating deep network training by reducing internal covariate shift. In: ICML. pp. 448–456. PMLR (2015)
  * [12] Isensee, F., Jaeger, P.F., Kohl, S.A., Petersen, J., Maier-Hein, K.H.: nnu-net: a self-configuring method for deep learning-based biomedical image segmentation. Nature methods 18(2), 203–211 (2021)
  * [13] Kanakis, M., Bruggemann, D., Saha, S., Georgoulis, S., Obukhov, A., Gool, L.V.: Reparameterizing convolutions for incremental multi-task learning without task interference. In: European Conference on Computer Vision. pp. 689–707 (2020)
  * [14] Kim, D., Bae, J., Jo, Y., Choi, J.: Incremental learning with maximum entropy regularization: Rethinking forgetting and intransigence. arXiv:1902.00829 (2019)
  * [15] Kim, K., Ji, B., Yoon, D., Hwang, S.: Self-knowledge distillation: A simple way for better generalization. arXiv:2006.12000 (2020)
  * [16] Kundu, J.N., Venkatesh, R.M., Venkat, N., Revanur, A., Babu, R.V.: Class-incremental domain adaptation. In: European Conference on Computer Vision. pp. 53–69 (2020)
  * [17] Li, K., Yu, L., Heng, P.A.: Domain-incremental cardiac image segmentation with style-oriented replay and domain-sensitive feature whitening. TMI (2022)
  * [18] Liu, P., Wang, X., Fan, M., Pan, H., Yin, M., Zhu, X., Du, D., Zhao, X., Xiao, L., Ding, L., Zhou, S.: Learning incrementally to segment multiple organs in a ct image. MICCAI (2022)
  * [19] Liu, X., Prince, J.L., Xing, F., Zhuo, J., Timothy, R., Stone, M., Fakhri, G.E., Woo, J.: Attentive continuous generative self-training for unsupervised domain adaptive medical image translation. Medical Image Analysis (2023)
  * [20] Liu, X., Xing, F., El Fakhri, G., Woo, J.: Memory consistent unsupervised off-the-shelf model adaptation for source-relaxed medical image segmentation. Medical Image Analysis 83, 102641 (2023)
  * [21] Liu, X., Xing, F., Shusharina, N., Lim, R., Jay Kuo, C.C., El Fakhri, G., Woo, J.: Act: Semi-supervised domain-adaptive medical image segmentation with asymmetric co-training. In: MICCAI (2022)
  * [22] Liu, X., Xing, F., Yang, C., El Fakhri, G., Woo, J.: Adapting off-the-shelf source segmenter for target medical image segmentation. In: MICCAI. pp. 549–559 (2021)
  * [23] Liu, X., Xing, F., You, J., Lu, J., Kuo, C.C.J., El Fakhri, G., Woo, J.: Subtype-aware dynamic unsupervised domain adaptation. IEEE TNNLS (2022)
  * [24] Liu, X., Yoo, C., Xing, F., Oh, H., El Fakhri, G., Kang, J.W., Woo, J., et al.: Deep unsupervised domain adaptation: a review of recent advances and perspectives. APSIPA Transactions on Signal and Information Processing 11(1) (2022)
  * [25] Liu, Y., Schiele, B., Sun, Q.: Adaptive aggregation networks for class-incremental learning. In: CVPR. pp. 2544–2553 (2021)
  * [26] Menze, B.H., Jakab, A., Bauer, S., Kalpathy-Cramer, J., Farahani, K., Kirby, J., Burren, Y., Porz, N., Slotboom, J., Wiest, R., et al.: The multimodal brain tumor image segmentation benchmark (BRATS). TMI 34(10), 1993–2024 (2014)
  * [27] Ozdemir, F., Fuernstahl, P., Goksel, O.: Learn the new, keep the old: Extending pretrained models with new anatomy and images. In: MICCAI. pp. 361–369 (2018)
  * [28] Shusharina, N., Söderberg, J., Edmunds, D., Löfman, F., Shih, H., Bortfeld, T.: Automated delineation of the clinical target volume using anatomically constrained 3d expansion of the gross tumor volume. Radiotherapy and Oncology (2020)
  * [29] Tajbakhsh, N., Jeyaseelan, L., Li, Q., Chiang, J.N., Wu, Z., Ding, X.: Embracing imperfect datasets: A review of deep learning solutions for medical image segmentation. Medical Image Analysis 63, 101693 (2020)
  * [30] Yang, G., Fini, E., Xu, D., Rota, P., Ding, M., Nabi, M., Alameda-Pineda, X., Ricci, E.: Uncertainty-aware contrastive distillation for incremental semantic segmentation. IEEE Transactions on Pattern Analysis and Machine Intelligence (2022)
  * [31] Yin, H., Molchanov, P., Alvarez, J.M., Li, Z., Mallya, A., Hoiem, D., Jha, N.K., Kautz, J.: Dreaming to distill: Data-free knowledge transfer via deepinversion. In: CVPR. pp. 8715–8724 (2020)
  * [32] You, C., Xiang, J., Su, K., Zhang, X., Dong, S., Onofrey, J., Staib, L., Duncan, J.S.: Incremental learning meets transfer learning: Application to multi-site prostate mri segmentation. arXiv:2206.01369 (2022)
  * [33] Yu, L., Liu, X., Van de Weijer, J.: Self-training for class-incremental semantic segmentation. IEEE Transactions on Neural Networks and Learning Systems (2022)
  * [34] Zhang, C.B., Xiao, J.W., Liu, X., Chen, Y.C., Cheng, M.M.: Representation compensation networks for continual semantic segmentation. In: CVPR (2022)
  * [35] Zhang, H., Zhang, Y., Jia, K., Zhang, L.: Unsupervised domain adaptation of black-box source models. arXiv:2101.02839 (2021)
  * [36] Zhang, J., Qi, L., Shi, Y., Gao, Y.: Generalizable semantic segmentation via model-agnostic learning and target-specific normalization. arXiv:2003.12296 (2020)
  * [37] Zhou, X.Y., Yang, G.Z.: Normalization in training u-net for 2-D biomedical semantic segmentation. IEEE Robotics and Automation Letters 4(2), 1792–1799 (2019)
  * [38] Zou, Y., Yu, Z., Liu, X., Kumar, B., Wang, J.: Confidence regularized self-training. In: ICCV. pp. 5982–5991 (2019)


marginparsep has been altered.
topmargin has been altered.
marginparwidth has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

On-device federated learning with Flower

Anonymous Authors1

###### Abstract

Federated Learning (FL) allows edge devices to collaboratively learn a shared
prediction model while keeping their training data on the device, thereby
decoupling the ability to do machine learning from the need to store data in
the cloud. Despite the algorithmic advancements in FL, the support for on-
device training of FL algorithms on edge devices remains poor. In this paper,
we present an exploration of on-device FL on various smartphones and embedded
devices using the Flower framework. We also evaluate the system costs of on-
device FL and discuss how this quantification could be used to design more
efficient FL algorithms.

††footnotetext:  1Anonymous Institution, Anonymous City, Anonymous Region,
Anonymous Country. Correspondence to: Anonymous Author
<EMAIL_ADDRESS>
Preliminary work. Under review by the Machine Learning and Systems (MLSys)
Conference. Do not distribute.

##  1 Introduction

We have seen remarkable progress in enabling the execution of deep learning
models on mobile and embedded devices to infer user contexts and behaviors
Warden & Situnayake (2019); Fromm et al. (2018); Chowdhery et al. (2019);
Malekzadeh et al. (2019); Lee et al. (2019). This has been powered by the
increasing computational abilities of edge devices as well as novel software
optimizations to enable cloud-scale models to run on resource-constrained
devices. However, when it comes to the training of these edge-focused models,
a working assumption has been that the models will be trained centrally in the
cloud, using training data aggregated from several users.

Federated Learning (FL) McMahan et al. (2017) aims to enable distributed edge
devices (or users) to collaboratively _train_ a shared prediction model while
keeping their personal data private. At a high level, this is achieved by
repeating three basic steps: i) local parameters update to a shared prediction
model on each edge device, ii) sending the local parameter updates to a
central server for aggregation, and iii) receiving the aggregated model back
for the next round of local updates.

A major bottleneck to FL research is the paucity of frameworks that support
federated training of workloads on mobile and embedded devices. While several
frameworks including Tensorflow Federated Google (2020); Abadi et al. (2016)
and LEAF Caldas et al. (2018) enable simulation of FL clients, they cannot be
used to understand the training dynamics and compute the system costs of FL on
edge devices. Edge devices exhibit significant heterogeneity in their software
stack, compute capabilities, and network bandwidth. All these system-related
factors, in combination with the choice of the FL clients and parameter
aggregation algorithms, can impact the accuracy and training time of models
trained in a federated setting.

In this paper, we present our exploration of on-device training of FL
workloads on Android smartphones and Nvidia Jetson series embedded devices,
using the Flower framework Beutel et al. (2020). Flower offers a stable
implementation of the core components of an FL system, and provides higher-
level abstractions to enable researchers to experiment and implement new ideas
on top of a reliable stack. We first demonstrate how we use the language-,
platform- and ML framework-agnostic capabilities of Flower to support on-
device training of FL workloads on edge devices with heterogeneous hardware
and software stacks. We then deploy these FL clients on various embedded
devices as well as on Android smartphones hosted in the Amazon AWS Device Farm
(https://aws.amazon.com/device-farm/). Finally, we present an evaluation to
compute various system-related metrics of FL and highlight how this
quantification could lead to the design of more efficient FL algorithms.

##  2 Related Work

McMahan et al. (2017) introduced the basic federated averaging (FedAvg)
algorithm and evaluated it in terms of communication efficiency. The
optimization of distributed training with and without federated concepts has
been covered from many angles Jia et al. (2018); Chahal et al. (2018).
Bonawitz et al. (2019) detail the design of a large-scale Google-internal FL
system. TFF Google (2020), PySyft Ryffel et al. (2018), LEAF Caldas et al.
(2018), FedML He et al. (2020) are other open-source frameworks that support
research and experimentation of FL workloads.

Given their relative hardware limitations, most of the works involving machine
learning on mobile devices were originally aimed at adapting existing models
to specific latency and storage constraints. For this purpose, optimized
versions of TensorFlow and PyTorch were developed David et al. (2020); Paszke
et al. (2019) and can now be considered mainstream. Federated Learning, on the
other hand, shifts the training burden from server to the client, which in
turn creates the need for developing the adequate supporting back-ends (e.g.
back-propagation) for low-power hardware as well. Recent works towards this
goal include using low-precision training Sun et al. (2020), controlled
updates of biases to reduce memory Cai et al. (2020), and early-exit models
that provide a trade-off between accuracy and compute Leontiadis et al.
(2021). In time, as hardware capabilities in edge devices improve, we expect
to see these advances being deployed by mainstream frameworks.

##  3 Primer on Flower

Figure 1: Flower framework architecture.

Flower is a novel client-agnostic federated learning framework. One of the
underlying design goals of Flower is to enable integrating with an inherently
heterogeneous and ever-evolving edge device landscape. There are multiple
dimensions of on-device heterogeneity, amongst them are operating systems,
machine learning frameworks, programming languages, connectivity, and hardware
accelerators.

The Flower core framework, shown in Figure 1, implements the infrastructure to
run these heterogeneous workloads at scale. On the server side, there are
three major components involved: the FL loop, the RPC server, and a (user
customizable) _Strategy_. Strategy here refers to the federated averaging
algorithms (e.g., FedAvg) used for aggregating the model parameters across
clients. Clients connect to the RPC server which is responsible for monitoring
these connections and for sending and receiving _Flower Protocol_ messages.
The FL loop is at the heart of the FL process: it orchestrates the learning
process and ensures that progress is made. It does not, however, make
decisions about _how_ to proceed, those decisions are delegated to the
currently configured _Strategy_ implementation.

A distinctive property of this architecture is a server which is unaware of
the nature of connected clients. This allows to train models across
heterogeneous client platforms and implementations, including workloads
comprised of different client-side ML frameworks. Furthermore, on-device
training and evaluation can be implemented in different programming languages,
a property especially important for research on mobile and embedded platforms.
These platforms often do not support Python, but rely on specific languages
(Java on Android, Swift on iOS) for idiomatic development, or native C/C++ for
some embedded devices. Flower achieves a fully language-agnostic interface by
offering protocol-level integration. The _Flower Protocol_ defines core
server-side messages such as _fit_ and _evaluate_ , which include the
(serialized) global model parameters and expect return messages from the
client that return either updated model parameters (or gradients) or
evaluation results. Each message contains additional user-customizable
metadata that allows the server to control on-device hyper-parameters, for
example, the number of on-device training epochs. Due to space constraints, we
refer the reader to Beutel et al. (2020) and https://flower.dev/ for more
details on Flower.

##  4 Flower Clients and on-device training

In this section, we describe two instances of on-device federated learning
with Flower. First, we present how Flower clients can be developed in Java and
deployed on Android phones in the AWS Device Farm for federated model
training. Next, we discuss the implementation of Flower clients in Python and
their deployment on heterogeneous embedded devices such as Nvidia Jetson
series and Raspberry Pi.

Figure 2: Flower Android client architecture.

###  4.1 Java Flower Clients for Android Smartphones

By design, Flower is language-agnostic and can work with any ML framework on
the FL client, which maximizes its ability to federate existing training
pipelines. However, it also means Flower inherits the limitations of these
frameworks that currently offer limited support for on-device training on
Android devices.

In the absence of a full-fledged model training library for Android, we employ
the TensorFlow Lite (TFLite) Model Personalization support on Android to
perform on-device federated learning. More specifically, as shown in Figure 2,
we first obtain a pre-trained _Base Model_ (e.g., MobileNetV2 without its top
layers). The parameters of this model are frozen (i.e., not updated during
training), and it is only used as a feature extractor. Next, we define a _Head
Model_ which corresponds to the task-specific classifier that we want to train
using federated learning. The input to the _Head Model_ are the features
extracted from the _Base Model_ , and its weights are randomly initialized.
Finally, we use the TFLite Transfer Convertor to port the Base and Head models
to TFLite and package them inside an Android application for on-device
training.

Inside the Android application, Flower Client is a class which coordinates
with the TFLite Model Personalization libraries and implements the three core
methods required for federated training with Flower, namely (i) get_weights
(.), which gets the current weights of the _Head Model_ to support server-side
aggregation requests, (ii) fit (.), which updates the parameters of the Head
Model through local training and (iii) evaluate (.), which computes test loss
on the local dataset and communicates it to the server.

The Flower Client spawns a background thread and sets up bi-directional
streaming RPC with the Flower server using the StreamObserver class. Upon
receiving messages from the Flower server, the background thread calls the
appropriate TFLite methods to update the parameters of the _Head Model_ using
the local training data (e.g., by optimizing a cross-entropy loss for
classification tasks). The updated parameters are sent to the Flower Server
where they are aggregated across all clients, before being sent back to the
clients for the next round of training.

###  4.2 Python Flower Clients on Embedded Devices

Contrary to Android, on-device training on Python-enabled embedded devices can
be performed using platform-specific ML frameworks (e.g., TensorFlow,
PyTorch). However, one of the open challenges to deploy on-device training
workloads arises due to the heterogeneity of computing resources present in
such devices. For example, NVIDIA Jetson devices provide GPU acceleration,
while Raspberry Pi is limited to CPU-only workloads. Implementing different
versions of an FL client for each target platform can quickly become
unpractical in real-world scenarios.

Flower overcomes this challenge by providing a platform-agnostic way of
writing FL clients. Figure 3 illustrates the design of Flower Clients for a
NVIDIA Jetson TX2 and a Raspberry Pi device. Despite the differences in these
two embedded platforms, we can use the same code to develop Flower Clients,
and run them inside a platform-specific Docker image, enabling the client code
to access platform-specific resources (e.g. GPU) when available.

More importantly, this platform-agnostic capability of Flower can be combined
with its language-agnostic capabilities (i.e., support for Python, Java, C++
clients) to enable on-device training of federated workloads in highly-
heterogeneous client setups, e.g., when the FL clients include Android phones,
Android watches, Raspberry Pis, ARM micro-controllers, Nvidia Jetson devices
etc.

Figure 3: Flower clients can easily be deployed to heterogeneous devices by
leveraging existing container-based frameworks (e.g. Docker) that interface
with the host’s hardware.

##  5 Evaluation

Our evaluation focuses on quantifying the system costs associated with running
FL on various edge devices. In doing so, we also explore how such
quantification could help FL developers to design novel algorithms that
effectively trade-off between system costs and FL accuracy.

Table 1: Android phones used from the AWS Device Farm Device Name | Type | OS Version
---|---|---
Google Pixel 4 | Phone | 10
Google Pixel 3 | Phone | 10
Google Pixel 2 | Phone | 9
Samsung Galaxy Tab S6 | Tablet | 9
Samsung Galaxy Tab S4 | Tablet | 8.1.0

Datasets. Two datasets are used in our evaluation, namely _CIFAR-10_ and
_Office-31_ Office31 (2020), both of which are examples of object recognition
datasets.

Deployment Setup. We run the Flower Server configured with the FedAvg strategy
and host it on a cloud virtual machine. Two sets of edge devices are used in
the evaluation: Android smartphones and Nvidia Jetson TX2 accelerators. To
scale our experiments to a reasonably large number of Android clients with
different OS versions, we deploy Flower Clients on the Amazon AWS Device Farm,
which enables testing applications on real Android devices accessed through
AWS. Table 1 list the smartphones from AWS Device Farm used in our evaluation.
Nvidia Jetson TX2 devices support full-fledged PyTorch – this means we could
successfully port existing PyTorch training pipelines to implement FL clients
on them.

Table 2: Flower supports implementation of FL clients on any device that has
on-device training support. Here we show various FL experiments on Android and
Nvidia Jetson devices.

| Local
---
Epochs (E)
Accuracy |  | Convergence
---
Time (mins)
| Energy
---
Consumed (kJ)
1 | 0.48 | 17.63 | 10.21
5 | 0.64 | 36.83 | 50.54
10 | 0.67 | 80.32 | 100.95
(a) Performance metrics with Nvidia Jetson TX2 as we vary the number of local
epochs. Number of clients $C$ is set to 10 and the model is trained for 40
rounds.

| Number
---
of Clients (C)
Accuracy |  | Convergence
---
Time (mins)
| Energy
---
Consumed (kJ)
4 | 0.84 | 30.7 | 10.4
7 | 0.85 | 31.3 | 19.72
10 | 0.87 | 31.8 | 28.0
(b) Performance metrics with Android clients as we vary the number of clients.
Local epochs $E$ is fixed to 5 in this experiment and the model is trained for
20 rounds.

System Costs of FL. In Table 2(b), we present various performance metrics
obtained on Nvidia TX2 and Android devices. First, we train a ResNet-18 model
on the CIFAR-10 dataset on 10 Nvidia TX2 clients. In Table 2(a), we vary the
number of local training epochs ($E$) performed on each client in a round of
FL. Our results show that choosing a higher $E$ results in better FL accuracy,
however it also comes at the expense of significant increase in total training
time and overall energy consumption across the clients. While the accuracy
metrics in Table 2(a) could have been obtained in a simulated setup,
quantifying the training time and energy costs on real clients would not have
been possible without a real on-device deployment. As reducing the energy and
carbon footprint of training ML models is a major challenge for the community,
Flower can assist researchers in choosing an optimal value of $E$ to obtain
the best trade-off between accuracy and energy consumption.

Next, we train a 2-layer DNN classifier (_Head Model_) on top of a pre-trained
MobileNetV2 _Base Model_ on Android clients for the Office-31 dataset. In
Table 2(b), we vary the number of Android clients ($C$) participating in FL,
while keeping the local training epochs ($E$) on each client fixed to 5. We
observe that by increasing the number of clients, we can train a more accurate
object recognition model. Intuitively, as more clients participate in the
training, the model gets exposed to more diverse training examples, thereby
increasing its generalizability to unseen test samples. However, this accuracy
gain comes at the expense of high energy consumption – the more clients we
use, the higher the total energy consumption of FL. Again, based on this
analysis obtained using Flower, researchers can choose an appropriate number
of clients to find a balance between accuracy and energy consumption.

Table 3: Effect of computational heterogeneity on FL training times. Using Flower, we can compute a hardware-specific cutoff $\tau$ (in minutes) for each processor, and find a balance between FL accuracy and training time. $\tau=0$ indicates no cutoff time. |  | GPU
---
($\tau$ = 0)
| CPU
---
($\tau$ = 0)
| CPU
---
($\tau$ = 2.23)
| CPU
---
($\tau$ = 1.99)
Accuracy | 0.67 | 0.67 | 0.66 | 0.63
| Training
---
time (mins)
80.32 |  | 102
---
(1.27x)
| 89.15
---
(1.11x)
| 80.34
---
(1.0x)

Computational Heterogeneity across Clients. FL clients could have vastly
different computational capabilities. While newer smartphones are now equipped
with GPUs, other devices may have a much less powerful processor. How does
this computational heterogeneity impact FL?

For this experiment, we use Nvidia Jetson TX2 as the client device, which has
one Pascal GPU and six CPU cores. We repeat the experiment shown in Table
2(a), but instead of using the embedded GPU for training, we train the
ResNet-18 model on a CPU. In Table 3, we show that CPU training with local
epochs $E=10$ would take $1.27\times$ more time to obtain the same accuracy as
the GPU training. This implies that even a single client device with low
compute resources (e.g., a CPU) can become a bottleneck and significantly
increase the FL training time.

Once we obtain this quantification of computational heterogeneity, we can
design better federated optimization algorithms. As an example, we implement a
modified version of FedAvg where each client device is assigned a cutoff time
($\tau$) after which it must send its model parameters to the server,
irrespective of whether it has finished its local epochs or not. This strategy
has parallels with the FedProx algorithm Li et al. (2018) which also accepts
partial results from clients. However, the key advantage of using Flower is
that we can compute and assign a processor-specific cutoff time for each
client. For example, on average it takes 1.99 minutes to complete an FL round
on the TX2 GPU. If we set the same time as a cutoff for CPU clients
($\tau=1.99$ mins) as shown in Table 3, we obtain the same convergence time as
GPU, at the expense of 3% accuracy drop. With $\tau=2.23$, a better balance
between accuracy and training time could be obtained on a CPU.

##  6 Conclusion

We presented our exploration of federated training of models on mobile and
embedded devices by leveraging the language-, framework- and platform-agnostic
capabilities of the Flower framework. We also presented early results
quantifying the system costs of FL on the edge and its implications for the
design of more efficient FL algorithms. Although our work is early stage, we
hope it will trigger discussions at the workshop related to on-device training
for FL and beyond.

## References

  * Abadi et al. (2016) Abadi, M., Barham, P., Chen, J., Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., Kudlur, M., Levenberg, J., Monga, R., Moore, S., Murray, D. G., Steiner, B., Tucker, P., Vasudevan, V., Warden, P., Wicke, M., Yu, Y., and Zheng, X.  Tensorflow: A system for large-scale machine learning.  In _12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16)_ , pp. 265–283, 2016.
  * Beutel et al. (2020) Beutel, D. J., Topal, T., Mathur, A., Qiu, X., Parcollet, T., de Gusmão, P. P. B., and Lane, N. D.  Flower: A friendly federated learning research framework.  _CoRR_ , abs/2007.14390, 2020.  URL https://arxiv.org/abs/2007.14390.
  * Bonawitz et al. (2019) Bonawitz, K., Eichner, H., Grieskamp, W., Huba, D., Ingerman, A., Ivanov, V., Kiddon, C. M., Konečný, J., Mazzocchi, S., McMahan, B., Overveldt, T. V., Petrou, D., Ramage, D., and Roselander, J.  Towards federated learning at scale: System design.  In _SysML 2019_ , 2019.
  * Cai et al. (2020) Cai, H., Gan, C., Zhu, L., and Han, S.  Tinytl: Reduce memory, not parameters for efficient on-device learning.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H. (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 11285–11297. Curran Associates, Inc., 2020.
  * Caldas et al. (2018) Caldas, S., Wu, P., Li, T., Konecný, J., McMahan, H. B., Smith, V., and Talwalkar, A.  LEAF: A benchmark for federated settings.  _CoRR_ , abs/1812.01097, 2018.  URL http://arxiv.org/abs/1812.01097.
  * Chahal et al. (2018) Chahal, K. S., Grover, M. S., and Dey, K.  A hitchhiker’s guide on distributed training of deep neural networks.  _CoRR_ , abs/1810.11787, 2018.  URL http://arxiv.org/abs/1810.11787.
  * Chowdhery et al. (2019) Chowdhery, A., Warden, P., Shlens, J., Howard, A., and Rhodes, R.  Visual wake words dataset.  _CoRR_ , abs/1906.05721, 2019.  URL http://arxiv.org/abs/1906.05721.
  * David et al. (2020) David, R., Duke, J., Jain, A., Reddi, V. J., Jeffries, N., Li, J., Kreeger, N., Nappier, I., Natraj, M., Regev, S., et al.  Tensorflow lite micro: Embedded machine learning on tinyml systems.  _arXiv preprint arXiv:2010.08678_ , 2020.
  * Fromm et al. (2018) Fromm, J., Patel, S., and Philipose, M.  Heterogeneous bitwidth binarization in convolutional neural networks.  In _Proceedings of the 32nd International Conference on Neural Information Processing Systems_ , NIPS’18, pp. 4010–4019, Red Hook, NY, USA, 2018. Curran Associates Inc.
  * Google (2020) Google.  Tensorflow federated: Machine learning on decentralized data.  https://www.tensorflow.org/federated, 2020.  accessed 25-Mar-20.
  * He et al. (2020) He, C., Li, S., So, J., Zhang, M., Wang, H., Wang, X., Vepakomma, P., Singh, A., Qiu, H., Shen, L., Zhao, P., Kang, Y., Liu, Y., Raskar, R., Yang, Q., Annavaram, M., and Avestimehr, S.  Fedml: A research library and benchmark for federated machine learning.  _arXiv preprint arXiv:2007.13518_ , 2020.
  * Jia et al. (2018) Jia, X., Song, S., He, W., Wang, Y., Rong, H., Zhou, F., Xie, L., Guo, Z., Yang, Y., Yu, L., Chen, T., Hu, G., Shi, S., and Chu, X.  Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes.  _CoRR_ , abs/1807.11205, 2018.  URL http://arxiv.org/abs/1807.11205.
  * Lee et al. (2019) Lee, T., Lin, Z., Pushp, S., Li, C., Liu, Y., Lee, Y., Xu, F., Xu, C., Zhang, L., and Song, J.  Occlumency: Privacy-preserving remote deep-learning inference using sgx.  In _The 25th Annual International Conference on Mobile Computing and Networking_ , MobiCom ’19, New York, NY, USA, 2019. Association for Computing Machinery.  ISBN 9781450361699.  doi: 10.1145/3300061.3345447.  URL https://doi.org/10.1145/3300061.3345447.
  * Leontiadis et al. (2021) Leontiadis, I., Laskaridis, S., Venieris, S. I., and Lane, N. D.  It’s always personal: Using early exits for efficient on-device cnn personalisation.  In _Proceedings of the 22nd International Workshop on Mobile Computing Systems and Applications_ , HotMobile ’21, pp. 15–21, New York, NY, USA, 2021. Association for Computing Machinery.  ISBN 9781450383233.  doi: 10.1145/3446382.3448359.  URL https://doi.org/10.1145/3446382.3448359.
  * Li et al. (2018) Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V.  Federated optimization in heterogeneous networks.  _arXiv preprint arXiv:1812.06127_ , 2018.
  * Malekzadeh et al. (2019) Malekzadeh, M., Athanasakis, D., Haddadi, H., and Livshits, B.  Privacy-preserving bandits, 2019.
  * McMahan et al. (2017) McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A.  Communication-efficient learning of deep networks from decentralized data.  In Singh, A. and Zhu, X. J. (eds.), _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FL, USA_ , volume 54 of _Proceedings of Machine Learning Research_ , pp. 1273–1282. PMLR, 2017\.  URL http://proceedings.mlr.press/v54/mcmahan17a.html.
  * Office31 (2020) Office31.  Office 31 dataset.  https://people.eecs.berkeley.edu/~jhoffman/domainadapt/, 2020.  accessed 10-Oct-20.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  Pytorch: An imperative style, high-performance deep learning library.  In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32_ , pp. 8026–8037. Curran Associates, Inc., 2019.
  * Ryffel et al. (2018) Ryffel, T., Trask, A., Dahl, M., Wagner, B., Mancuso, J., Rueckert, D., and Passerat-Palmbach, J.  A generic framework for privacy preserving deep learning.  _CoRR_ , abs/1811.04017, 2018.  URL http://arxiv.org/abs/1811.04017.
  * Sun et al. (2020) Sun, X., Wang, N., Chen, C.-Y., Ni, J., Agrawal, A., Cui, X., Venkataramani, S., El Maghraoui, K., Srinivasan, V. V., and Gopalakrishnan, K.  Ultra-low precision 4-bit training of deep neural networks.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H. (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 1796–1807. Curran Associates, Inc., 2020.
  * Warden & Situnayake (2019) Warden, P. and Situnayake, D.  _Tinyml: Machine learning with tensorflow lite on arduino and ultra-low-power microcontrollers_.  ” O’Reilly Media, Inc.”, 2019.