miscarticle

11institutetext: School of Biomedical Engineering & Imaging Sciences, King’s
College London, UK 22institutetext: Department of Development & Regeneration,
KU Leuven, Belgium 33institutetext: Radiology Department, Oxford University
Hospitals NHS Foundation Trust, UK 44institutetext: Department of
Cardiovascular Sciences, KU Leuven, Belgium

# Adaptive Multi-scale Online Likelihood Network for AI-assisted Interactive
Segmentation

Muhammad Asad 11 Helena Williams 22 Indrajeet Mandal 33 Sarim Ather 33 Jan
Deprest 22 Jan D’hooge 44 Tom Vercauteren 11

###### Abstract

Existing interactive segmentation methods leverage automatic segmentation and
user interactions for label refinement, significantly reducing the annotation
workload compared to manual annotation. However, these methods lack quick
adaptability to ambiguous and noisy data, which is a challenge in CT volumes
containing lung lesions from COVID-19 patients. In this work, we propose an
adaptive multi-scale online likelihood network (MONet) that adaptively learns
in a data-efficient online setting from both an initial automatic segmentation
and user interactions providing corrections. We achieve adaptive learning by
proposing an adaptive loss that extends the influence of user-provided
interaction to neighboring regions with similar features. In addition, we
propose a data-efficient probability-guided pruning method that discards
uncertain and redundant labels in the initial segmentation to enable efficient
online training and inference. Our proposed method was evaluated by an expert
in a blinded comparative study on COVID-19 lung lesion annotation task in CT.
Our approach achieved 5.86% higher Dice score with 24.67% less perceived NASA-
TLX workload score than the state-of-the-art. Source code is available at:
https://github.com/masadcv/MONet-MONAILabel

##  1 Introduction

Deep learning methods for automatic lung lesion segmentation from CT volumes
have the potential to alleviate the burden on clinicians in assessing lung
damage and disease progression in COVID-19 patients [20, 21, 22]. However,
these methods require large amounts of manually labeled data to achieve the
level of robustness required for their clinical application [5, 8, 27, 23].
Manual labeling of CT volumes is time-consuming and may increase the workload
of clinicians. Additionally, applying deep learning-based segmentation models
to data from new unseen sources can result in suboptimal lesion segmentation
due to unseen acquisition devices/parameters, variations in patient pathology,
or future coronavirus variants resulting in new appearance characteristics or
new lesion pathologies [16]. To address this challenge, interactive
segmentation methods that can quickly adapt to such changing settings are
needed. These can be used either by end-users or algorithm developers to
quickly expand existing annotated datasets and enable agile retraining of
automatic segmentation models [4].

####  1.0.1 Related work:

Interactive segmentation methods for Artificial Intelligence (AI) assisted
annotation have shown promising applications in the existing literature [14,
18, 25, 26]. BIFSeg [26] utilizes a bounding box and scribbles with
convolutional neural network (CNN) image-specific fine-tuning to segment
potentially unseen objects of interest. MIDeepSeg [14] incorporates user-
clicks with the input image using exponential geodesic distance. However,
BIFSeg, MIDeepSeg, and similar deep learning-based methods exploit large
networks that do not adapt rapidly to new data examples in an online setting
due to the elevated computational requirements.

Due to their quick adaptability and efficiency, a number of existing online
likelihood methods have been applied as interactive segmentation methods [2,
3, 24]. DybaORF [24] utilizes hand-crafted features with dynamically changing
weights based on interactive labels’ distribution to train a Random Forest
classifier. ECONet [2] improves online learning with a shallow CNN that
jointly learns both features and classifier to outperform previous online
likelihood inference methods. While ECONet is, to the best of our knowledge,
the only online learning method that addresses COVID-19 lung lesion
segmentation, it is limited to learning from user scribbles only. This means
that it requires a significant amount of user interaction to achieve expert-
level accuracy. Additionally, the model uses a single convolution for feature
extraction, limiting its accuracy to a specific scale of pathologies. For each
CT volume, the model is trained from scratch, resulting in lack of prior
knowledge about lesions.

Figure 1: Adaptive online training weights: (a) input, (b) foreground /
background scribbles, (c) foreground and (d) background weights using
$\tau$=$0.2$ in Eq. (2).

####  1.0.2 Contributions:

To overcome limitations of existing techniques, we propose adaptive multi-
scale online likelihood network (MONet) for AI-assisted interactive
segmentation of lung lesions in CT volumes from COVID-19 patients. Our
contributions are three-fold, we propose:

  1. i.

Multi-scale online likelihood network (MONet), consisting of a multi-scale
feature extractor, which enables relevant features extraction at different
scales for improved accuracy;

  2. ii.

Adaptive online loss that uses weights from a scaled negative exponential
geodesic distance from user-scribbles, enabling adaptive learning from both
initial segmentation and user-provided corrections (Fig. 1);

  3. iii.

Probability-guided pruning approach where uncertainty from initial
segmentation model is used for pruning ambiguous online training data.

MONet enables human-in-the-loop online learning to perform AI-assisted
annotations and should not be mistaken for an end-to-end segmentation model.

We perform expert evaluation which shows that adaptively learned MONet
outperforms existing state-of-the-art, achieving 5.86% higher Dice score with
24.67% less perceived NASA-TLX workload score evaluated.

##  2 Method

Given an input image volume, $I$, a pre-trained CNN segmentation model
generates an automatic segmentation $C$ with associated probabilities $P$.
When using data from a new domain, the automated network may fail to properly
segment foreground/background objects. To improve this, the user provides
scribbles-based interaction indicating corrected class labels for a subset of
voxels in the image $I$. Let $\mathcal{S}=\mathcal{S}^{f}\cup\mathcal{S}^{b}$
represent these set of scribbles, where $\mathcal{S}^{f}$ and
$\mathcal{S}^{b}$ denote the foreground and background scribbles,
respectively, and $\mathcal{S}^{f}\cap\mathcal{S}^{b}=\emptyset$. Fig. 2 (a)
shows scribbles $\mathcal{S}$, along with the initial segmentation $C$ and
probabilities $P$.

Figure 2: Adaptive learning for interactive segmentation: (a) training and
inference of MONet using adaptive loss and probability-guided pruning; (b)
architecture of our multi-scale online likelihood network (MONet).

###  2.1 Multi-scale Online Likelihood Network

Our proposed multi-scale online likelihood network (MONet), shown in Fig. 2
(b), uses a multi-scale feature extractor that applies a 3D convolution at
various kernel sizes to capture spatial information at different scales. The
output of each scale is concatenated and fed to a fully-connected classifier,
which infers the likelihood for background/foreground classification of the
central voxel in the input patch. Each layer in MONet is followed by batch
normalization and ReLU activation.

###  2.2 Adaptive Loss for Online Learning

The scribbles $\mathcal{S}$ only provide sparse information for online
learning. However, these corrections are likely also applicable to neighboring
voxels with similar appearance features, thereby providing an extended source
of training information. Concurrently, the initial automated segmentation $C$
will often provide reliable results away from the scribbles. To extend the
influence of the scribbles $\mathcal{S}$ while preserving the quality of the
initial segmentation $C$, we propose a spatially-varying adaptive online loss:

$\mathcal{L}=-\sum_{i}\left[(1-W_{i})\mathcal{L}^{C}_{i}+W_{i}\mathcal{L}^{\mathcal{S}}_{i}\right],$
(1)

where $i$ is a voxel index, $\mathcal{L}^{C}$ and $\mathcal{L}^{\mathcal{S}}$
are individual loss terms for learning from the automated segmentation $C$ and
the user-provided correction scribbles $\mathcal{S}$ respectively. $W$ are
spatially-varying interaction-based weights defined using the geodesic
distance $D$ between voxel $i$ and the scribbles $\mathcal{S}$:

$\displaystyle W_{i}$
$\displaystyle=\exp\left({\frac{-D(i,S,I)}{\tau}}\right),$ (2)

where the temperature term $\tau$ controls the influence of $W$ in $I$. The
geodesic distance to the scribbles is defined as
$D(i,\mathcal{S},I)={\mathrm{min}}_{j\in\mathcal{S}}\,d(i,j,I)$ where
$d(i,j,I)={\mathrm{min}}_{p\in\mathcal{P}_{i,j}}\int_{0}^{1}\left\|\nabla
I(p(x))\cdot\mathbf{u}(x)\right\|\,\,dx$ and $\mathcal{P}_{i,j}$ is the set of
all possible differentiable paths in $I$ between voxels $i$ and $j$. A
feasible path $p$ is parameterized by $x\in[0,1]$. We denote
$\mathbf{u}(x)=p^{\prime}(x)/\left\|p^{\prime}(x)\right\|$ the unit vector
tangent to the direction of the path $p$. We further let $D=\infty$ for
$\mathcal{S}=\emptyset$.

####  2.2.1 Dynamic Label-balanced Cross-Entropy Loss:

User-scribbles for online interactive segmentation suffer from dynamically
changing class imbalance [2]. Moreover, lung lesions in CT volumes usually
occupy a small subset of all voxels, introducing additional label imbalance
and hence reducing their impact on imbalanced online training. To address
these challenges, we utilize a label-balanced cross-entropy loss [2, 10, 11],
with dynamically changing class weights from segmentations and scribbles
distribution. Given an online model with parameters $\theta$, the foreground
likelihood from this model is $p_{i}=P(s_{i}=1|I,\theta)$. Then, the
segmentations-balanced and scribbles-balanced cross-entropy terms are:

$\mathcal{L}^{C}_{i}=\alpha^{f}y^{C}_{i}\log
p_{i}+\alpha^{b}(1-y^{C}_{i})\log(1-p_{i}),$ (3)
$\mathcal{L}^{S}_{i}=\beta^{f}y^{\mathcal{S}}_{i}\log
p_{i}+\beta^{b}(1-y^{\mathcal{S}}_{i})\log(1-p_{i}),$ (4)

where $\alpha$ and $\beta$ are class weights for labels $C$ and scribbles
$\mathcal{S}$ that are defined by labels and scribbles distributions during
online interaction as:
$\alpha^{f}={\left|\mathcal{T}\right|}/{\left|C^{f}\right|}$,
$\alpha^{b}={\left|\mathcal{T}\right|}/{\left|C^{b}\right|}$,
$\beta^{f}={\left|\mathcal{T}\right|}/{\left|\mathcal{S}^{f}\right|}$,
$\beta^{b}={\left|\mathcal{T}\right|}/{\left|\mathcal{S}^{b}\right|}$ and
$\left|\mathcal{T}\right|=\left|C\right|+\left|\mathcal{S}\right|$.
$y^{C}_{i}$ and $y^{\mathcal{S}}_{i}$ represent labels in $C$ and
$\mathcal{S}$, respectively.

