a continual learning framework for uncertainty-aware … · 2021. 1. 6. · badrinarayanan, and...

9
A Continual Learning Framework for Uncertainty-Aware Interactive Image Segmentation Ervine Zheng, Qi Yu * , Rui Li, Pengcheng Shi, Anne Haake Rochester Institute of Technology Abstract Deep learning models have achieved state-of-the-art perfor- mance in semantic image segmentation, but the results pro- vided by fully automatic algorithms are not always guar- anteed satisfactory to users. Interactive segmentation offers a solution by accepting user annotations on selective areas of the images to refine the segmentation results. However, most existing models only focus on correcting the current im- age’s misclassified pixels, with no knowledge carried over to other images. In this work, we formulate interactive image segmentation as a continual learning problem and propose a framework to effectively learn from user annotations, aim- ing to improve the segmentation on both the current image and unseen images in future tasks while avoiding deterio- rated performance on previously-seen images. It employs a probabilistic mask to control the neural network’s kernel ac- tivation and extract the most suitable features for segmenting images in each task. We also design a task-aware architec- ture to automatically infer the optimal kernel activation for initial segmentation and subsequent refinement. Interactions with users are guided through multi-source uncertainty esti- mation so that users can focus on the most important areas to minimize the overall manual annotation effort. Extensive experiments are performed on both medical and natural im- age datasets to illustrate the proposed framework’s effective- ness on basic segmentation performance, forward knowledge transfer, and backward knowledge transfer. Introduction Deep learning (DL) based methods have been increasingly applied to semantic image segmentation in recent years with superior performance (Long, Shelhamer, and Darrell 2015; Ronneberger, Fischer, and Brox 2015; Chen et al. 2017; Milletari, Navab, and Ahmadi 2016). However, large-scale annotated images are usually required to properly train a complex DL model to ensure a good segmentation perfor- mance. Fully annotating an image is a lengthy and laborious process. In many specialized domains (e.g., medicine and biology), image annotation may only be performed by ex- pert users with special expertise. Thus, the availability of large-scale fully annotated images is usually rather limited in these domains, which poses a key challenge of applying * Corresponding author. Copyright © 2021, Association for the Advancement of Artificial Intelligence (www.aaai.org). All rights reserved. state-of-the-art DL models. Furthermore, segmentation re- sults provided by fully automatic algorithms are not always guaranteed satisfactory to users. For example, in medical image analysis, physicians may form different opinions on the boundary of segments. It is vital to provide an effective means for users to adjust the segmentation results. Finally, domain shift may cause severe deterioration of segmentation results if training images vary significantly from the testing ones due to different imaging devices, environments, and other factors, which may frequently occur in many appli- cations. Interactive segmentation offers a promising direction by leveraging users’ knowledge to refine the segmentation re- sults. It allows users to interact with the system by annotat- ing a few pixels, which are then used by the model to update the segmentation results. However, most existing algorithms process images and user annotations as isolated cases. They mainly focus on correcting the current image’s misclassified pixels but ignore carrying over knowledge learned from user annotations to other images. In a practical setting, new im- ages are usually provided to a model in sequential order. The model should effectively learn from user annotations to im- prove segmentation performance and reduce users’ (repeti- tive) annotation effort in the future. Besides, most existing algorithms leave the decision of selective annotation alto- gether to the users, i.e., users may make edits on any areas without knowing which edit is most useful to the model for improving the results. A fundamental challenge of interac- tive segmentation is how to provide users informative guid- ance to improve the segmentation accuracy with minimum annotation effort. To address the above challenges, we formulate interactive segmentation as a continual learning problem. The segmen- tation tasks are organized as a sequence. In each task, the model is trained with only one or a few images and along with limited user annotations. This is different from con- ventional continual learning, where one task usually con- tains a large batch of data instances, because interactive seg- mentation aims to leverage limited user annotations to im- prove performance. We also design a training-testing pro- tocol to evaluate how a model can effectively learn from user annotations to improve the segmentation performance on 1) the current image and 2) the unseen images in fu- ture tasks, while mitigating the deteriorated performance on PRELIMINARY VERSION: DO NOT CITE The AAAI Digital Library will contain the published version some time after the conference

Upload: others

Post on 10-Sep-2021

2 views

Category:

Documents


0 download

TRANSCRIPT

Page 1: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

A Continual Learning Framework for Uncertainty-AwareInteractive Image Segmentation

Ervine Zheng, Qi Yu *, Rui Li, Pengcheng Shi, Anne HaakeRochester Institute of Technology

Abstract

Deep learning models have achieved state-of-the-art perfor-mance in semantic image segmentation, but the results pro-vided by fully automatic algorithms are not always guar-anteed satisfactory to users. Interactive segmentation offersa solution by accepting user annotations on selective areasof the images to refine the segmentation results. However,most existing models only focus on correcting the current im-age’s misclassified pixels, with no knowledge carried over toother images. In this work, we formulate interactive imagesegmentation as a continual learning problem and propose aframework to effectively learn from user annotations, aim-ing to improve the segmentation on both the current imageand unseen images in future tasks while avoiding deterio-rated performance on previously-seen images. It employs aprobabilistic mask to control the neural network’s kernel ac-tivation and extract the most suitable features for segmentingimages in each task. We also design a task-aware architec-ture to automatically infer the optimal kernel activation forinitial segmentation and subsequent refinement. Interactionswith users are guided through multi-source uncertainty esti-mation so that users can focus on the most important areasto minimize the overall manual annotation effort. Extensiveexperiments are performed on both medical and natural im-age datasets to illustrate the proposed framework’s effective-ness on basic segmentation performance, forward knowledgetransfer, and backward knowledge transfer.

IntroductionDeep learning (DL) based methods have been increasinglyapplied to semantic image segmentation in recent years withsuperior performance (Long, Shelhamer, and Darrell 2015;Ronneberger, Fischer, and Brox 2015; Chen et al. 2017;Milletari, Navab, and Ahmadi 2016). However, large-scaleannotated images are usually required to properly train acomplex DL model to ensure a good segmentation perfor-mance. Fully annotating an image is a lengthy and laboriousprocess. In many specialized domains (e.g., medicine andbiology), image annotation may only be performed by ex-pert users with special expertise. Thus, the availability oflarge-scale fully annotated images is usually rather limitedin these domains, which poses a key challenge of applying

*Corresponding author.Copyright © 2021, Association for the Advancement of ArtificialIntelligence (www.aaai.org). All rights reserved.

state-of-the-art DL models. Furthermore, segmentation re-sults provided by fully automatic algorithms are not alwaysguaranteed satisfactory to users. For example, in medicalimage analysis, physicians may form different opinions onthe boundary of segments. It is vital to provide an effectivemeans for users to adjust the segmentation results. Finally,domain shift may cause severe deterioration of segmentationresults if training images vary significantly from the testingones due to different imaging devices, environments, andother factors, which may frequently occur in many appli-cations.

Interactive segmentation offers a promising direction byleveraging users’ knowledge to refine the segmentation re-sults. It allows users to interact with the system by annotat-ing a few pixels, which are then used by the model to updatethe segmentation results. However, most existing algorithmsprocess images and user annotations as isolated cases. Theymainly focus on correcting the current image’s misclassifiedpixels but ignore carrying over knowledge learned from userannotations to other images. In a practical setting, new im-ages are usually provided to a model in sequential order. Themodel should effectively learn from user annotations to im-prove segmentation performance and reduce users’ (repeti-tive) annotation effort in the future. Besides, most existingalgorithms leave the decision of selective annotation alto-gether to the users, i.e., users may make edits on any areaswithout knowing which edit is most useful to the model forimproving the results. A fundamental challenge of interac-tive segmentation is how to provide users informative guid-ance to improve the segmentation accuracy with minimumannotation effort.

To address the above challenges, we formulate interactivesegmentation as a continual learning problem. The segmen-tation tasks are organized as a sequence. In each task, themodel is trained with only one or a few images and alongwith limited user annotations. This is different from con-ventional continual learning, where one task usually con-tains a large batch of data instances, because interactive seg-mentation aims to leverage limited user annotations to im-prove performance. We also design a training-testing pro-tocol to evaluate how a model can effectively learn fromuser annotations to improve the segmentation performanceon 1) the current image and 2) the unseen images in fu-ture tasks, while mitigating the deteriorated performance on

PRELIMINARY VERSION: DO NOT CITE The AAAI Digital Library will contain the published

version some time after the conference

Page 2: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

3) previously-seen images. We refer to the three objectivesmentioned above as basic performance, forward knowledgetransfer, and backward knowledge transfer.

To our best knowledge, this is the first work that inte-grates interactive segmentation with continual learning. Wepropose a Continual Learning framework for InteractiveSegmentation (referred to as CLIS) to achieve the threeobjectives simultaneously. It employs a task-specific prob-abilistic mask, generated from a Bayesian non-parametricIndian Buffet Process (IBP), to control the network’s con-volutional kernel activation. A different set of kernels areactivated for each task, aiming to extract the most suit-able features for segmenting the corresponding images. Be-sides, the IBP prior encourages kernels frequently activatedin previous tasks to be more likely activated in the currenttask, effectively leveraging knowledge learned from the past.We design a task-aware architecture that leverages the un-annotated image to determine the initial kernel activation.In particular, the image data will feed into a multi-layerperceptron, which can be trained to encode the importantknowledge learned from previous tasks. The kernel activa-tion mask can be inferred with only new image data for ini-tial segmentation. If the results deviate from a user’s opin-ion, s/he may choose to annotate on selective areas. Themodel will be updated based on the annotations and gen-erate refined segmentation results. This interaction contin-ues until users are satisfied with the results, and then moveon to the next task. Furthermore, interactions with users areguided through multi-source uncertainty estimation so thatusers can focus on the most important areas to minimize theoverall manual annotation effort.

Figure 1 illustrates a typical usage of the proposed frame-work in one task and demonstrates how user annotations areleveraged to refine the segmentation on a dermatological im-age. The model consists of three modules: a task-aware seg-mentation module, an uncertainty estimation module, and aninteractive learning module. The segmentation module per-forms pixel-level classification and infers a binary mask thatactivates a subset of the convolutional kernels to extract themost suitable features for segmentation in each task. The un-certainty estimation module generates a map that visualizesthe model’s uncertainty on its prediction. A user may chooseto provide additional annotations to refine the segmentation.As shown in Figure 1 (C), the model makes unconfident pre-dictions on the marked areas. Users are guided to edit thoseareas, and then the interactive learning module propagatesfrom user annotations to other pixels and updates the net-work. This process may iterate until the refined result is ac-ceptable to the user.

The major contribution of this paper is three-fold: (i)formulating interactive segmentation as a continual learn-ing problem and designing an evaluation protocol regardingbasic performance and forward/backward knowledge trans-fer, (ii) a task-aware segmentation framework with a binarymask that activates a subset of the convolutional kernelsto extract the most suitable features in each task, and (iii)introducing uncertainty maps to provide users informativeguidance and effectively improve the segmentation accuracywith minimum annotation effort.

Figure 1: Schematic view of the proposed framework (Top)and an illustrative example of errors in initial segmentationof a dermatological image (Bottom). Different colors in seg-mentation map denotes different classes: healthy skin (gray),crusts (purple), bullaes (blue). Areas with high uncertaintyin the uncertainty map are colored red.