The patch-based training approach from [2] is used to first extract
K$\times$K$\times$K patches from $I$ centered around each voxel in
$\mathcal{S}$ and $C$ and train MONet using Eq. (1). Once learned, efficient
online inference from MONet is achieved by applying it to the whole input CT
volumes as a fully convolutional network [12].

###  2.3 Improving Efficiency with Probability-guided Pruning

MONet is applied as an online likelihood learning method, where the online
training happens with an expert human annotator in the loop, which makes
online training efficiency critical. We observe that the automatic
segmentation models provide dense labels $C$ which may significantly impact
online training and inference performance. $C$ may contain ambiguous
predictions for new data, and a number of voxels in $C$ may provide redundant
labels. To improve online efficiency while preserving accuracy during
training, we prune labels as $C^{*}=\mathcal{M}\odot C$ where:
$\mathcal{M}_{i}$ is set to $1$ if $P_{i}\geq\zeta$ and $U_{i}\geq\eta$ and
$0$ otherwise. $\zeta\in[0,1]$ is the minimum confidence to preserve a label,
$U_{i}\in[0,1]$ is a uniformly distributed random variable, and $\eta\in[0,1]$
is the fraction of samples to prune.

Table 1: State-of-the-art evaluated comparison methods, showing improvement in
accuracy (Dice and ASSD) when using different features. Features in blue text
are proposed in this paper. Key: OL - online learning, PP - post-processing.

Method |  Technique |  Initial |  Multi |  Adaptive |  Temp. |  Dice |  ASSD
---|---|---|---|---|---|---|---
|  |  Seg. |  Scale |  Loss |  ($\tau$) |  (%) |
MONet (proposed) |  OL |  ✓ |  ✓ |  ✓ |  ✓ |  77.77 |  11.82
MONet-NoMS |  OL |  ✓ |  ✗ |  ✓ |  ✓ |  77.06 |  13.01
ECONet[2] |  OL |  ✗ |  ✗ |  ✗ |  ✗ |  77.02 |  20.19
MIDeepSegTuned[14] |  PP |  ✓ |  ✗ |  ✗ |  ✓ |  76.00 |  20.16
MIDeepSeg[14] |  PP |  ✓ |  ✗ |  ✗ |  ✗ |  56.85 |  33.25
IntGraphCut |  PP |  ✓ |  ✗ |  ✗ |  ✗ |  68.58 |  28.64

##  3 Experimental Validation

Table 1 outlines the different state-of-the-art interactive segmentation
methods and their extended variants that we introduce for fair comparison. We
compare our proposed MONet with ECONet [2] and MIDeepSeg [14]. As our proposed
Eq. (2) is inspired by the exponential geodesic distance from MIDeepSeg [14],
we introduce MIDeepSegTuned, which utilizes our proposed addition of a
temperature term $\tau$. Moreover, to show the importance of multi-scale
features, we include MONet-NoMS which uses features from a single 3D
convolution layer. We utilize MONAI Label to implement all online likelihood
methods [7]. For methods requiring an initial segmentation, we train a 3D UNet
[6] using MONAI [17] with features $[32,32,64,128,256,32]$. Output from each
method is regularized using GraphCut optimization. We also compare against a
baseline interactive GraphCut (IntGraphCut) implementation, that updates UNet
output with scribbles based on [3] and then performs GraphCut optimization.
The proposed method is targeted for online training and inference, where quick
adaptability with minimal latency is required. Note that incorporating more
advanced deep learning methods in this context would result in a considerable
decrease in online efficiency, rendering the method impractical for online
applications [2]. We utilize a GPU-based implementation of geodesic distance
transform [1] in Eq. (2), whereas MIDeepSeg uses a CPU-based implementation.
We use NVIDIA Tesla V100 GPU with 32 GB memory for all our experiments.
Comparison of accuracy for each method is made using Dice similarity (Dice)
and average symmetric surface distance (ASSD) metrics against ground truth
annotations [2, 14]. Moreover, we compare performance using execution time
(Time), including online training and inference time, average full annotation
time (FA-Time), and number of voxels with scribbles (S) needed for a given
accuracy.

####  3.0.1 Data:

To simulate a scenario where the automatic segmentation model is trained on
data from a different source than it is tested on, we utilize two different
COVID-19 CT datasets. The dataset from the COVID-19 CT lesion segmentation
challenge [21] is used for training and validation of 3D UNet for automatic
segmentation task and patch-based pre-training of MONet/MONet-NoMS/ECONet.
This dataset contains binary lung lesions segmentation labels for 199 CT
volumes (160 training, 39 validation). We use UESTC-COVID-19 [27], a dataset
from a different source, for the experimental evaluation of interactive
segmentation methods (test set). This dataset contains 120 CT volumes with
lesion labels, from which 50 are by expert annotators and 70 are by non-expert
annotators. To compare robustness of the proposed method against expert
annotators, we only use the 50 expert labelled CT volumes.

Method |  Dice (%) |  ASSD |  Time (s) |  Scribbles
---|---|---|---|---
MONet (proposed) |  77.77$\pm$6.84 |  11.82$\pm$12.83 |  6.18$\pm$2.42 |  20$\pm$24
MONet-NoMS |  77.06$\pm$7.27 |  13.01$\pm$15.29 |  7.76$\pm$8.16 |  20$\pm$24
ECONet[2] |  77.02$\pm$6.94 |  20.19$\pm$14.71 |  1.46$\pm$1.22 |  2283$\pm$2709
MIDeepSegTuned[14] |  76.00$\pm$7.37 |  20.16$\pm$22.57 |  7.97$\pm$2.47 |  23$\pm$17
MIDeepSeg[14] |  56.85$\pm$14.25 |  33.25$\pm$25.26 |  6.26$\pm$1.46 |  436$\pm$332
IntGraphCut |  68.58$\pm$9.09 |  28.64$\pm$27.36 |  0.11$\pm$0.04 |  480$\pm$359

Table 2: Quantitative comparison of interactive segmentation methods using
synthetic scribbler shows mean and standard deviation of Dice, ASSD, Time and
Synthetic Scribbles Voxels.

Figure 3: Validation accuracy using synthetic scribbles.

####  3.0.2 Training Parameters:

Training of 3D UNet utilized a learning rate (lr) of $1e^{-4}$ for 1000 epochs
and MONet/MONet-NoMS offline pre-training used 50 epochs, and
$\text{lr}=1e^{-3}$ dropped by 0.1 at the 35th and 45th epoch. Online training
for MONet, MONet-NoMS and ECONet [2] used 200 epochs with $\text{lr}=1e^{-2}$
set using cosine annealing scheduler [13]. Dropout of 0.3 was used for all
fully-connected layers in online models. Each layer size in ECONet and MONet-
NoMS was selected by repeating line search experiments from [2]: (i) input
patch/3D convolution kernel size of $K=9$, (ii) 128 input 3D convolution
filters and (iii) fully-connected sizes of 32$\times$16$\times$2\. For MONet,
we utilize four input 3D convolution with multi-scale kernel sizes
$K=[1,3,5,9]$ with each containing 32 filters (i.e., a total of 128 filters,
same as (ii)). We utilize the same fully-connected sizes as in (iii) above.
Parameters $\zeta=0.8$ and $\eta=0.98$ are selected empirically. We utilize
$\tau=0.3$ for MONet, MIDeepSegTuned and MONet-NoMS. We use GraphCut
regularization, where $\lambda=2.5$ and $\sigma=0.15$ [3]. Search experiments
used for selecting $\tau$, $\lambda$, $\sigma$ are shown in Fig. 1 and 2 in
supplementary material.

###  3.1 Quantitative Comparison using Synthetic Scribbler

We employ the synthetic scribbler method from [2, 25] where mis-segmented
regions in the inferred segmentations are identified by comparison to the
ground truth segmentations. Table 2 and Fig. 3 present quantitative comparison
of methods using synthetic scribbler. They show that MONet outperforms all
existing state-of-the-art in terms of accuracy with the least number of
synthetic scribbled voxels. In particular, MONet outperforms both MIDeepSeg
[14] and MIDeepSegTuned, where adaptive online learning enables it to quickly
adapt and refine segmentations. In terms of efficiency, online training and
inference of the proposed MONet takes around 6.18 seconds combined, which is
22.4% faster as compared to 7.97 seconds for MIDeepSeg. However, it is slower
than ECONet and ISeg. MIDeepSeg performs the worst as it is unable to adapt to
large variations and ambiguity within lung lesions from COVID-19 patients,
whereas by utilizing our proposed Eq. (2) in MIDeepSegTuned, we improve its
accuracy. When comparing to online learning methods, MONet outperforms MONet-
NoMS, where the accuracy is improved due to MONet’s ability to extract multi-
scale features. Existing state-of-the-art online method ECONet [2] requires
significantly more scribbled voxels as it only relies on user-scribbles for
online learning.

|  MONet (proposed) |  MIDeepSeg-Tuned[14]
---|---|---
Finished |  100% |  33.33 %
NASA-TLX |  52.33 |  77.00
Dice (%) |  88.53$\pm$2.27 |  82.67$\pm$12.36
ASSD |  2.91$\pm$1.58 |  11.30$\pm$20.09
FA-Time (s) |  507.11 |  567.43

Table 3: Workload validation by expert user, shows Dice (%), ASSD, full
annotation time, FA-Time (s), overall NASA-TLX perceived workload score and
the % of data successfully annotated by expert.

NASA-TLX weighted scores |  MONet (proposed) |  MIDeepSeg-Tuned[14]
---|---|---
Effort |  14.67 |  21.33
Frustration |  7.67 |  16.67
Mental Demand |  13.33 |  15.00
Performance |  10.00 |  17.00
Physical Demand |  4.67 |  7.00
Temporal Demand |  2.00 |  0.00
Total workload |  52.33 |  77.00

Table 4: NASA-TLX perceived workload by expert user, shows total workload and
individual sub-scale scores. The method with low score requires less effort,
frustration, mental, temporal and physical demands with high perceived
performance.

###  3.2 Performance and Workload Validation by Expert User

This experiment aims to compare the performance and _perceived_ subjective
workload of the proposed MONet with the best performing comparison method
MIDeepSegTuned based on [14]. We asked an expert, with 2 years of experience
in lung lesion CT from Radiology Department, Oxford University Hospitals NHS
Foundation Trust, to utilize each method for labelling the following
pathologies as lung lesions in 10 CT volumes from UESTC-COVID-19 expert set
[27]: ground glass opacity, consolidation, crazy-paving, linear opacities. One
CT volume is used by the expert to practice usage of our tool. The remaining 9
CT volumes were presented in a random order, where the perceived workload was
evaluated by the expert at half way (after 5 segmentations) and at the end. We
use the National Aeronautics and Space Administration Task Load Index (NASA-
TLX) [9] as per previous interactive segmentation studies [15, 19, 28]. The
NASA-TLX asks the expert to rate the task based on six factors, being
performance, frustration, effort, mental, physical and temporal demand. The
weighted NASA-TLX score is then recorded as the expert answers 15 pair-wise
questions rating factors based on importance. In addition, we also recorded
accuracy metrics (Dice and ASSD) against ground truth labels in [27], time
taken to complete annotation and whether the expert was able to successfully
complete their task within 10 minutes allocated for each volume.

Figure 4: Visual comparison of interactive segmentation results from Section
3.2. Segmentations are shown with contours on axial plane slices from
different cases.

Table 3 presents an overview for this experiment, where using the proposed
MONet, the expert was able to complete 100% of the labelling task, whereas
using MIDeepSegTuned they only completed 33.33% within the allocated time. In
addition, MONet achieves better accuracy with lower time for complete
annotation and less overall perceived workload with NASA-TLX of 52.33% as
compared to 77.00% for MIDeepSegTuned. Table 4 shows the individual scores
that contribute to overall perceived workload. It shows that using the
proposed MONet, the expert perceived reduced workload in all sub-scale scores
except temporal demand. We believe this is due to the additional online
training/inference overhead for MONet application. Fig. 4 visually compares
these results where MONet results in more accurate segmentation as compared to
MIDeepSegTuned. We also note that MONet’s ability to apply learned knowledge
on the whole volume enables it to also infer small isolated lesions, which
MIDeepSegTuned fails to identify.

##  4 Conclusion

We proposed a multi-scale online likelihood network (MONet) for scribbles-
based AI-assisted interactive segmentation of lung lesions in CT volumes from
COVID-19 patients. MONet consisted of a multi-scale feature extractor that
enabled extraction of relevant features at different scales for improved
accuracy. We proposed an adaptive online loss that utilized adaptive weights
based on user-provided scribbles that enabled adaptive learning from both an
initial automated segmentation and user-provided label corrections.
Additionally, we proposed a dynamic label-balanced cross-entropy loss that
addressed dynamic class imbalance, an inherent challenge for online
interactive segmentation methods. Experimental validation showed that the
proposed MONet outperformed the existing state-of-the-art on the task of
annotating lung lesions in COVID-19 patients. Validation by an expert showed
that the proposed MONet achieved on average 5.86% higher Dice while achieving
24.67% less perceived NASA-TLX workload score than the MIDeepSegTuned method
[14].

##  5 Acknowledgments

This project has received funding from the European Union’s Horizon 2020
research and innovation programme under grant agreement No 101016131 (icovid
project). This work was also supported by core and project funding from the
Wellcome/EPSRC [WT203148/Z/16/Z; NS/A000049/1; WT101957; NS/A000027/1]. This
project utilized scribbles-based interactive segmentation tools from
opensource project MONAI Label 111https://github.com/Project-MONAI/MONAILabel
[7].

## References

  * [1] Muhammad Asad, Reuben Dorent and Tom Vercauteren  “FastGeodis: Fast Generalised Geodesic Distance Transform”  In _arXiv preprint arXiv:2208.00001_ , 2022
  * [2] Muhammad Asad, Lucas Fidon and Tom Vercauteren  “ECONet: Efficient Convolutional Online Likelihood Network for Scribble-based Interactive Segmentation”  In _Medical Imaging with Deep Learning_ , 2022
  * [3] Yuri Y Boykov and M-P Jolly  “Interactive graph cuts for optimal boundary & region segmentation of objects in ND images”  In _Proceedings eighth IEEE international conference on computer vision. ICCV 2001_ 1, 2001, pp. 105–112  IEEE
  * [4] Samuel Budd, Emma C Robinson and Bernhard Kainz  “A survey on active learning and human-in-the-loop deep learning for medical image analysis”  In _Medical Image Analysis_ 71 Elsevier, 2021, pp. 102062
  * [5] Guillaume Chassagnon et al.  “AI-Driven CT-based quantification, staging and short-term outcome prediction of COVID-19 pneumonia”  In _arXiv preprint arXiv:2004.12852_ , 2020
  * [6] Özgün Çiçek et al.  “3D U-Net: learning dense volumetric segmentation from sparse annotation”  In _International conference on medical image computing and computer-assisted intervention_ , 2016, pp. 424–432  Springer
  * [7] Andres Diaz-Pinto et al.  “Monai label: A framework for ai-assisted interactive labeling of 3d medical images”  In _arXiv preprint arXiv:2203.12362_ , 2022
  * [8] Camila Gonzalez et al.  “Detecting when pre-trained nnU-Net models fail silently for Covid-19 lung lesion segmentation”  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , 2021, pp. 304–314  Springer
  * [9] Sandra G Hart  “NASA-task load index (NASA-TLX); 20 years later”  In _Proceedings of the human factors and ergonomics society annual meeting_ 50.9, 2006, pp. 904–908  Sage publications Sage CA: Los Angeles, CA
  * [10] Yaoshiang Ho and Samuel Wookey  “The real-world-weight cross-entropy loss function: Modeling the costs of mislabeling”  In _IEEE Access_ 8 IEEE, 2019, pp. 4806–4813
  * [11] Matjaz Kukar and Igor Kononenko  “Cost-sensitive learning with neural networks.”  In _ECAI_ 15, 1998, pp. 88–94  Citeseer
  * [12] Jonathan Long, Evan Shelhamer and Trevor Darrell  “Fully convolutional networks for semantic segmentation”  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2015, pp. 3431–3440
  * [13] Ilya Loshchilov and Frank Hutter  “Sgdr: Stochastic gradient descent with warm restarts”  In _arXiv preprint arXiv:1608.03983_ , 2016
  * [14] Xiangde Luo et al.  “MIDeepSeg: Minimally interactive segmentation of unseen objects from medical images using deep learning”  In _Medical Image Analysis_ 72 Elsevier, 2021, pp. 102102
  * [15] Hari McGrath et al.  “Manual segmentation versus semi-automated segmentation for quantifying vestibular schwannoma volume on MRI”  In _International Journal of Computer Assisted Radiology and Surgery_ 15 Springer, 2020, pp. 1445–1455
  * [16] Thomas A McLaren, James F Gruden and Daniel B Green  “The bullseye sign: A variant of the reverse halo sign in COVID-19 pneumonia”  In _Clinical Imaging_ 68 Elsevier, 2020, pp. 191–196
  * [17] MONAI Consortium  “MONAI: Medical Open Network for AI”, 2020  DOI: 10.5281/zenodo.4323058
  * [18] Martin Rajchl et al.  “Deepcut: Object segmentation from bounding box annotations using convolutional neural networks”  In _IEEE transactions on medical imaging_ 36.2 IEEE, 2016, pp. 674–683
  * [19] Anjana Ramkumar et al.  “Using GOMS and NASA-TLX to evaluate human–computer interaction process in interactive segmentation”  In _International Journal of Human–Computer Interaction_ 33.2 Taylor & Francis, 2017, pp. 123–134
  * [20] Marie-Pierre Revel et al.  “Study of Thoracic CT in COVID-19: The STOIC Project”  In _Radiology_ 301.1 Radiological Society of North America, 2021, pp. E361–E370
  * [21] Holger Roth et al.  “Rapid Artificial Intelligence Solutions in a Pandemic-The COVID-19-20 Lung CT Lesion Segmentation Challenge”, Research Square preprint, 2021
  * [22] Geoffrey D Rubin et al.  “The role of chest imaging in patient management during the COVID-19 pandemic: a multinational consensus statement from the Fleischner Society”  In _Radiology_ 296.1 Radiological Society of North America, 2020, pp. 172–180
  * [23] Sofie Tilborghs et al.  “Comparative study of deep learning methods for the automatic segmentation of lung, lesion and lesion type in CT scans of COVID-19 patients”  In _arXiv preprint arXiv:2007.15546_ , 2020
  * [24] Guotai Wang et al.  “Dynamically balanced online random forests for interactive scribble-based segmentation”  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , 2016, pp. 352–360  Springer
  * [25] Guotai Wang et al.  “DeepIGeoS: a deep interactive geodesic framework for medical image segmentation”  In _IEEE transactions on pattern analysis and machine intelligence_ 41.7 IEEE, 2018, pp. 1559–1572
  * [26] Guotai Wang et al.  “Interactive medical image segmentation using deep learning with image-specific fine tuning”  In _IEEE transactions on medical imaging_ 37.7 IEEE, 2018, pp. 1562–1573
  * [27] Guotai Wang et al.  “A noise-robust framework for automatic segmentation of COVID-19 pneumonia lesions from CT images”  In _IEEE Transactions on Medical Imaging_ 39.8 IEEE, 2020, pp. 2653–2663
  * [28] Helena Williams et al.  “Interactive segmentation via deep learning and B-spline explicit active surfaces”  In _Medical Image Computing and Computer Assisted Intervention–MICCAI 2021: 24th International Conference, Strasbourg, France, September 27–October 1, 2021, Proceedings, Part I 24_ , 2021, pp. 315–325  Springer


# Mobile V-MoEs: Scaling Down Vision Transformers
via Sparse Mixture-of-Experts

Erik Daxberger Correspondence to: Erik Daxberger<EMAIL_ADDRESS>Floris
Weers  Bowen Zhang  Tom Gunter  Ruoming Pang  Marcin Eichner  Michael
Emmersberger  Yinfei Yang  Alexander Toshev  Xianzhi Du

Apple

###### Abstract

Sparse Mixture-of-Experts models (MoEs) have recently gained popularity due to
their ability to decouple model size from inference efficiency by only
activating a small subset of the model parameters for any given input token.
As such, sparse MoEs have enabled unprecedented scalability, resulting in
tremendous successes across domains such as natural language processing and
computer vision. In this work, we instead explore the use of sparse MoEs to
scale-down Vision Transformers (ViTs) to make them more attractive for
resource-constrained vision applications. To this end, we propose a simplified
and mobile-friendly MoE design where entire images rather than individual
patches are routed to the experts. We also propose a stable MoE training
procedure that uses super-class information to guide the router. We
empirically show that our sparse Mobile Vision MoEs (V-MoEs) can achieve a
better trade-off between performance and efficiency than the corresponding
dense ViTs. For example, for the ViT-Tiny model, our Mobile V-MoE outperforms
its dense counterpart by 3.39% on ImageNet-1k. For an even smaller ViT variant
with only 54M FLOPs inference cost, our MoE achieves an improvement of 4.66%.