Related WorksContinual Learning Continual learning aims to train mod-els to learn over time by encoding new knowledge while re-taining previously learned knowledge. Catastrophic forget-ting is a long-standing issue for continual learning (Thrunand Mitchell 1995). It usually occurs when new data differssignificantly from previous data and causes learned knowl-edge as the network parameters to be overwritten (French1999). Existing works address this issue through regulariza-tion, isolation, or rehearsal.

Isolation-based methods allocate different neural re-sources to encode the information for different tasks. Forexample, progressive networks (Rusu et al. 2017) freeze net-works modules trained on previous tasks and expand thenetwork architecture with additional modules for new tasks.PackNet (Mallya and Lazebnik 2018) uses iterative pruningto exploit the redundancies in large deep networks and freeup parameters for learning new tasks. However, isolation-based approaches usually lack a principled way to deter-mine the optimal sets of parameter when task informationis unknown for new data, or requires storing all the iso-lated parameters for each task. Regularization-based modelsimpose constraints on updating the neural weights and pe-nalizing the changes. Learning without forgetting (LwF) (Liand Hoiem 2017) uses knowledge distillation to enforce thenetwork predictions with old parameters to be similar tothose with updated parameters. Elastic weight consolidation(EWC) (Kirkpatrick et al. 2017) and Synaptic Intelligence(SI) (Zenke, Poole, and Ganguli 2017) impose a weightedquadratic penalty on the update of the parameters, where the

Page 3: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

weight encodes the relative importance of parameters to themodel’s performance on previous tasks. However, regular-ization approaches may suffer from the trade-off betweenthe model’s performance on the old and new tasks. The pro-posed framework is closely related to the works of varia-tional continual learning (Nguyen et al. 2018; Kessler et al.2019), which leverages Bayesian neural networks to retain adistribution over model parameters, and IBP priors to auto-matically determine the complexity of network. We developa task-aware architecture to infer the activated kernels foreach task automatically. The cost of storage does not signif-icantly grow as tasks accumulate. It also achieves good per-formance on both old and new tasks by balancing the reuseof learned kernels and the generation of new kernels.

Interactive Segmentation Interactive segmentation al-low user to make edits by clicks (Xu et al. 2016), scrib-bles (Grady et al. 2005), bounding boxes (Rajchl et al.2016; Castrejon et al. 2017), or extreme points (Maniniset al. 2018; Khan et al. 2019). Interactive segmentation mod-els need to be pre-trained using a hold-out labeled datasetto make initial predictions before user interactions. Afteruser annotations are collected, a majority of the models re-fine segmentation through spatial regularization using post-processing techniques such as conditional random fields andgraph cuts (Wang et al. 2018; Dhara et al. 2018; Zhou,Chen, and Wang 2019). Since the network parameters arenot updated, they do not extract knowledge from user anno-tations for segmenting other images. A few models propa-gate user annotations to unannotated pixels and retrain themodel (Lin et al. 2016), but they do not incorporate spe-cific mechanisms for knowledge transfer and thus may suf-fer from catastrophic forgetting.

Uncertainty quantifies the degree to which a machinelearning model is unconfident about its predictions and im-plies whether users can trust the results (Gal 2016). TheBayesian convolutional network is a popular technique ofuncertainty estimation for semantic segmentation (Kendall,Badrinarayanan, and Cipolla 2015; Jena and Awate 2019).The model may report high uncertainty on visually-difficultareas, possibly due to low contrast or brightness, and out-of-distribution areas due to the different distributions be-tween training and testing images. Interactive segmentationposes a new challenge as uncertainty estimation shall be fur-ther used to inform users for selective annotation. However,integrating uncertainty with deep learning-based interactivesegmentation is under-explored. One notable work leveragesuncertainty to actively query labels on an entire image (Yanget al. 2017), but strictly speaking, it is not an interactive seg-mentation model as it does not refine segmentation results.

Besides continual learning and interactive segmentation,the proposed framework is also related to active learning andonline learning. We provide additional discussions on theirkey differences in the Appendix.

The Framework of Continual Learning forInteractive Segmentation

In this section, we start by presenting the overall architectureof the proposed CLIS framework. We then describe task for-

Figure 2: Architecture of the CLIS Framework (new compo-nents are highlighted in blue)

mulation in continual learning and a unique training-testingprotocol to best support interactive image segmentation. Fi-nally, we present the details for continual knowledge learn-ing through Bayesian nonparametric modeling.

Overview of the Architecture. Segmentation modelssuch as fully convolutional network (Long, Shelhamer, andDarrell 2015) and U-Net (Ronneberger, Fischer, and Brox2015) usually apply an encoder-decoder architecture. Dif-ferent from conventional segmentation models, the proposedframework introduces several new components, uniquelydesigned for interactive segmentation and continual learn-ing, which are highlighted in Figure 2 and discussed below.

The segmentation module consists of an encoder-decoderarchitecture (Block A1-A4 and Block B1-B3). We pro-pose to use Bayesian nonparametric modeling to incorporateknowledge learned from previous tasks as probabilistic priorto regularize parameter updates in the current task. We intro-duce binary masks to control convolutional kernel activationof the encoders and decoders, aiming to activate the optimalset of kernels for each task. The task-aware module extractsinformation from images to infer the optimal masks with orwithout user annotations. The uncertainty estimation mod-ule is nested in Blocks A1-A4 and B1-B3 to generate an un-certainty map from an input image to guide user annotations.The interactive learning module is inspired by Markov Ran-dom Fields (MRF) to propagate user annotations to unanno-tated pixels by minimizing a customized energy function.