##  1 Introduction

Figure 1: Accuracy vs. FLOPs for ViTs of different sizes. Labels (e.g.
12$\times$192, which is ViT-Tiny) refer to the number of ViT layers (e.g. 12)
and the hidden embedding dimension (e.g. 192). The sparse MoEs outperform
their corresponding dense baselines across different model scales. Fig. 3(a)
lists all numerical results.

(a) Dense ViT

(b) Regular V-MoE

(c) Mobile V-MoE

(d) Layer types

Figure 2: Model architectures. (a) The dense ViT baseline model uses dense ViT
layers throughout. (b) Regular sparse V-MoE with layer-wise per-patch routers.
(c) Our proposed sparse Mobile V-MoE design with a single per-image router. In
both (b) and (c), dense ViT layers are followed by MoE-ViT layers (here, $k=1$
out of $E=3$ experts are activated per input). (d) In contrast to dense ViT
layers [19], MoE-ViT layers have a separate MLP per expert (preceded by a
router) while all other parts of the layer are shared across all experts [13].

The trade-off between performance and efficiency of neural networks (NNs)
remains a challenge, especially in settings where computational resources are
limited. Recently, sparsely-gated Mixture-of-Experts models (sparse MoEs) have
gained popularity as they provide a promising solution to this problem by
enabling the decoupling of model size from inference efficiency [3]. MoEs are
NNs that are partitioned into “experts”, which are trained jointly with a
router to specialize on subsets of the data. In MoEs, each input is processed
by only a small subset of model parameters (aka _conditional computation_). In
contrast, traditional dense models activate all parameters for each input.

Sparse MoEs were popularized in deep learning by [16], which introduced sparse
MoE-layers as drop-in replacements for standard NN layers. Most recent MoEs
are based on the Transformer [19], which processes individual input tokens; in
accordance, recent MoEs also route individual input tokens to experts, i.e.,
image patches in the case of Vision Transformers (ViTs) [13, 2] (see Fig. 2).
Conditional computation as implemented by sparse MoEs has enabled the training
of Transformers of unprecedented size [4]. As a result, MoEs have achieved
impressive successes across various domains including language [10, 4], vision
[13], speech [20] and multi-modal learning [12], and currently hold state-of-
the-art results on many benchmarks [21].

The ability to increase model capacity while keeping inference cost low is
also appealing for resource-constrained vision problems. While Transformers
are getting increasingly established as the de-facto standard architecture for
large-scale visual modeling [2, 13], virtually all mobile-friendly models
still leverage convolutions due to their efficiency [6, 15, 5, 1, 11, 18]. As
such, conditional computation could potentially enable attention-based models
to reduce the gap to convolutional models in the small-scale regime. However,
Transformer-based MoEs have not yet been explored for resource-constrained
settings; this might be due to two main weaknesses of recently-popularized
MoEs [16].

Firstly, while per-token routing increases the flexibility to learn an optimal
computation path through the model, it makes inference inefficient, as many
(or even all) experts need to be loaded for a single input image. Secondly,
recent MoEs train the routers jointly with the rest or the model in an end-to-
end fashion. To avoid collapse to just a few experts while ignoring all
others, one needs to use load balancing mechanisms [3] such as dedicated
auxiliary losses [16]. However, the resulting complex optimization objectives
often lead to training instabilities / divergence [10, 4, 21, 12].

In this work, we investigate the potential of sparse MoEs to scale-down ViTs
for resource-constrained vision applications via an MoE design and training
procedure that addresses the aforementioned issues. Our contributions are:

  1. 1.

We propose a simplified, mobile-friendly sparse MoE design in which a single
router assigns entire images (rather than image patches) to the experts (see
Fig. 2).

  2. 2.

We develop a simple yet robust training procedure in which expert imbalance is
avoided by leveraging semantic super-classes to guide the router training.

  3. 3.

We empirically show that our proposed sparse MoE approach allows us to scale-
down ViT models by improving their performance vs. efficiency trade-off.

##  2 Scaling down ViTs via sparse MoEs

###  2.1 Conditional computation with sparse MoEs

An MoE implements conditional computation by activating different subsets of a
NN (so-called experts) for different inputs. We consider an MoE layer with $E$
experts as

$\text{MoE}(\mathbf{x})=\sum_{i=1}^{E}g(\mathbf{x})_{i}e_{i}(\mathbf{x}),$ (1)

where $\mathbf{x}\in\mathbb{R}^{D}$ is the input to the layer,
$e_{i}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ is the function computed by
expert $i$, and $g:\mathbb{R}^{D}\rightarrow\mathbb{R}^{E}$ is the routing
function which computes an input-dependent weight for each expert [16]. In a
ViT-based MoE, each expert $e_{i}$ is parameterized by a separate multi-layer
perceptron (MLP) within the ViT layer, while the other parts are shared across
experts (see Fig. 2). We use the routing function

$g(\mathbf{x})=\text{TOP}_{k}(\text{softmax}(\mathbf{W}\mathbf{x})),$ (2)

where the operation $\text{TOP}_{k}(\mathbf{x})$ sets all elements of
$\mathbf{x}$ to zero except those with the $k$ largest values [13]. In a
sparse MoE, we have $k\ll E$, s.t. we only need to load and compute the $k$
experts with the largest routing weights. This allows us to scale-up the
overall model capacity (determined by $E$) without increasing the inference
cost (determined by $k$).

###  2.2 Efficient and robust MoEs for small-scale ViTs

Per-image routing. Recent large-scale sparse MoEs use per-patch routing (i.e.
the inputs $\mathbf{x}$ are individual image patches). This generally requires
a larger number of experts to be activated for each image. For example, [13]
show that in their MoE with per-patch routing, “most images use –on aggregate
by pooling over all their patches– most of the experts” [13, Appendix E.3].
Thus, per-patch routing can increase the computational and memory overhead of
the routing mechanism and reduce the overall model efficiency. We instead
propose to use per-image routing (i.e., the inputs $\mathbf{x}$ are entire
images) to reduce the number of activated experts per image, as also done in
early works on MoEs [7, 9].

Super-class-based routing. Previous works on sparse MoEs jointly train the
router end-to-end together with the experts and the dense ViT backbone, to
allow the model to learn the optimal assignment from inputs to experts based
on the data [13]. While learning the optimal routing mechanism from scratch
can result in improved performance, it often leads to training instabilities
and expert collapse, where most inputs are routed to only a small subset of
the experts, while all other experts get neglected during training [3]. Thus,
an additional auxiliary loss is typically required to ensure load-balancing
between the different experts, which can increase the complexity of the
training process [3].

In contrast, we propose to group the classes of the dataset into super-classes
and explictly train the router to make each expert specialize on one super-
class. To this end, we add an additional cross-entropy loss $\mathcal{L}_{g}$
between the router output $g(\mathbf{x})$ in Eq. 2 and the ground truth super-
class labels to the regular classification loss $\mathcal{L}_{C}$ to obtain
the overall weighted loss $\mathcal{L}=\mathcal{L}_{C}+\lambda\mathcal{L}_{g}$
(we use $\lambda=0.3$ in our experiments, which we found to work well). Such a
super-class division is often readily provided with the dataset (e.g. for
CIFAR-10/100 or MS-COCO). If a dataset does not come with a super-class
division, we can easily obtain one as follows: 1) we first train a dense
baseline model on the dataset; 2) we then compute the model’s confusion matrix
over a held-out validation set; 3) we finally construct a confusion graph from
the confusion matrix and apply a graph clustering algorithm to obtain the
super-class division [8]. This approach encourages the super-classes to
contain semantically similar images that the model often confuses.
Intuitively, by allowing the different MoE experts to specialize on the
different semantic data clusters, performance on the highly-confused classes
should be improved. We use this approach in our experiments on ImageNet-1k,
computing the confusion matrix via a dense ViT-S/16 model. The resulting
super-class division for $E=10$ experts is shown in Tab. 1; the super-classes
contain semantically related classes.

ID | Classes | Super-class
---|---|---
0 | boxer, pug, Rottweiler | dogs
1 | orangutan, weasel, panda | other mammals
2 | toucan, flamingo, ostrich | birds
3 | eel, scorpion, hammerhead | other animals
4 | minivan, ambulance, taxi | land vehicles
5 | submarine, canoe, pirate | sea vehicles
6 | guacamole, hotdog, banana | food
7 | backpack, pyjama, kimono | clothes
8 | monitor, iPod, photocopier | tech devices
9 | xylophone, harp, trumpet | instruments
Table 1: Super-class division for $E=10$. For each super-class, we list three
randomly chosen class names (which turn out to be semantically related)
together with a possible super-class name.

##  3 Experiments

We now present empirical results on the standard ImageNet-1k classification
benchmark [14]. We train all models from scratch on the ImageNet-1k training
set of 1.28M images, and then evaluate their top-1 accuracy on the held-out
validation set of 50K images. In Sec. 3.1, we first evaluate our proposed
sparse Mobile V-MoE across a range of model scales and show that they achieve
better performance vs. efficiency trade-offs than the respective dense ViT
baselines. In Sec. 3.2, we then conduct several ablation studies to get a
better understanding of the properties of our proposed sparse MoE model design
and training procedure.

###  3.1 Accuracy vs. efficiency across ViT scales

Model | FLOPs | Top-1 Accuracy
---|---|---
|  | Dense | MoE | $\mathbf{\Delta}$
12$\times$384 | 2297M | 71.88 | 74.23 | +2.35
9$\times$384 | 1752M | 69.94 | 72.47 | +2.53
6$\times$384 | 1207M | 63.21 | 66.91 | +3.70
12$\times$192111This corresponds to the ViT-Tiny model [17] with patch size $32\times 32$. | 618M | 59.51 | 62.90 | +3.39
9$\times$192 | 478M | 56.50 | 59.52 | +3.02
6$\times$192 | 338M | 51.18 | 55.69 | +4.51
12$\times$96 | 176M | 53.79 | 55.39 | +1.60
9$\times$96 | 140M | 51.27 | 52.99 | +1.72
6$\times$96 | 103M | 46.54 | 50.28 | +3.74
12$\times$64 | 88M | 42.90 | 46.07 | +3.17
9$\times$64 | 71M | 40.46 | 43.95 | +3.49
6$\times$64 | 54M | 36.64 | 41.30 | +4.66
(a) Accuracy vs. efficiency across ViT scales.

$\mathbf{E}$ | Router | MoE | $\mathbf{\Delta}$
---|---|---|---
5 | 86.43 | 72.33 | +1.64
7 | 87.43 | 73.13 | +2.44
10 | 87.12 | 73.52 | +2.83
15 | 84.16 | 73.10 | +2.41
20 | 84.08 | 73.36 | +2.67
(b) Total number of experts $E$.