Task Formulation for Interactive SegmentationInteractive segmentation models usually require pre-trainingto make initial predictions given new images. We first use ahold-out dataset that consists of images and pixel-wise la-bels to pre-train the network. We then formulate interactiveimage segmentation as a continual learning problem witha sequence of T tasks as {T1, T2, ..., TT }. In task Tt, themodel is provided with S new imagesXt = {x1t , x2t , ..., xSt }and makes initial predictions, and the user provides selec-tive annotations At = {a1t , a2t , ..., aSt }, which are consid-

Page 4: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

ered ground truth and used to update the model’s parame-ters. This setting can be categorized as continual learningwith new instances (Lomonaco and Maltoni 2017) whereall classes are known in pre-training while each subsequenttask contains instances with a different subset of semanticclasses.

We evaluate an interactive segmentation framework’s per-formance from three aspects: basic performance, forwardknowledge transfer, and backward knowledge transfer. Thebasic performance measures how the model learns from userannotations to improve the current image’s segmentation.The model is trained at task Tt using the images and the cor-responding annotations and evaluated using the ground truthof unannotated pixels Ut = {u1t , u2t , ..., uSt }. The forwardknowledge transfer measures how the model leverages userannotations to improve the segmentation of other unseen im-ages. After training at task Tt, the model is evaluated usingunseen images from future tasks {Tt+1, ..., TT } in terms ofinitial segmentation. The backward knowledge transfer mea-sures how the model prevents deteriorated performance onpreviously-seen images. The model is trained at task Tt andevaluated using images from previous task {T1, ..., Tt−1} interms of initial prediction. In summary, at task Tt, the train-ing set St and evaluation set Qt are

St = {Xt, At},Basic :Qt = {Xt, Ut}Forward :Qt = {Xt+1,...,T , Ut+1,...,T }

Backward :Qt = {X1,...,t−1, U1,...,t−1}

Continual Knowledge Learning through BayesianNonparametric ModelingThe proposed framework organizes tasks in a sequential or-der to facilitate interactive segmentation for knowledge ac-quisition from users. Thus, it is critical to properly maintainimportant (latent) features and knowledge learned throughprior tasks for future unseen tasks. Furthermore, as the tasksmay be highly heterogeneous, not all the learned featuresare relevant to a given task. Instead, a specific subset of fea-tures should be selected to optimize the segmentation per-formance.

To address these challenges, we propose to allocate theneural network resources dynamically through Bayesiannonparametric modeling. The proposed model allows the la-tent feature space to continuously grow with new tasks whilebeing capable of extracting a subset of latent features mostrelevant to a given task through posterior inference. In par-ticular, for each convolutional layer l, we set the truncationthreshold Kl to a large number to accommodate a potentiallarge feature space. We further introduce a layer-wise binarymask to determine whether a kernel is activated or not givena specific image. Given the l-th layer’s input as hl−1, thek-th convolutional kernel W l

k and the binary mask zlk, theoutput feature map hl is

hl = {hlk}k, hlk = zlk(W lk ~ hl−1) (1)

where ~ denotes the convolution operator. If zlk = 0, thecorresponding channel hlk is 0, indicating that the kerneldoes not contribute to this task.

To achieve dynamic architecture evolution and task-specific feature extraction, we place an Indian Buffet process(IBP) (Ghahramani and Griffiths 2006) prior on the binarymask for each layer:

vlk ∼ Beta(a0, b0), πlk ∼

k∏j=1

vlj , zlk ∼ Bern(πl

k) (2)

where zlk is a Bernoulli variable indicating whether a kernelis activated or not. For each convolutional kernel W l

k, weplace an element-wise Gaussian prior

W lk = {wl

ki}i, wlki ∼ N(µ0, (σ0)2) (3)

where i is the index of elements within the kernel.The IBP prior allows potentially infinite latent features

to jointly explain the ever-increasing data instances (im-ages/tasks in our case). It is analogous to customers arriv-ing at an Indian Buffet restaurant with unlimited dishes. Inour interactive segmentation case, each layer is placed anIBP prior, where the tasks correspond to customers, and theconvolutional kernels correspond to the dishes. An impor-tant characteristic of the IBP is that a new customer is morelikely to take dishes that have been taken by a lot of previouscustomers, while it is still possible to take some new dishes.This characteristic is useful in our framework to automati-cally determine the optimal number of kernels and encour-age reusing existing kernels in a new task. If the new taskis similar to some previous tasks, the masks should be simi-lar to the corresponding previous masks, indicating that thekernel activations are mostly the same. If the task is differ-ent from any previous tasks, the masks shall have some newnon-zero entries, meaning that some new kernels are acti-vated for feature extraction and interpretation.

Variational Inference Since exact inference of the pos-terior of zlk and W l

k is intractable, we introduce a varia-tional distribution q to approximate those posteriors. Letφ = {v, z,W} denote all the corresponding parameters,and Dt denote the data that the network is fitting at task t.Assume q(φ) is factorized as

q(φ) =∏k,l

q(vlk)q(zlk|vlk)∏i

q(wlki) (4)

The key idea of variational continual learning is using thevariational posterior distribution of task Tt−1 as a prior dis-tribution for task Tt: qt−1(φ) = pt(φ). For the first taskT1,the prior is set to the parameters learned from pre-training.The total loss function for task t is given as the negative ev-idence lower bound:

Lt = KL[qt(φ)||qt−1(φ)]− Eqt(φ)[ln p(Dt|φ)] (5)

where the first term is the KL divergence that regularizesthe change of parameters between task Tt and task Tt−1,and the second term is the expectation of the log likelihood.Plugging in (4), the loss is expanded as

Lt = KL[qt(v)||qt−1(v)] +KL[qt(z|v)||qt−1(z|v)]

+KL[qt(W)||qt−1(W)]− Eqt(φ)[ln p(Dt|z,W)](6)

Page 5: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