$\mathbf{L}$ | Router | MoE | $\mathbf{\Delta}$
---|---|---|---
1 | 90.17 | 72.14 | +1.45
2 | 87.12 | 73.52 | +2.83
4 | 82.12 | 71.67 | +0.98
6 | 77.70 | 70.07 | -0.62
8 | 72.09 | 64.47 | -6.22
(c) Number of MoE layers $L$.

$\mathbf{k}$ | FLOPs | Dense | MoE | $\mathbf{\Delta}$
---|---|---|---|---
1 | 2534M | 70.79 | 73.44 | +2.65
2 | 2769M | 71.70 | 74.42 | +2.72
3 | 3005M | 73.58 | 74.44 | +0.86
5 | 3476M | 74.87 | 74.32 | -0.55
10 | 4653M | 75.10 | 74.37 | -0.73
(d) Number of experts $k$ per image.

Routing | Input | Acc. | $\mathbf{\Delta}$
---|---|---|---
Dense | N/A | 71.70 |
Super-class | image | 74.42 | +2.72
Rand. class | image | 69.22 | -2.48
End-to-end | image | 73.57 | +1.87
End-to-end | token | 74.85 | +3.15
(e) Routing strategies.

Figure 3: Empirical results. (a) Our Mobile V-MoEs outperform the respective
dense ViTs across model scales. Model names (e.g. 12$\times$192) refer to the
number of layers (12) and the embedding size (192). (b-e) Ablation studies
using DeiT-Ti/16 [17], with $k=1$, $E=10$, $L=2$ by default. Best performance
vs. efficiency trade-off is achieved with (b) $E=10$ experts total, (c) $L=2$
MoE layers (out of 12 layers total), (d) $k=1$ or $k=2$ experts activated per
image, (e) our semantic super-class routing; the settings used in (a) are
bolded.

We consider ViT models (both MoEs and corresponding dense baselines) of
different sizes by scaling the total number of layers (we use 12, 9 or 6) and
the hidden embedding size (we use 384, 192, 96 or 64). The number of multi-
head self-attention heads is (6, 3, 3, 2) for the different hidden embedding
sizes. The embedding size of the MLP is $4\times$ the hidden embedding size,
as is common practice. We use $E=10$ experts in total for the MoE, out of
which $k=1$ is activated per input image. Our MoEs comprise of $L=2$ MoE-ViT
layers preceded by (10, 7 or 4) dense ViT layers (see Fig. 2). We use a patch
size of $32\times 32$ for all models. This is because the the patch size
effectively controls the trade-off between FLOPs and number of model
parameters: as we aim to optimize for FLOPs, a larger patch size (resulting in
a fewer number of patches) is beneficial. We also tried using a smaller patch
size of $16\times 16$, where the result trends were basically the same (but
where the number of FLOPs was higher relative to the model capacity and thus
accuracy). For the ViTs with hidden sizes 384 and 192, we use the DeiT
training recipe [17], while for hidden sizes 96 and 64, we use the standard
ViT training recipe [2] to avoid underfitting. Figs. 1 and 3(a) compare top-1
validation accuracy vs. FLOPs. Our Mobile V-MoEs outperform the corresponding
dense ViT baselines across all model sizes.

###  3.2 Ablation studies

We train DeiT-Tiny [17] (12 layers total, 192 embedding size, $16\times 16$
patch size) with $k=1$ out of $E=10$ experts per input, and with $L=2$ MoE
layers (unless noted otherwise); the dense ViT baseline achieves 70.79%
accuracy.

Total number of experts $E$. We consider different widths of the MoE, i.e.,
different numbers of experts $E$ (and thus super-classes), ranging between
$E=5$ and $E=20$. We report both the accuracy of the entire MoE model (i.e.,
on the 1,000-way classification task), as well as the accuracy of the router
(i.e., on the $E$-way super-classification task). Fig. 3(b) shows that overall
performance improves until $E=10$, from which point onwards it stagnates. The
router accuracy also drops beyond $E=10$ due to the increased difficulty of
the $E$-way super-classification problem.

Number of MoE layers $L$. We consider different depths of the MoE, i.e.,
different numbers of MoE layers $L$, ranging between $L=1$ and $L=8$ (out of
12 ViT layers in total). We again report both the full MoE and router
accuracies. Fig. 3(c) shows that overall performance peaks at $L=2$, and
rapidly decreases for larger $L$. This is due to the router accuracy, which
declines with increasing $L$ as the router gets less information (from the
$12-L$ ViT layers).

Number of experts $k$ per image. We vary the number of experts $k$ activated
per image. We compare against dense baselines that use an MLP with hidden
dimension scaled by $k$ to match the MoE’s inference FLOPs. Fig. 3(d) shows
that $k=1$ and $k=2$ perform best (relative to the dense baseline), with
decreasing performance delta for larger $k$.

Routing strategies. We compare our proposed semantic super-class per-image
routing vs. end-to-end-learned routing (both per-image and per-token) and a
baseline with random super-classes (for $k$=2). Fig. 3(e) shows that our
method (Fig. 2) is better, except for learned per-token routing (as in the
regular V-MoE [13], Fig. 2), which however needs to activate many more experts
and thus model parameters for each input image (up to 11.05M, vs. 6.31M for
ours).

##  4 Conclusions and future work

We showed that sparse MoEs can improve the performance vs. efficiency trade-
off compared to dense ViTs, in an attempt to make ViTs more amenable to
resource-constrained applications. In the future, we aim to apply our MoE
design to models that are more mobile-friendly than ViTs, e.g., light-weight
CNNs such as MobileNets [6, 15, 5] or ViT-CNN hybrids [1, 11, 18]. We also aim
to consider other vision tasks, e.g., object detection. Finally, we aim to get
actual on-device latency measurements for all models.

## References

  * [1] Yinpeng Chen, Xiyang Dai, Dongdong Chen, Mengchen Liu, Xiaoyi Dong, Lu Yuan, and Zicheng Liu.  Mobile-former: Bridging mobilenet and transformer.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5270–5279, 2022.
  * [2] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [3] William Fedus, Jeff Dean, and Barret Zoph.  A review of sparse expert models in deep learning.  arXiv preprint arXiv:2209.01667, 2022.
  * [4] William Fedus, Barret Zoph, and Noam Shazeer.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  J. Mach. Learn. Res, 23:1–40, 2021.
  * [5] Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, et al.  Searching for mobilenetv3.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 1314–1324, 2019.
  * [6] Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  arXiv preprint arXiv:1704.04861, 2017.
  * [7] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton.  Adaptive mixtures of local experts.  Neural computation, 3(1):79–87, 1991.
  * [8] Ruochun Jin, Yong Dou, Yueqing Wang, and Xin Niu.  Confusion graph: Detecting confusion communities in large scale image classification.  In IJCAI, pages 1980–1986, 2017.
  * [9] Michael I Jordan and Robert A Jacobs.  Hierarchical mixtures of experts and the em algorithm.  Neural computation, 6(2):181–214, 1994.
  * [10] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen.  Gshard: Scaling giant models with conditional computation and automatic sharding.  arXiv preprint arXiv:2006.16668, 2020.
  * [11] Sachin Mehta and Mohammad Rastegari.  Mobilevit: light-weight, general-purpose, and mobile-friendly vision transformer.  arXiv preprint arXiv:2110.02178, 2021.
  * [12] Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby.  Multimodal contrastive learning with limoe: the language-image mixture of experts.  arXiv preprint arXiv:2206.02770, 2022.
  * [13] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby.  Scaling vision with sparse mixture of experts.  Advances in Neural Information Processing Systems, 34:8583–8595, 2021.
  * [14] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  International journal of computer vision, 115:211–252, 2015.
  * [15] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4510–4520, 2018.
  * [16] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  arXiv preprint arXiv:1701.06538, 2017.
  * [17] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  In International conference on machine learning, pages 10347–10357. PMLR, 2021.
  * [18] Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, and Anurag Ranjan.  An improved one millisecond mobile backbone.  arXiv preprint arXiv:2206.04040, 2022.
  * [19] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [20] Zhao You, Shulin Feng, Dan Su, and Dong Yu.  Speechmoe: Scaling to large acoustic models with dynamic routing mixture of experts.  arXiv preprint arXiv:2105.03036, 2021.
  * [21] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus.  Designing effective sparse expert models.  arXiv preprint arXiv:2202.08906, 2022.


# Local and correlated studies of humidity-mediated ferroelectric thin film
surface charge dynamics

Iaroslav Gaponenko  Department of Quantum Matter Physics, University of
Geneva, 1211, Geneva, Switzerland  G.W. Woodruff School of Mechanical
Engineering, Georgia Institute of Technology, Atlanta, GA, 30332, USA
<EMAIL_ADDRESS>Loïc Musy  Department of Quantum Matter Physics,
University of Geneva, 1211, Geneva, Switzerland  Neus Domingo  Catalan
Institute of Nanoscience and Nanotechnology (ICN2), CSIC and BIST, Campus UAB,
Bellaterra, 08193 Barcelona, Spain  Nicolas Stucki  University of Applied
Sciences and Arts Western Switzerland (HES-SO/HEPIA), 1213 Geneva, Switzerland
Albert Verdaguer  Institut de Ciència de Materials de Barcelona (ICMAB), CSIC,
Campus UAB, Bellaterra, 08193 Barcelona, Spain  Nazanin Bassiri-Gharb  G.W.
Woodruff School of Mechanical Engineering, Georgia Institute of Technology,
Atlanta, GA, 30332, USA  School of Materials Science and Engineering, Georgia
Institute of Technology, Atlanta, GA, 30332, USA  Patrycja Paruch  Department
of Quantum Matter Physics, University of Geneva, 1211, Geneva, Switzerland

###### Abstract

Electrochemical phenomena in ferroelectrics are of particular interest for
catalysis and sensing applications, with recent studies highlighting the
combined role of the ferroelectric polarisation, applied surface voltage and
overall switching history. Here, we present a systematic Kelvin probe
microscopy study of the effect of relative humidity and polarisation switching
history on the surface charge dissipation in ferroelectric Pb(Zr0.2Ti0.8)O3
thin films. We analyze the interaction of surface charges with ferroelectric
domains through the framework of physically constrained unsupervised machine
learning matrix factorization, Dictionary Learning, and reveal a complex
interplay of voltage-mediated physical processes underlying the observed
signal decays. Additional insight into the observed behaviours is given by a
Fitzhugh-Nagumo reaction-diffusion model, highlighting the lateral spread and
charge passivation process contributors within the Dictionary Learning
analysis.

## Introduction

Surface and bulk electrochemical phenomena are particularly important in
ferroelectrics, whose switchable remanent polarisation modulates their
electrochemical reactivity, determining fundamental material responses and
giving rise to promising catalysis and sensing applications [1, 2].
Electrochemistry has been linked to all aspects of ferroelectric switching and
screening [3, 4, 5], with recent studies highlighting its influence on
catalytic effects and functionality-modifying phenomena [6, 7]. These effects
can be amplified by the additional presence of adsorbates, such as surface
water ubiquitous under ambient conditions [8, 9].

As ferroelectric surfaces exhibit varying polarisation orientation and
chemical composition, the interplay of electrochemistry with surface water
generates a wide range of complex phenomena. Structurally, preferential
adsorption can modulate the usual humidity dependent molecular arrangement
observed on simple dielectric materials [10] of 1–2 monolayers of ice-like
water, followed by less ordered liquid layers. Chemically, varying
dissociation rates into OH- and protons vs. molecular water and even
modulation of the surface and bulk material composition as a function of
polarisation orientation and relative humidity have been reported [11, 12].

The observed effects are further enhanced under the application of an electric
field, with radical water-mediated changes and time-dependent processes taking
place, as also seen in non-polar materials [13]. In addition, sufficiently
high fields will enhance oxidation reactions at the surface and promote charge
injection and/or polarisation switching [14, 15, 16], leading to a strong
dependence of electrochemical response on the electrostatic (switching and
charging) history of the material. The combination of all these phenomena
gives rise to a highly nontrivial mixed response of ferroelectrics to changes
in relative humidity or surface bias, with widely ranging timescales and
magnitudes, as demonstrated previously by local probe surface potential
imaging [12, 17].

Over the last decade, investigation of complex phenomena with multiple
parameters have been made possible by incremental advances in experimental
techniques and acquisition hardware, yielding rich datasets of correlated
responses at a broad range of timescales. Additionally, tremendous progress
has been made in tools designed to analyze the resulting hyperdimensional
datasets, based on Big Data approaches such as machine learning and
dimensional reduction. These have been successfully applied to studies of
ferroelectric materials, enabling a better understanding of the competing and
correlated processes involved in switching and surface screening under ambient
conditions [18, 19, 20].

Here, we investigate the mechanisms of surface charge dissipation on
ferroelectric thin film surfaces as a function of time, applied voltage,
writing history from as-grown through final polarisation, and relative
humidity. Local probe imaging of the surface potential was performed on two
Pb(Zr0.2Ti0.8)O3 (PZT) thin films of opposite polarisation, after the writing
of a predefined structure. The resulting time-dependent datasets were
dimensionally-stacked and analyzed by means of Dictionary Learning in order to
unravel the underlying physical and chemical behaviours. A simplified
reaction-diffusion model complements the experimental work, allowing the
results of the machine learning analysis to be interpreted in terms of
physical/chemical processes. Our observations unequivocally demonstrate that
whilst both positive and negative surface charges decay with time, only
negative charges spread laterally on the surface, diffusing far beyond the
originally written areas.

## Results and discussion

### Imaging surface charge dissipation

Surface charge dissipation was imaged by Kelvin probe force microscopy (KPFM)
on 100 nm thick PZT films (presenting either up- or down-oriented as-grown
monodomain polarisation) and tracked as a function of relative humidity,
polarisation, applied voltages, and time. Domain structures were patterned by
scanning the biased atomic force microscope tip in contact with the surface:
first, to define three lateral stripe regions with positive (+8 V), zero, and
negative (-8 V) tip bias; then, three overlaid vertical stripe regions, again
with positive, zero, and negative tip bias, as illustrated in Fig. 1 (a). The
sequential writing is performed in order to take into account possible effects
of polarisation switching history and varying intensity of charge injection
(areas scanned twice with negative or positive voltages). Once the writing was
completed, KPFM imaging was carried out until dissipation of the surface
potential contrast, approximately 9 to 12 hours. An illustration of resulting
KPFM image is shown in Fig. 1 (b), with the writing structure overlaid. We
note that whilst the piezoresponse force microscopy image of the ferroelectric
domains generated through the writing process follows the writing structure
exactly, the KPFM usually extends outwards due to charge diffusion, especially
at high humidity. 111We note that a 20% overscan generated by the atomic force
microscope used in this study, required for the operation of the closed loop
scanner, results in the extension of the stripe regions along their length,
giving rise to additional sharp contrast regions outside the formally defined
writing structure. The procedure was repeated at each relative humidity
setpoint (ten different values spanning 6% to 80%), for both samples (up- and
down-oriented as-grown polarisation). The control of relative humidity was
performed with a in-house-built flow-based, low noise humidity controller,
which enables fast and precise in-situ control in the atomic force microscope
chamber [21, 22].

Figure 1: Experimental protocol for mapping charge dissipation on the
ferroelectric thin film surface. In each sample and at each humidity setpoint,
a biased atomic force microscope tip was scanned over the surface: first
defining lateral stripe regions with positive (+8 V), zero, and negative bias
(-8 V) at 0∘; then defining vertical stripe regions with identical positive,
zero, and negative bias at 90∘. (a) Schema of the target writing pattern. (b)
Typical surface potential (KPFM) image of the resulting structure, with
overlaid writing bias pattern. (c) Ferroelectric domain structure resulting
from voltage application on the down-polarised sample (PFM phase image).

The first five resulting surface potential images (approximately 70 minutes
after writing) are shown in Fig. 2 for both samples at all ten humidity levels
probed. Several trends can be identified from simple visual inspection. First,
in both samples, faster surface charge dissipation and larger spatial spread
of the KPFM signal beyond the boundaries of the stripe domains can be observed
at higher humidities. Second, the positive charges seem to be much more
localized (bright yellow contrast) than the negative charges (dark blue
contrast), regardless of the initial sample polarisation. Finally, the charge
contrast seems to be more intense in areas where the polarisation was switched
from its as-grown state - i.e. where negative tip bias was applied for the
down-polarised film and positive tip bias for the up-polarised film - hinting
at the role of switching history in the dissipation. These observations are in
line with previous investigations of the dynamics of charged species on
ferroelectric thin films [12], showing a strong effect of relative humidity on
charge dissipation. However, the additional parameter of history yields new
insight: in both samples the first writing (horizontal stripes) can be readily
distinguished at low humidities, but is not discernible at high humidities,
suggesting faster dynamics in presence of high water content at the surface.

Figure 2: First five KPFM images for the complete humidity dataset, as a
function of time and humidity for the (a) down-polarised and (b) up-polarised
samples. A new structure is written at each relative humidity setpoint and
then monitored continuously by KPFM until complete surface potential
dissipation. At low humidities the charges seem to be more localized and
retained longer than at high humidities. Additionally, at higher humidities
the horizontal lines corresponding to the first writing step are not visible,
suggesting much faster charge dynamics in presence of high water content at
the surface. Lastly, the written pattern remains visible in the up-polarised
films at much higher humidity levels than in the down-polarised films.

### Dictionary Learning analysis

Due to the complex and correlated nature of the full dataset, any analysis
beyond simple visual inspection is challenging as it will introduce human
error and bias: regions over which observed behaviour is to be averaged may be
selected arbitrarily, outliers and/or unexpected data might be ignored, and
overall the spatial distribution of information within the images will be lost
as most analysis would concentrate on average signals. To address these
issues, machine learning-based techniques were leveraged for the correlative
analysis of the acquired data.

The KPFM images for all ten humidities and two samples (one up-polarised and
one down-polarised) were stacked along the spatial dimension in order to
induce a time correlation, and are shown in the Supplementary Materials,
Figure S1. Such stacking imposes a constraint of identical behaviour(s) across
polarisation, time and humidity values [20], thus enabling a direct comparison
of the evolution of the behaviours as a function of initial sample
polarisation and relative humidity. Dictionary Learning was selected as the
machine learning dimensional reduction technique for the further analysis of
the dataset due to its intrinsically sparse nature [23], with the expectation
of only a few behaviours being spatially co-located. The number of components
$n$ is a free parameter and was evaluated from two to ten, allowing us to
track which behaviours are associated with the different regions as they
successively segregate out, as detailed in the Supplementary Materials, Figure
S3.

The resulting decomposition, here presented for $n=10$, shows three distinct
types of components, each associated with the relevant spatial distribution
weight maps. The first component, C1, has a low magnitude and is spatially
located in the regions that can be identified as a background - i.e.,
locations outside of the written structure or where zero tip bias was applied
during writing - and is identifiable as a distinct component for $n\geq 8$.
This component is also present within the written areas at higher humidities,
and specifically in the locations where the surface charge has already
dissipated. An inspection of its time-evolution indicates that this component
is essentially stable with a very slight linear drift, possibly caused by the
change in tip-sample interactions over time, as shown in the exponential-plus-
linear fit C1 in Fig. 3d. The second component, C2, decays with time from an
initial high positive value, and is present in areas of positive voltage
application only, regardless of the initial polarisation state of the films.
These areas separate out as distinct from those written with negative bias at
$n=6$. Increasing $n$ does not change the C2 component. The last three
components present negative decays, starting at a low surface potential value
and evolving towards the background potential as the surface charges
dissipate. These three components are co-located within each other and spread
outwards around areas of negative voltage application, with the highest
intensity decay (C3) on the inside, and the lowest intensity decay (C5) on the
outside of the negative charge distributions - a structure reminiscent of
Matryoshka dolls, and suggesting the possible presence of lateral diffusion.
This characteristic structure is apparent from $n=7$, and increasing $n$
simply refines the areas written with negative bias into increasingly fine
concentric segments.

To quantify the time evolution of the positive and negative decays, the
corresponding components - C2 to C5 - were fitted to a double exponential
decay function, $CPD(t)=A_{0}+A_{1}\exp{t/\tau_{1}}+A_{2}\exp{t/\tau_{2}}$,
with $CPD$ the contact potential difference, and $A_{0}$ the offset
corresponding to the final CPD value, $A_{1},A_{2}$ the amplitudes, and
$\tau_{1},\tau_{2}$ the time constant of the decays. From the parameters shown
in Fig. 3d, two behaviors seem to be present in all decaying components, each
with two substantially different time constants. The latter differ by at least
an order of magnitude, as reported previously for this material, indicating
the coexistence of a material-dependent fast electrochemical process and a
slow process mediating the surface charge dissipation [12]. The fast process,
of the order of 1000 seconds, can be related to the re-equilibration of
surface screening charges [24], whereas the slow process, of the order of
10000 seconds, is specifically attributed to the dynamics of tip-injected
surface or near-surface charges.