Eq (6) is the general expression of the loss function. How-ever, each term may take different forms on different stagesof interactive segmentation, which are detailed in the follow-ing sections.

We use reparameterization tricks to make the variationalparameters learnable by the network through gradient de-scent. At task Tt, we plug in the parameters alk,t−1, b

lk,t−1

learned from the previous task as a prior and notice thatqt(v

lk) ∼ Beta(alk,t, b

lk,t). The first KL divergence term of

(6) is approximated as

KL[qt(v)||qt−1(v)] =∑k,l

lnB(alk,t−1, b

lk,t−1)

B(alk,t, blk,t)

+ (alk,t − alk,t−1)ψ(alk,t) + (blk,t − blk,t−1)ψ(blk,t)

+ (alk,t−1 + blk,t−1 − alk,t − blk,t)ψ(alk,t + blk,t)

(7)

where B and ψ denote beta and digamma functions, re-spectively. The discrete nature of Bernoulli variables makesbackpropagation infeasible. We relax the hard constraintof Bernoulli variables with continuous ones as a concreteBernoulli distribution (Maddison, Mnih, and Teh 2016),which ranges between [0, 1] and is peaked at 0 and 1. Givenπlk in the forward pass, zlk is now sampled from

zlk ={

1 + exp[−τ−1(lnπlk − ln(1− πl

k) + ε)]}−1 (8)

where τ is the temperature controlling the smoothness ofconcerte distribution, and ε ∼ Logistic(0, 1) is randomlysampled from a logistic distribution (Balakrishnan 1991).With the relaxation, the second term of (6) is approximatedas

KL[qt(z|v)||qt−1(z|v)] =∑k,l

πlk,t(lnπ

lk,t − lnπl

k,t−1)

+ (1− lnπlk,t)(ln(1− lnπl

k,t)− ln(1− lnπlk,t−1))

(9)At task Tt, we plug in the parameters µl

ki,t−1, σlki,t−1

learned from the previous task as Gaussian priors, and thethird KL divergence term is estimated similarly.

Before user interaction, the model is pre-trained using asmall set of data to learn the initial weights of the network,and the total loss is defined in (6), where the log-likelihoodterm is the expectation of cross-entropy loss of pixel-wiselabels. During pre-training, we use Monte-Carlo samplingto draw samples for each layers’ mask and kernels using (2)and (3), and estimate the output using (1). The total loss canbe back-propagated to optimize z and W.

Predicting Initial SegmentationGiven new images from task Tt, the proposed frameworkfirst generates initial predictions for user reference. At thistime, user annotations are not provided yet, and thus theground-truth labels are not available. It poses a question ofhow the optimal mask for this task can be learned throughbackpropagation. To address this problem, we propose toleverage the image data to infer the optimal mask. Givena new image, we use Monte Carlo (MC) sampling to draw

samples for each layers’ kernels using (3). The initial esti-mation of kernel activation is generated by feeding the in-put data via the task-aware module, which consists of multi-layer perceptron. The forward pass goes through the en-coders and decoders to generate an initial segmentation anduncertainty map.

Interactive SegmentationOnce the initial segmentation is generated, users may start toadjust the initial result through annotations. Our frameworkaccompanies the prediction with important uncertainty in-formation to direct users to the most informative parts of thesegmentation map and lessen the overall annotation effort.After the annotations are provided, they will be used to re-fine the model to generate an improved segmentation moreconsistent with users’ expectations.

Uncertainty Estimation The proposed framework gen-erates uncertainty maps using Monte-Carlo sampling withpixel-wise variance estimation (Kendall and Gal 2017). Thisuncertainty estimation method can be implemented withinthe proposed architecture of the segmentation module. Itdecomposes the total uncertainty into two distinct sources:epistemic and aleatoric. Epistemic uncertainty is estimatedby using MC sampling to generate multiple predictions andevaluate the variance, while the aleatoric uncertainty is eval-uated by adding a head to the decoder’s last layer beforesoftmax to estimate the variance. The total uncertainty isevaluated using the predictive entropy:

Um = −(θm)T ln θm (10)

where θm is predicted probability vector of pixel m aver-aged across MC samples.

Propagating User Annotations Given the uncertaintymap, users are guided to focus on segmentation areas withhigh uncertainty, which are more likely to be inaccurate. Itreduces users’ effort to check the entire image to locate in-correct segmentation areas. Then the user annotations arepropagated from annotated areas to unannotated regions togenerate an updated segmentation map θ∗. Intuitively, theupdated segmentation map should not deviate too muchfrom the original result while being consistent with user an-notations. Following the work with Markov random field forpropagating user annotations (Lin et al. 2016), we introducean energy function G, which consists of a unary term Guna,an annotation consistency term Gann, and a pairwise simi-larity term Gpar.

G =∑m

Guna(m)+∑m

Gpar(m)+∑m,n

Gann(m,n) (11)

where m and n are indices of pixels. The unary term en-courages θ∗ to be similar to the original prediction θ foreach pixel m. It is defined as the KL divergence betweentwo categorical distributions parameterized by θ∗

m and θm:

Guna(m) = (θ∗m)T (lnθ∗

m − ln θm) (12)

The pairwise term encourages visually similar pixels to havesimilar propagated probability vectors. It takes the form of

Page 6: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

weighted symmetric KL divergence:

Gpar(m,n) = exp(−‖rm − rn‖22)× [(θ∗m)T (lnθ∗

m

− lnθ∗n) + (θ∗

n)T (lnθ∗n − lnθ∗

m)]/2(13)

where (m,n) denotes a pair of pixels and rm denotes theLAB color vector of pixel m. The annotation consistencyterm enforces a labeled pixel has a probability vector almostthe same as the one-hot user annotation.

Gann(m) = −λm(am)T lnθ∗m (14)

where am is user annotation at pixel m, the weight λm isset to a sufficiently large value if pixel m is annotated andλm = 0 if otherwise.

Learning from User Annotations The propagated mapis used to evaluate the loss and update the network pa-rameters to adapt to user annotations. The total loss is de-fined in (6). Notice that z and W now denote the binarymasks and kernels from all the encoders and decoders, andDt = {Xt, Y

∗t } where Y ∗

t denotes the propagated maps,and the term Eqt(φ) equals the cross-entropy Lce:

Eqt(φ)[ln p(Y∗t |φ)] = −Eqt(φ)[

1

M

∑m

y∗m ln(ym)] (15)

where y∗m is the one-hot vector of pixel m’s propagated la-

bel, ym is the corresponding predicted probability vector.The total loss is backpropagated to refine the network

parameters. With the updated parameters of z and W, themodel performs another feed-forward to generate a refinedsegmentation map and the corresponding uncertainty esti-mation for user reference. This process may iterate multiplerounds until the user is satisfied.

ExperimentsIn this section, we report our experimental results on threereal-world image datasets to demonstrate the proposedframework’s effectiveness.

Datasets. We evaluate the performance of the proposedframework using three image datasets. The first is from thePascal VOC challenge (Everingham et al. 2015) that in-cludes 2913 images with 21 semantic classes. We randomlyselect 20% images for pre-training, because a segmentationmodel trained with sufficiently large data usually performswell on initial predictions and make user annotations lesscritical. The second dataset is Cityscapes (Cordts et al. 2016)containing street scenes from 50 different cities, with pixel-level annotations of 5000 frames and 8 major semantic cate-gories. Similarly, we randomly sample 20% images for pre-training. The third dataset includes 255 images collected in aprior study through data elicitation experiments. Each imagecorresponds to a multiclass segmentation map that partitionsthe image based on primary morphology, including macules,papules, plaques, nodules, vesicles, bullae, scale, and crusts.We use 80% of the data (204 images) for pre-training due tothis dataset’s small size and the rest for evaluation througha random split. Data augmentation is performed in a similarway as (Ronneberger, Fischer, and Brox 2015) using hori-zontal flipping, random shifting, and rotation.

Experimental setup. For all three datasets, we formulatetwenty task sequences, each of which contains eight con-secutive tasks. In each task, we randomly sample one imageand simulate user interactions with the proposed framework.Given the segmentation result, sixteen erroneous areas withthe highest uncertainty are selected and annotated, and themodel is updated to generate a refined segmentation. Weset the maximum iteration of user interactions to two, andrecord the basic performance. The performance on forwardand backward knowledge transfer is evaluated after tasks 4and 8. All reported results are averaged among 20 task se-quences. The hyper-parameters are set to α0 = 1, β0 = 1,µ0 = 0, σ0 = 1, τ = 2. The Adam optimizer is used forgradient-based optimization. For propagating user annota-tions, we use the Slic algorithm (Achanta et al. 2012) togroup visually-similar pixels into superpixels and reduce thecomputational cost. It should be noted that simulated userinteractions are suitable for quantitative evaluation becausethey provide a controlled environment and eliminate the fac-tors that could affect human actions. In the Appendix, weinclude examples of actual user interactions to illustrate theframework’s usage in practice1.

Comparison baselines and evaluation metrics. Wecompare with LwF (Li and Hoiem 2017), EWC (Kirkpatricket al. 2017) and PackNet (Mallya and Lazebnik 2018).LwF and EWC are representative regularization-based ap-proaches that introduce additional loss terms for protectingconsolidated knowledge. In contrast, the proposed frame-work applies variational inference to the task-specific prob-abilistic mask and uses KL divergence to effectively regu-larize kernels and masks. PackNet is an isolation-based ap-proach that implements network pruning to allocate differentneurons for different tasks, while the proposed frameworkrelies on a task-specific mask to control kernel activation.We use mean accuracy (Acc.) and mean intersect-over-union(IoU) as evaluation metrics.

Performance comparison. We first present qualitativecomparisons with the baselines in Figure 3. After trainingthe model at Task 4, we visualize the refined segmenta-tion map and visualize the initial performance at Task 6for forward knowledge transfer and Task 2 for backwardknowledge transfer. Compared with other baselines, the pro-posed framework usually generates better boundaries, i.e.,the boundaries of horse and the television (Dataset 1 Tasks6 and 2), buildings (Dataset 2 Task 2), bullae and crusts(Dataset 3 Tasks 4 and 6). Furthermore, the proposed frame-work performs better in identifying smaller objects, i.e.,papules (Dataset 3 Task 2). The superior performance maybe attributed to (1) uncertainty-guided user interactions thatprovide informative annotations and (2) task-specific maskthat activates suitable kernels for feature extraction and in-terpretation.

Quantitative comparisons are provided in Table 1. In mostcases, the proposed CLIS framework outperforms the com-peting baselines. LwF tends to strongly regularize the updateof parameters and thus makes it harder to refine segmenta-

1The appendix and the source code are available athttps://github.com/ritmininglab/CLIS

Page 7: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

Figure 3: Illustrative Examples of Segmentation Results From Dataset 1 (Left) and 2 (Middle) and 3 (Right)

Table 1: Quantitative Comparison of Basic Performance,Forward and Backward Knowledge Transfer

Dataset 1 Dataset 2 Dataset 3Acc. IoU Acc. IoU Acc. IoU

Basic PerformanceCLIS 86.7 77.8 91.5 85.0 82.8 72.1LwF 84.8 73.9 86.9 78.5 80.7 67.9EWC 85.5 75.0 89.5 82.4 81.2 69.5

PackNet 85.9 76.1 90.4 83.2 81.9 70.4Forward Knowledge Transfer