In all cases, however, the overall mechanisms (background, positive and
negative decay) are present in both films and share the same parameters,
varying in intensity with humidity and initial polarisation, but mostly
depending on the polarity of the applied voltage. The positive voltage related
component, C2, for instance is found in both films below 50% RH, but seems to
be more present at the higher humidity only in the up-polarised films, where a
positive voltage application switches the polarisation. Conversely, negative
voltage related components, C3, C4, and C5 seem to be more present within the
down-polarised film, where such a voltage polarity switches the polarisation.
The coexistence of the same physical dissipation processes (as fitted above)
in both samples suggests that the initial polarisation orientation does not
affect the dissipation mechanisms which seem to be primarily dominated by
voltage and surface electrochemistry. However, the history of polarisation is
important, as switching processes during voltage application will result in
higher surface charge density and/or longer charge retention, especially at
higher humidities.

Figure 3: Dictionary Learning spectral decomposition of the complete dataset
of surface potential evolution measured in the two samples and at ten
humidities into five components. The weight maps for the (a) down-polarised
and (b) up-polarised sample demonstrate spatially localized features, which
once correlated with the (c) time dependent components yield a clear
separation between the background response (C1), and the areas of positive
(C2) and negative (C3-5) charge. The negative charge region is split into co-
located concentric structures, reminiscent of Matryoshka dolls. Components C2
– C5 can be fitted with double exponential decays, each yielding (d) two very
different time constants.

### Reaction-diffusion simulation

To better understand the specific behaviors observed in the experimental data
and the resulting Machine Learning analysis, a simplified reaction-diffusion
model was constructed, based on the FitzHugh-Nagumo system. This is an
excitable system model, ideal for describing the distributions of different
species during their relaxation to equilibrium, as for example in neural
synapses [25]. As the experimental data suggests the presence of lateral
spread acting together with surface and bulk electrochemistry, the parameters
of the model were tuned to reproduce the observations qualitatively, with a
spatially localized and negligibly diffusive positive charge distribution, and
an initially more spread and highly diffusive negative charge distribution.

The resulting time evolution of the combined distributions is shown in Fig.
4a. As a function of time, the intensity of both distributions appears to
decay, as illustrated in the plot of the evolution of each distribution’s
center point intensity. However, the difference in diffusivity of the two
distributions is apparent when a cross-section is taken at the center of each
image. Whilst the positive charges (red) seem to simply decay (i.e. reduce in
intensity as a function of time), the negative charge distribution spreads
outwards via lateral diffusion, as indicated by the characteristic crossing
points.

Figure 4: FitzHugh-Nagumo reaction-diffusion model of the surface charge
dissipation. (a) Time evolution of the positive (red) and negative (blue)
surface charge distributions shows spatially localised vs. highly diffusive
behavior, qualitatively reproducing the experimental observations. Maxima (b)
and cross-sections (c) of the positive and negative charge distributions show
that both positive and negative charge intensities decay with time, but that
only negative charges show a lateral diffusion, as evidenced by the increased
width of the signal, leading to characteristic crossing points (green
circles).

The model data were subjected to the same computational treatment as the
experimental data with Dictionary Learning analysis. The resulting components
and weight maps are shown in Fig. 5, separated into behaviors for positively
and negatively charged regions. For the positive regions, only simple decays
are observed, with the second component identifying mostly the interaction of
charges at the boundary between positively and negatively charged regions. In
comparison, the negative charge decay shows a striking resemblance to the
experimental data analysis, with concentric shell structures in the weight
maps.

Lastly, the components associated with the weight maps are also consistent
with the original experimental observations: the most and least intense
behaviours are located at the center and at the periphery of each charged
regions, respectively. Per model parameters, the species that shows this
behavior is associated with a reaction and diffusion process - as opposed to
the other species which is constrained to reaction only. Thus, the
introduction of diffusion into the model brings out the characteristic
concentric shell structure of decreasing intensity in the Dictionary Learning
analysis. This result confirms the intuition that lateral diffusion is taking
place within our experimental data based on the similarity to the machine
learning analysis features, making this - to the authors’ knowledge - the
first reported observation of a signature of such a behaviour observed by
machine learning methods.

Figure 5: Dictionary Learning analysis of Reaction-Diffusion analysis. The
maps shown in (a) and corresponding components shown in (b) demonstrate two
specific behaviours: positive and negative decays. The latter show a
decreasing amplitude following a characteristic concentric structure,
consistent with the presence of a diffusion process.

## Conclusion

In conclusion, in this investigation of the interplay of time, relative
humidity, polarisation and voltage with the surface charge dissipation of
ferroelectric thin films, polarised regions written by scanning probe
microscopy were analyzed with Dictionary Learning spectral decomposition. The
results highlight two separate voltage-mediated behaviours, characteristic of
a combined surface/bulk electrochemical process and of surface charge
passivation. Negative charges are additionally observed to diffuse laterally,
as demonstrated by the characteristic concentric shell structure in machine
learning analysis, while positive charges are mostly confined to their
original location with very limited lateral diffusion. Our interpretation is
corroborated by a reaction-diffusion model set up to reproduce the
experimental KPFM observations, whose resulting time evolution was analyzed
with the same spectral decomposition technique, yielding similar
characteristic concentric structures in the case of the diffusing species. The
cumulative evidence points to a non-trivial behaviour of surface charges, with
the positive and negative charges exhibiting a substantially different
diffusion/dissipation mechanism that is modulated by the underlying
ferroelectric polarisation. The coexistence of these electrochemical processes
opens a pathway towards polarisation-directed catalysis, with a highly
localized control of charges.

## Materials and methods

### Materials growth

Films of 100 nm thick Pb(Zr0.2Ti0.8)O3 (PZT) have been grown on Nb:SrTiO3
single crystalline substrates by RF-magnetron sputtering, as described
elsewhere [26]. The films used in this study were up- and down-polarised as
grown, respectively, with a smooth sub-nanometer roughness topography.

### Kelvin Probe Force Microscopy

The scanning probe microscopy imaging was performed on an Asylum Research
Cypher microscope equipped with an environmental control chamber, using Bruker
SCM-PtSi cantilevers with a resonance frequency of $\sim 75kHz$ and a spring
constant of $\sim 2.8N/m$. Sample temperature was kept constant at
$30^{\circ}C$, and relative humidity was varied through the use of a home-
build controller [21, 22]. The humidity was stabilized for at least an hour
prior to each measurement series. Topographic imaging was performed in
intermittent contact mode (AC Mode), with a free oscillation amplitude of
$\sim 10nm$ and a setpoint of $\sim 5nm$. The contact potential difference
between the tip and the sample was acquired over $12x12\mu m$ areas at a scan
speed of $0.5Hz$ in single pass Kelvin probe force microscopy mode with the
following parameters: $3V$ excitation voltage, $5kHz$ frequency, $-90^{\circ}$
phase, $5$ proportional gain, $10$ integral gain. Domain writing was done in
contact mode with a $0.5V$ setpoint over $8x8\mu m$ areas at a scan speed of
$0.5Hz$ and applied voltages of $\pm 8V$, as detailed in Fig. 1(a).

### Dictionary Learning analysis

Time-dependent data was stacked into data cubes, and the data cubes at each
humidity and for each sample were concatenated along the spatial dimension x,
yielding a correlated dataset. A similar stacking method was performed for the
reaction-diffusion modeling dataset. The Dictionary Learning implementation of
the scikit-learn library was then used on the flattened data set with the
following parameters: $3-10$ n_components, $1000000$ max_iter, $8$ n_jobs,
lasso_lars transform_algorithm, $15$ transform_n_nonzero_coefs, lars
fit_algorithm, and $0$ random_state. The code and data used for the analysis
is freely available on the paper repository [CITE yareta_or_zenodo].

### Reaction-diffusion modeling

The reaction-diffusion model used is based on a 2D Fitzhugh-Nagumo excitable
system model [25]. The implementation used is adapted from existing Python
code, recipe 12.4 in [27]. The specific form of the model is a set of two
differential equations:

$\frac{\partial u}{\partial t}=a\Delta u+u-u^{3}-v+k$ $\tau\frac{\partial
v}{\partial t}=b\Delta v+u-v$

with $u,v$ denoting the two species, $a,b$ their corresponding diffusion
factors, $k$ the intensity of the external stimulus, and $tau$ the time
constant providing the separation of the timescales of the evolution of the
two species. For the comparison to surface charge dissipation, the two species
were initially set up in separate bands and allowed to relax. The parameter
set, selected up to allow for single species diffusion only, is: $a=0$,
$b=0.1$, $\tau=0.1$, and $k=-.005$. The resulting spatiotemporal evolution was
analyzed through the same Dictionary Learning workflow as the experimetnally
obtained data set. The code and data used for the simulation and analysis is
freely available on the paper repository [CITE yareta_or_zenodo].