CLIS 79.2 65.5 84.5 75.0 77.8 63.4LwF 77.5 64.0 83.5 72.4 76.1 61.8EWC 78.1 63.9 83.7 73.1 76.5 61.7

PackNet 78.0 64.3 82.9 71.4 75.7 60.9Backward Knowledge Transfer

CLIS 81.7 69.3 86.6 77.9 78.9 65.3LwF 80.4 68.1 85.5 75.7 77.5 63.9EWC 80.2 67.5 85.7 76.2 76.9 63.0

PackNet 79.7 67.0 84.9 74.0 76.4 62.1

tion based on user annotations. PackNet incorporates a flexi-ble architecture and performs well on refining segmentation,but lacks a principled way of identifying task-specific ker-nel activation, and thus the performance may deteriorate onprevious tasks.

Interpretability. Using a layer-wise binary mask to con-trol kernel activation also improves the framework’s inter-pretability because we can track the activated kernels at eachtask. In Figure 4, we visualize the kernel activation as a gray-scale matrix for two tasks. The highlighted fifth kernel of thetenth layer is activated in the first task (the correspondingparameter π is close to 1) but not in the second task. We vi-sualize the corresponding output feature maps of this kernelin both tasks. It shows that the kernel captures features re-

Figure 4: Illustrative Example of Kernel Activation for Task1 (Top) and 2 (Bottom): Due to space limit, the kernel acti-vation map only visualize first 20 kernels of each layer. Darkgrids indicate the corresponding kernel is activated.

lated to plaques, which is important disease morphology forthe image in the first task.

Ablation study. We conduct an ablation study to furtherevaluate some key components of the framework, includinguncertainty guided interaction and task-aware kernel activa-tion for initial segmentation. Please refer to the Appendixfor details.

ConclusionIn this work, we formulate interactive segmentation as a con-tinual learning problem and propose a framework to effec-tively learn from user annotations to improve the segmen-tation performance on current and future tasks and preventcatastrophic forgetting on previous tasks. Uncertainty infor-mation is leveraged to provide users with informative guid-ance to improve the segmentation with minimum annotationeffort. A future direction is to extend the allowed interac-tion from clicks to other types such as scribbles and boxesto make the interaction process more efficient.

Page 8: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

AcknowledgementThis research was partially supported by NSF IIS award IIS-1814450 and ONR award N00014-18-1-2875. Any opin-ions, findings, and conclusions or recommendations ex-pressed in this paper are those of the authors and do not nec-essarily reflect the official views of any funding agency. Wewould like to thank the anonymous reviewers for reviewingthe manuscript.

Ethical ImpactThis work will provide both theoretical underpinning andempirical evaluation on the design of human-in-the-looplearning in interactive segmentation, aiming to infuse hu-man expertise into the training-evaluation-feedback loopand make the model more accurate, efficient, and cus-tomized to human needs. To the best of our understanding,we do not identify the negative societal implications of theproposed framework.

ReferencesAchanta, R.; Shaji, A.; Smith, K.; Lucchi, A.; Fua, P.; andSusstrunk, S. 2012. SLIC superpixels compared to state-of-the-art superpixel methods. IEEE transactions on patternanalysis and machine intelligence 34(11): 2274–2282.

Balakrishnan, N. 1991. Handbook of the logistic distribu-tion. CRC Press.

Castrejon, L.; Kundu, K.; Urtasun, R.; and Fidler, S. 2017.Annotating object instances with a polygon-rnn. In Proceed-ings of the IEEE conference on computer vision and patternrecognition, 5230–5238.

Chen, L.-C.; Papandreou, G.; Kokkinos, I.; Murphy, K.; andYuille, A. L. 2017. Deeplab: Semantic image segmentationwith deep convolutional nets, atrous convolution, and fullyconnected crfs. IEEE transactions on pattern analysis andmachine intelligence 40(4): 834–848.

Cordts, M.; Omran, M.; Ramos, S.; Rehfeld, T.; Enzweiler,M.; Benenson, R.; Franke, U.; Roth, S.; and Schiele, B.2016. The Cityscapes Dataset for Semantic Urban SceneUnderstanding. In Proc. of the IEEE Conference on Com-puter Vision and Pattern Recognition (CVPR).

Dhara, A. K.; Ayyalasomayajula, K. R.; Arvids, E.;Fahlstrom, M.; Wikstrom, J.; Larsson, E.-M.; and Strand,R. 2018. Segmentation of post-operative glioblastoma inmri by u-net with patient-specific interactive refinement.In International MICCAI Brainlesion Workshop, 115–122.Springer.

Everingham, M.; Eslami, S. M. A.; Van Gool, L.; Williams,C. K. I.; Winn, J.; and Zisserman, A. 2015. The Pascal Vi-sual Object Classes Challenge: A Retrospective. Interna-tional Journal of Computer Vision 111(1): 98–136.

French, R. M. 1999. Catastrophic forgetting in connectionistnetworks. Trends in cognitive sciences 3(4): 128–135.

Gal, Y. 2016. Uncertainty in deep learning. University ofCambridge 1: 3.