## References

  * [1] Hwang, J. _et al._ Tuning perovskite oxides by strain: Electronic structure, properties, and functions in (electro)catalysis and ferroelectricity.  _Materials Today_ DOI: https://doi.org/10.1016/j.mattod.2019.03.014 (2019).
  * [2] Yang, S. M. _et al._ Mixed electrochemical–ferroelectric states in nanoscale ferroelectrics.  _Nature Physics_ 13, 812–818, DOI: 10.1038/nphys4103 (2017).
  * [3] Neumayer, S. M. _et al._ Surface Chemistry Controls Anomalous Ferroelectric Behavior in Lithium Niobate.  _ACS Applied Materials & Interfaces_ 10, 29153–29160, DOI: {10.1021/acsami.8b09513} (2018).
  * [4] Fabiano, S. _et al._ Ferroelectric polarization induces electronic nonlinearity in ion-doped conducting polymers.  _Science Advances_ 3, DOI: 10.1126/sciadv.1700345 (2017).  https://advances.sciencemag.org/content/3/6/e1700345.full.pdf.
  * [5] Ievlev, A. V. _et al._ Nanoscale electrochemical phenomena of polarization switching in ferroelectrics.  _ACS Applied Materials & Interfaces_ 10, 38217–38222, DOI: 10.1021/acsami.8b13034 (2018).  https://doi.org/10.1021/acsami.8b13034.
  * [6] Kakekhani, A. & Ismail-Beigi, S.  Ferroelectric-based catalysis: Switchable surface chemistry.  _ACS Catalysis_ 5, 4537–4545, DOI: 10.1021/acscatal.5b00507 (2015).
  * [7] Garrity, K., Kakekhani, A., Kolpak, A. & Ismail-Beigi, S.  Ferroelectric surface chemistry: First-principles study of the pbtio3 surface.  _Phys. Rev. B_ 88, 045401, DOI: 10.1103/PhysRevB.88.045401 (2013).
  * [8] Tian, Y. _et al._ Water printing of ferroelectric polarization.  _Nature Communications_ 9, DOI: 10.1038/s41467-018-06369-w (2018).
  * [9] Ievlev, A. V. _et al._ Chemical state evolution in ferroelectric films during tip-induced polarization and electroresistive switching.  _ACS Applied Materials & Interfaces_ 8, 29588–29593, DOI: 10.1021/acsami.6b10784 (2016).  https://doi.org/10.1021/acsami.6b10784.
  * [10] Asay, D. B. & Kim, S. H.  Evolution of the adsorbed water layer structure on silicon oxide at room temperature.  _The Journal of Physical Chemistry B_ 109, 16760–16763, DOI: 10.1021/jp053042o (2005).  PMID: 16853134, https://doi.org/10.1021/jp053042o.
  * [11] Cordero-Edwards, K. _et al._ Water affinity and surface charging at the z-cut and y-cut linbo3 surfaces: An ambient pressure x-ray photoelectron spectroscopy study.  _The Journal of Physical Chemistry C_ 120, 24048–24055, DOI: 10.1021/acs.jpcc.6b05465 (2016).  https://doi.org/10.1021/acs.jpcc.6b05465.
  * [12] domingo, n. _et al._ Surface charged species and electrochemistry of ferroelectric thin films.  _Nanoscale_ –, DOI: 10.1039/C9NR05526F (2019).
  * [13] Verdaguer, A. _et al._ Charging and discharging of graphene in ambient conditions studied with scanning probe microscopy.  _Applied Physics Letters_ 94, DOI: {10.1063/1.3149770} (2009).
  * [14] Segura, J. J., Domingo, N., Fraxedas, J. & Verdaguer, A.  Surface screening of written ferroelectric domains in ambient conditions.  _Joural of Applied Physics_ 113, DOI: {10.1063/1.4801983} (2013).  21st IEEE International Symposium on Applications of Ferroelectrics held jointly with 11th European Conference on the Applications of Polar Dielectrics and 4th Conference on Piezoresponse Force Microscopy and Nanoscale Phenomena in Polar Materials, Univ Aveiro, Aveiro, PORTUGAL, JUL 09-13, 2012.
  * [15] Ievlev, A. V., Morozovska, A. N., Shur, V. Y. & Kalinin, S. V.  Humidity effects on tip-induced polarization switching in lithium niobate.  _Applied Physics Letters_ 104, 092908, DOI: 10.1063/1.4867979 (2014).  https://doi.org/10.1063/1.4867979.
  * [16] Ievlev, A. V. _et al._ Intermittency, quasiperiodicity and chaos in probe-induced ferroelectric domain switching.  _Nature Physics_ 10, 59–66, DOI: 10.1038/nphys2796 (2013).
  * [17] Shishkin, E. I., Shur, V. Y., Schlaphof, F. & Eng, L. M.  Observation and manipulation of the as-grown maze domain structure in lead germanate by scanning force microscopy.  _Applied Physics Letters_ 88, 252902, DOI: 10.1063/1.2183369 (2006).  https://doi.org/10.1063/1.2183369.
  * [18] Vasudevan, R. K. _et al._ Multidimensional dynamic piezoresponse measurements: Unraveling local relaxation behavior in relaxor-ferroelectrics via big data.  _Journal of Applied Physics_ 118, 072003, DOI: 10.1063/1.4927803 (2015).  https://doi.org/10.1063/1.4927803.
  * [19] Li, L. _et al._ Machine learning–enabled identification of material phase transitions based on experimental data: Exploring collective dynamics in ferroelectric relaxors.  _Science Advances_ 4, eaap8672, DOI: 10.1126/sciadv.aap8672 (2018).
  * [20] Griffin, L. A., Gaponenko, I., Zhang, S. & Bassiri-Gharb, N.  Smart machine learning or discovering meaningful physical and chemical contributions through dimensional stacking.  _npj Computational Materials_ 5, DOI: 10.1038/s41524-019-0222-z (2019).
  * [21] Gaponenko, I., Gamperle, L., Herberg, K., Muller, S. C. & Paruch, P.  Low-noise humidity controller for imaging water mediated processes in atomic force microscopy.  _Review of Scientific Instruments_ 87, 063709, DOI: 10.1063/1.4954285 (2016).  https://doi.org/10.1063/1.4954285.
  * [22] Gaponenko, I., Musy, L., Muller, S. C. & Paruch, P.  Open source standalone relative humidity controller for laboratory applications.  _Engineering Research Express_ 1, 025042, DOI: 10.1088/2631-8695/ab5771 (2019).
  * [23] Dumitrescu, B. & Irofti, P.  _Dictionary Learning Algorithms and Applications_ (Springer International Publishing, 2018).
  * [24] Tong, S. _et al._ Mechanical removal and rescreening of local screening charges at ferroelectric surfaces.  _Physical Review Applied_ 3, DOI: 10.1103/physrevapplied.3.014003 (2015).
  * [25] Rocşoreanu, C., Georgescu, A. & Giurgiţeanu, N.  _The FitzHugh-Nagumo Model_ (Springer Netherlands, 2000).
  * [26] Gariglio, S., Stucki, N., Triscone, J.-M. & Triscone, G.  Strain relaxation and critical temperature in epitaxial ferroelectric pb(zr0.20ti0.80)o3 thin films.  _Applied Physics Letters_ 90, 202905, DOI: 10.1063/1.2740171 (2007).  https://doi.org/10.1063/1.2740171.
  * [27] Rossant, C.  _IPython Interactive Computing and Visualization Cookbook, Second Edition : Over 100 hands-on recipes to sharpen your skills in high-performance numerical computing and data science in the Jupyter Notebook_ (Packt Publishing, Birmingham, 2018).
  * [28] Gaponenko, I. _et al._ Local and correlated studies of humidity-mediated ferroelectric thin film surface charge dynamics, DOI: 10.26037/yareta:wq6me5d75jb6vkz7rhva45ht6i (2020).

## Data availability

Data and analysis code is available at the long term storage Yareta project
repository, hosted at the University of Geneva [28].

## Acknowledgements

The authors acknowledge Dr Sergei V. Kalinin of Oak Ridge National Laboratory,
for helpful discussions about machine learning and the initial suggestion to
explore reaction-diffusion modeling. This work was supported by Division II of
the Swiss National Science Foundation under project 200021_178782.

## Author contributions

The authors declare no competing financial interests.

I.G. performed the measurements and simulations on films grown by N.S. I.G.
and L.M. analyzed the data. I.G. and P.P. wrote the manuscript. All authors
contributed to the scientific discussion and manuscript revisions.

## Supplementary Materials

### Complete humidity-sample-time dataset

The complete humidity-sample-time dataset is shown in Fig. S1, for the two
samples (up and down polarised) and ten humidities.

Figure S1: Complete KPFM dataset as a function of sample, relative humidity
and time.

### Spectral decomposition principle

The principle behind spectral decomposition methods such as DictionaryLearning
is shown in Fig. S2. The complete dataset is split into a set of one-
dimensional components with corresponding two-dimensional weight maps.
Dimensional stacking was used to integrate the relative humidity and initial
sample polarisation dependence into the maps, thus correlating the data sets
and enabling quantitative comparison of time-dependent decays.

Figure S2: Principle behind machine learning based spectral decomposition.

### Dependence on the number of components

A systematic investigation of the influence of the number of components on the
DictionaryLearning analysis is shown in Fig. S3. As the number of components
increases, so does the number of physical behaviours. At first, when $n=2$,
only a single physical behaviour is visible, alongside a Gaussian noise
component with a null weight map. When the number of allowed components is
increased up to $n=6$, the negative voltage process splits from the main
component. Increasing the number of components to $n=7$, the concentric
structure for the negative voltage region starts to form with the split of the
latter into two distinct map/component pairs. At $n=8$, a background component
splits from the positive voltage component, and finally at $n=10$, a third
concentric negative component makes its appearance.

Figure S3: Dictionary Learning based spectral decomposition of the correlated
KPFM dataset as a function of the number of components. Increasing number of
components refines the separation into further sub-processes.

### Reaction-diffusion modeling

The reaction-diffusion modeling, here based on a Fitzhugh-Nagumo excitation-
relaxation process, provides insight into the contribution of the various
physical processes to the machine learning based analysis. Such a sandbox
approach enables a unique way to tailor the interactions present within the
two species. In particular, the reaction-diffusion modeling highlights the
importance of sparse approaches such as DictionaryLearning when dealing within
physical materials presenting very different co-located behaviors.

#### Initial conditions choice

Initial species distributions were chosen to mimic the positive and negative
charges as deposited by biased-tip atomic force microscopy writing. The
positive charge species $u$ shown in Fig. S4(a) was initialized with only a
reaction interaction with the negative charge species $v$, and no diffusion
component. The negative charge species $v$, shown in Fig. S4(b) is susceptible
to both reaction and diffusion processes. The total distribution $u-v$ is
shown in Fig. S4(c), and qualitatively resembles the Kelvin probe force
microscopy data in Fig. 2.

Figure S4: Initial distributions of the two reacting/diffusing species (a) $u$
and (b) $v$. The resulting potential is computed in (c) as $u-v$ and processed
further with machine learning techniques.

#### Principal component analysis

Whilst principal component analysis (PCA) does not provide a physically
interpretable spectral decomposition, it can be used to estimate the
complexity of the decomposition. The PCA weight maps for the first 50
components are shown in Fig. S5(a), illustrating the concentric-like pattern
generated by the diffusive nature of the negative charge species distribution
$v$. Indeed, the positive charge species distribution $u$ is only present in
the first four maps of the decomposition. The corresponding PCA components
shown in Fig. S5(b) additionally shows a periodic behaviour with increasing
frequency for the first 26 components, as evidenced by their Fourier transform
in Fig. S5(c).

Figure S5: Principal component analysis of the reaction-diffusion dataset up
to 50 components, with (a) the maps and (b) their corresponding components.
(c) The Fourier transform in the components shows an increasing oscillation
frequency for the increasing order components.

#### Non-negative matrix factorization

A non-negative matrix factorization (NMF) approach was tried concurrent to
DictionaryLearning, with results for $n=4$ and $n=10$ components shown in Fig.
S6(b,c) respectively. Although not able to separate the physical processes as
cleanly as DictionaryLearning, NMF also highlights the concentric
decomposition of the negative voltage species with an increasing number of
components.

Figure S6: Comparison between PCA (a) and NMF (b) $n=4$ and (c) $n=10$
methods. In all cases - just like in the DictionaryLearning analysis - the
negative voltage species generate concentric components due to their diffusive
nature.