Ghahramani, Z.; and Griffiths, T. L. 2006. Infinite latentfeature models and the Indian buffet process. In Advancesin neural information processing systems, 475–482.Grady, L.; Schiwietz, T.; Aharon, S.; and Westermann, R.2005. Random walks for interactive organ segmentation intwo and three dimensions: Implementation and validation.In International Conference on Medical Image Computingand Computer-Assisted Intervention, 773–780. Springer.Jena, R.; and Awate, S. P. 2019. A bayesian neural net tosegment images with uncertainty estimates and good cali-bration. In International Conference on Information Pro-cessing in Medical Imaging, 3–15. Springer.Kendall, A.; Badrinarayanan, V.; and Cipolla, R. 2015.Bayesian segnet: Model uncertainty in deep convolu-tional encoder-decoder architectures for scene understand-ing. arXiv preprint arXiv:1511.02680 .Kendall, A.; and Gal, Y. 2017. What uncertainties do weneed in bayesian deep learning for computer vision? InAdvances in neural information processing systems, 5574–5584.Kessler, S.; Nguyen, V.; Zohren, S.; and Roberts, S. 2019.Indian Buffet Neural Networks for Continual Learning.arXiv preprint arXiv:1912.02290 .Khan, S.; Shahin, A. H.; Villafruela, J.; Shen, J.; and Shao,L. 2019. Extreme Points Derived Confidence Map as a Cuefor Class-Agnostic Interactive Segmentation Using DeepNeural Network. In International Conference on MedicalImage Computing and Computer-Assisted Intervention, 66–73. Springer.Kirkpatrick, J.; Pascanu, R.; Rabinowitz, N.; Veness, J.; Des-jardins, G.; Rusu, A. A.; Milan, K.; Quan, J.; Ramalho, T.;Grabska-Barwinska, A.; et al. 2017. Overcoming catas-trophic forgetting in neural networks. Proceedings of thenational academy of sciences 114(13): 3521–3526.Li, Z.; and Hoiem, D. 2017. Learning without forgetting.IEEE transactions on pattern analysis and machine intelli-gence 40(12): 2935–2947.Lin, D.; Dai, J.; Jia, J.; He, K.; and Sun, J. 2016. Scribble-sup: Scribble-supervised convolutional networks for seman-tic segmentation. In Proceedings of the IEEE Conference onComputer Vision and Pattern Recognition, 3159–3167.Lomonaco, V.; and Maltoni, D. 2017. Core50: a new datasetand benchmark for continuous object recognition. arXivpreprint arXiv:1705.03550 .Long, J.; Shelhamer, E.; and Darrell, T. 2015. Fully convo-lutional networks for semantic segmentation. In Proceed-ings of the IEEE conference on computer vision and patternrecognition, 3431–3440.Maddison, C. J.; Mnih, A.; and Teh, Y. W. 2016. The con-crete distribution: A continuous relaxation of discrete ran-dom variables. arXiv preprint arXiv:1611.00712 .Mallya, A.; and Lazebnik, S. 2018. Packnet: Adding mul-tiple tasks to a single network by iterative pruning. In Pro-ceedings of the IEEE Conference on Computer Vision andPattern Recognition, 7765–7773.

Page 9: A Continual Learning Framework for Uncertainty-Aware … · 2021. 1. 6. · Badrinarayanan, and Cipolla 2015; Jena and Awate 2019). The model may report high uncertainty on visually-difficult

Maninis, K.-K.; Caelles, S.; Pont-Tuset, J.; and Van Gool,L. 2018. Deep extreme cut: From extreme points to objectsegmentation. In Proceedings of the IEEE Conference onComputer Vision and Pattern Recognition, 616–625.Milletari, F.; Navab, N.; and Ahmadi, S.-A. 2016. V-net:Fully convolutional neural networks for volumetric medicalimage segmentation. In 2016 Fourth International Confer-ence on 3D Vision (3DV), 565–571. IEEE.Nguyen, C. V.; Li, Y.; Bui, T. D.; and Turner, R. E. 2018.Variational Continual Learning. In International Conferenceon Learning Representations.Rajchl, M.; Lee, M. C.; Oktay, O.; Kamnitsas, K.; Passerat-Palmbach, J.; Bai, W.; Damodaram, M.; Rutherford, M. A.;Hajnal, J. V.; Kainz, B.; et al. 2016. Deepcut: Object seg-mentation from bounding box annotations using convolu-tional neural networks. IEEE transactions on medical imag-ing 36(2): 674–683.Ronneberger, O.; Fischer, P.; and Brox, T. 2015. U-net: Con-volutional networks for biomedical image segmentation. InInternational Conference on Medical image computing andcomputer-assisted intervention, 234–241. Springer.Rusu, A. A.; Vecerık, M.; Rothorl, T.; Heess, N.; Pascanu,R.; and Hadsell, R. 2017. Sim-to-real robot learning frompixels with progressive nets. In Conference on Robot Learn-ing, 262–270.Thrun, S.; and Mitchell, T. M. 1995. Lifelong robot learning.Robotics and autonomous systems 15(1-2): 25–46.Wang, G.; Zuluaga, M. A.; Li, W.; Pratt, R.; Patel, P. A.;Aertsen, M.; Doel, T.; David, A. L.; Deprest, J.; Ourselin, S.;et al. 2018. DeepIGeoS: a deep interactive geodesic frame-work for medical image segmentation. IEEE transactionson pattern analysis and machine intelligence 41(7): 1559–1572.Xu, N.; Price, B.; Cohen, S.; Yang, J.; and Huang, T. S. 2016.Deep interactive object selection. In Proceedings of theIEEE Conference on Computer Vision and Pattern Recog-nition, 373–381.Yang, L.; Zhang, Y.; Chen, J.; Zhang, S.; and Chen, D. Z.2017. Suggestive annotation: A deep active learning frame-work for biomedical image segmentation. In Internationalconference on medical image computing and computer-assisted intervention, 399–407. Springer.Zenke, F.; Poole, B.; and Ganguli, S. 2017. Continual learn-ing through synaptic intelligence. Proceedings of machinelearning research 70: 3987.Zhou, B.; Chen, L.; and Wang, Z. 2019. Interactive DeepEditing Framework for Medical Image Segmentation. In In-ternational Conference on Medical Image Computing andComputer-Assisted Intervention, 329–337. Springer.