Original: AdaWAC : A DAPTIVELY W EIGHTED A UGMENTATION C ONSISTENCY R EGULARIZATION FOR V OLUMETRIC M EDICAL I MAGE S EGMENTATION
Revised: AdaWAC : A DAPTIVELY W EIGHTED A UGMENTATION C ONSISTENCY R EGULARIZATION FOR V OLUMETRIC M EDICAL I MAGE S EGMENTATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Sample reweighting is an effective strategy for learning from training data coming from a mixture of different subpopulations.
Revised: Sample reweighting is an effective strategy for learning from training data coming from a mixture of different subpopulations.

Original: However, existing reweighting algorithms do not fully take advantage of the particular type of data distribution encountered in volumetric medical image segmentation, where the training data images are uniformly distributed but their associated data labels fall into two subpopulations—“label-sparse” and “label-dense”—depending on whether the data image occurs near the beginning/end of the volumetric scan or the middle.
Revised: However, existing reweighting algorithms do not fully take advantage of the particular type of data distribution encountered in volumetric medical image segmentation, where the training data images are uniformly distributed but their associated data labels fall into two subpopulations—“label-sparse” and “label-dense”—depending on whether the data image occurs near the beginning/end of the volumetric scan or the middle.

Original: For this setting, we propose AdaWAC as an adaptive weighting algorithm that assigns label-dense samples to supervised cross-entropy loss and label-sparse samples to unsupervised consistency regularization.
Revised: For this setting, we propose AdaWAC as an adaptive weighting algorithm that assigns label-dense samples to supervised cross-entropy loss and label-sparse samples to unsupervised consistency regularization.

Original: We provide a convergence guarantee for AdaWAC by appealing to the theory of online mirror descent on saddle point problems.
Revised: We provide a convergence guarantee for AdaWAC by appealing to the theory of online mirror descent on saddle point problems.

Original: Moreover, we empirically demonstrate that AdaWAC not only enhances segmentation performance and sample efficiency but also improves robustness to the subpopulation shift in labels.
Revised: Moreover, we empirically demonstrate that AdaWAC not only enhances segmentation performance and sample efficiency but also improves robustness to the subpopulation shift in labels.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Modern machine learning has been revolutionizing the field of medical imaging, especially in computer-aided diagnosis with Computed Tomography (CT) and Magnetic Resonance Imaging (MRI) scans.
Revised: Modern machine learning has been revolutionizing the field of medical imaging, especially in computer-aided diagnosis with Computed Tomography (CT) and Magnetic Resonance Imaging (MRI) scans.

Original: While the successes of most classical learning algorithms ( e.g. , empirical risk minimization (ERM)) build upon the assumption that training samples are independently and identically distributed ( i.i.d. ), real-world volumetric medical images rarely fit into this picture.
Revised: While the successes of most classical learning algorithms ( e.g. , empirical risk minimization (ERM)) build upon the assumption that training samples are independently and identically distributed ( i.i.d. ), real-world volumetric medical images rarely fit into this picture.

Original: Specifically for medical image segmentation, as instantiated in Figure 1, the segmentation labels corresponding to different cross-sections of organs within a given volume tend to have distinct distributions.
Revised: Specifically for medical image segmentation, as instantiated in Figure 1, the segmentation labels corresponding to different cross-sections of organs within a given volume tend to have distinct distributions.

Original: That is, the slices toward the beginning/end of the volume that contain no target organs have very few, if any, segmentation labels (which we refer to as “label-sparse”); whereas segmentation labels are prolific in the slices toward the middle of the volume (“label-dense”).
Revised: That is, the slices toward the beginning/end of the volume that contain no target organs have very few, if any, segmentation labels (which we refer to as “label-sparse”); whereas segmentation labels are prolific in the slices toward the middle of the volume (“label-dense”).

Original: Such discrepancy in labels results in distinct difficulty levels measured by the training cross-entropy (Wang et al., 2021b) and leads to various training schedulers (Hacohen & Weinshall, 2019; Tang et al., 2018; Tullis & Benjamin, 2011).
Revised: Such discrepancy in labels results in distinct difficulty levels measured by the training cross-entropy (Wang et al., 2021b) and leads to various training schedulers (Tullis & Benjamin, 2011; Tang et al., 2018; Hacohen & Weinshall, 2019).

Original: Motivated by the separation between label-sparse and label-dense samples, we explore the following questions in this work:
Revised: Motivated by the separation between label-sparse and label-dense samples, we explore the following questions in this work:

Original: What is the effect of separation between sparse and dense labels on segmentation?
Revised: What is the effect of separation between sparse and dense labels on segmentation?

Original: Can we leverage such separation to improve the segmentation accuracy?
Revised: Can we leverage such separation to improve the segmentation accuracy?

Original: We first formulate the mixture of label-sparse and label-dense samples as a subpopulation shift in the conditional distribution of labels given images P ( y | x ) .
Revised: We first formulate the mixture of label-sparse and label-dense samples as a subpopulation shift in the conditional distribution of labels given images P ( y | x ) .

Original: As illustrated in Figure 1, such subpopulation shift induces a separation in supervised cross-entropy between sparse and dense labels despite the uniformity of data images.
Revised: As illustrated in Figure 1, such subpopulation shift induces a separation in supervised cross-entropy between sparse and dense labels despite the uniformity of data images.

Original: Prior works address the subpopulation shift issue by utilizing hard thresholding algorithms such as Trimmed Loss Estimator (Shen & Sanghavi, 2019), MKL-SGD (Shah et al., 2020), Ordered
Revised: Prior works address the subpopulation shift issue by utilizing hard thresholding algorithms such as Trimmed Loss Estimator (Shen & Sanghavi, 2019), MKL-SGD (Shah et al., 2020), Ordered

Original: SGD (Kawaguchi & Lu, 2020), and quantile-based Kacmarz algorithm
Revised: SGD (Kawaguchi & Lu, 2020), and quantile-based Kacmarz algorithm

Original: (Haddock et al., 2020).
Revised: (Haddock et al., 2020).

Original: However, these trimmed-loss-based methods discard the samples from some subpopulations ( e.g. samples with label corruption estimated by their losses) at each iteration, which results in loss of information in the discarded data, leading to reduced sample efficiency.
Revised: However, these trimmed-loss-based methods discard the samples from some subpopulations ( e.g. samples with label corruption estimated by their losses) at each iteration, which results in loss of information in the discarded data, leading to reduced sample efficiency.

Original: Relaxing the hard thresholding operator to soft thresholding is proposed to incorporate the information from both subpopulations (Sagawa et al., 2020; Wang et al., 2018).
Revised: Relaxing the hard thresholding operator to soft thresholding is proposed to incorporate the information from both subpopulations (Wang et al., 2018; Sagawa et al., 2020).

Original: However, lowering the weights assigned to some subpopulations of data according to the properties of their labels reduces the importance of the data and labels simultaneously, suggesting that we may further improve the learning efficiency by exploiting the uniformity of data and the separation of labels separately.
Revised: However, lowering the weights assigned to some subpopulations of data according to the properties of their labels reduces the importance of the data and labels simultaneously, suggesting that we may further improve the learning efficiency by exploiting the uniformity of data and the separation of labels separately.

Original: Instead of thresholding out or down-weighting the label-sparse samples, we exploit the image inputs of these samples via augmentation consistency regularization .
Revised: Instead of thresholding out or down-weighting the label-sparse samples, we exploit the image inputs of these samples via augmentation consistency regularization .

Original: Consistency regularization (Bachman et al., 2014; Laine & Aila, 2016; Sohn et al., 2020) aims to learn the proximity between data augmentations of the same samples; crucially, this set-up does not involve the data labels, and hence consistency regularization has become an essential strategy for utilizing unlabeled data.
Revised: Consistency regularization (Bachman et al., 2014; Laine & Aila, 2016; Sohn et al., 2020) aims to learn the proximity between data augmentations of the same samples; crucially, this set-up does not involve the data labels, and hence consistency regularization has become an essential strategy for utilizing unlabeled data.

Original: For medical imaging tasks, consistency regularization has been extensively studied in the semi-supervised learning setting (Basak et al., 2022;
Revised: For medical imaging tasks, consistency regularization has been extensively studied in the semi-supervised learning setting (Bortsova et al., 2019;

Original: Bortsova et al., 2019; Li et al., 2020; Wang et al., 2021a; Zhang et al., 2021;
Revised: Zhao et al., 2019; Li et al., 2020; Wang et al., 2021a; Zhang et al., 2021;

Original: Zhao et al., 2019; Zhou et al., 2021).
Revised: Zhou et al., 2021; Basak et al., 2022).

Original: By contrast, we will explore its potency in the fully supervised setting—leveraging the spare information in all image inputs, regardless of their label subpopulations.
Revised: By contrast, we will explore its potency in the fully supervised setting—leveraging the spare information in all image inputs, regardless of their label subpopulations.

Original: Moreover, in light of the uniformity of unsupervised consistency on different slices throughout each volume, the augmentation consistency of the encoder layer outputs serves as a natural reference for separating samples from different subpopulations.
Revised: Moreover, in light of the uniformity of unsupervised consistency on different slices throughout each volume, the augmentation consistency of the encoder layer outputs serves as a natural reference for separating samples from different subpopulations.

Original: Whereby, we introduce the weighted augmentation consistency (WAC) regularization —a minimax formulation that not only incorporates the consistency regularization but also leverages the consistency regularization as a reference for reweighting the cross-entropy and the augmentation consistency terms corresponding to different samples.
Revised: Whereby, we introduce the weighted augmentation consistency (WAC) regularization —a minimax formulation that not only incorporates the consistency regularization but also leverages the consistency regularization as a reference for reweighting the cross-entropy and the augmentation consistency terms corresponding to different samples.

Original: At the saddle point, the WAC regularization automatically separates samples from different label subpopulations by assigning all weight to the consistency terms for label-sparse samples, and all weight to the cross-entropy terms for label-dense samples.
Revised: At the saddle point, the WAC regularization automatically separates samples from different label subpopulations by assigning all weight to the consistency terms for label-sparse samples, and all weight to the cross-entropy terms for label-dense samples.

Original: Furthermore, as an algorithm for solving the minimax problem posed by the WAC regularization, we propose AdaWAC —an adaptive weighting scheme inspired by a mirror-descent-based algorithm for distributionally-robust optimization (Sagawa et al., 2020).
Revised: Furthermore, as an algorithm for solving the minimax problem posed by the WAC regularization, we propose AdaWAC —an adaptive weighting scheme inspired by a mirror-descent-based algorithm for distributionally-robust optimization (Sagawa et al., 2020).

Original: By adaptively adjusting the weights between the cross-entropy and consistency terms of different samples, AdaWAC comes with both a convergence guarantee and empirical success.
Revised: By adaptively adjusting the weights between the cross-entropy and consistency terms of different samples, AdaWAC comes with both a convergence guarantee and empirical success.

Original: Overall, we summarize the main contributions of this work as follows:
Revised: Overall, we summarize the main contributions of this work as follows:

Original: • We cast the discrepancy between the sparse and dense labels within each volume as a subpopulation shift in the conditional distribution P ( y | x ) (Section 2).
Revised: • We cast the discrepancy between the sparse and dense labels within each volume as a subpopulation shift in the conditional distribution P ( y | x ) (Section 2).

Original: • We propose WAC regularization which uses the consistency of encoder layer outputs (in a UNet architecture) as a natural reference to incentivize separation between samples with sparse and dense labels (Section 3), along with an adaptive weighting algorithm— AdaWAC —for solving the WAC regularization problem with a convergence guarantee (Section 4).
Revised: • We propose WAC regularization which uses the consistency of encoder layer outputs (in a UNet architecture) as a natural reference to incentivize separation between samples with sparse and dense labels (Section 3), along with an adaptive weighting algorithm— AdaWAC —for solving the WAC regularization problem with a convergence guarantee (Section 4).

Original: • We empirically demonstrate the potency of AdaWAC not only in enhancing segmentation performance and sample efficiency but also in improving distributional robustness (Section 5).
Revised: • We empirically demonstrate the potency of AdaWAC not only in enhancing segmentation performance and sample efficiency but also in improving distributional robustness (Section 5).

Original: 1. R ELATED W ORK
Revised: 1. R ELATED W ORK

Original: Sample reweighting.
Revised: Sample reweighting.

Original: Sample reweighting is a popular strategy for coping with subpopulation shifts in training data where different weights are assigned to samples from different subpopulations.
Revised: Sample reweighting is a popular strategy for coping with subpopulation shifts in training data where different weights are assigned to samples from different subpopulations.

Original: In particular, the distributionally-robust optimization (DRO) (Ben-Tal et al., 2013; Duchi& Namkoong, 2018; Duchi et al., 2016;  Sagawa et al., 2020) considers a collection of training sample groups from different distributions where, with the explicit grouping of samples, the goal is to minimize the worst-case loss over the groups.
Revised: In particular, the distributionally-robust optimization (DRO) (Ben-Tal et al., 2013; Duchi et al., 2016; Duchi & Namkoong, 2018; Sagawa et al., 2020) considers a collection of training sample groups from different distributions where, with the explicit grouping of samples, the goal is to minimize the worst-case loss over the groups.

Original: Without prior knowledge on sample grouping, importance sampling (Alain et al., 2015; Gopal, 2016; Katharopoulos & Fleuret, 2018; Loshchilov & Hutter, 2015; Needell et al., 2014; Zhao & Zhang, 2015), iterative trimming (Kawaguchi & Lu, 2020; Shen & Sanghavi, 2019), and empirical-loss-based reweighting (Wu et al., 2022) are commonly incorporated in the stochastic optimization process for adaptive reweighting and separation of samples from different subpopulations.
Revised: Without prior knowledge on sample grouping, importance sampling (Needell et al., 2014; Zhao & Zhang, 2015; Alain et al., 2015; Loshchilov & Hutter, 2015; Gopal, 2016; Katharopoulos & Fleuret, 2018), iterative trimming (Kawaguchi & Lu, 2020; Shen & Sanghavi, 2019), and empirical-loss-based reweighting (Wu et al., 2022) are commonly incorporated in the stochastic optimization process for adaptive reweighting and separation of samples from different subpopulations.

Original: Consistency regularization.
Revised: regularization.

Original: Consistency regularization (Bachman et al., 2014;Berthelot et al., 2019; Laine & Aila, 2016; Sohn et al., 2020) is a popular way to exploit data augmentations that encourage the model to learn the vicinity among augmentations of the same sample, with the assumption that data augmentations generally preserve the semantic information in data.
Revised: Consistency regularization (Bachman et al., 2014; Laine & Aila, 2016; Sohn et al., 2020; Berthelot et al., 2019) is a popular way to exploit data augmentations that encourage the model to learn the vicinity among augmentations of the same sample, with the assumption that data augmentations generally preserve the semantic information in data.

Original: For medical imaging, consistency regularization is generally leveraged as a semi-supervised learning tool (Basak et al., 2022; Bortsova et al., 2019;
Revised: For medical imaging, consistency regularization is generally leveraged as a semi-supervised learning tool (Bortsova et al., 2019;

Original:  Li et al., 2020; Wang et al., 2021a; Zhang et al., 2021;Zhao et al., 2019;
Revised: Zhao et al., 2019; Li et al., 2020; Wang et al., 2021a; Zhang et al., 2021;

Original: Zhou et al., 2021).
Revised: Zhou et al., 2021; Basak et al., 2022).

Original: In efforts to incorporate consistency regularization in medical image segmentation with augmentation-sensitive labels, Li et al.
Revised: In efforts to incorporate consistency regularization in medical image segmentation with augmentation-sensitive labels, Li et al.

Original: (2020) encourages transformation consistency between predictions with augmentations applied to the image inputs and the segmentation outputs.
Revised: (2020) encourages transformation consistency between predictions with augmentations applied to the image inputs and the segmentation outputs.

Original: Basak et al.
Revised: Basak et al.

Original: (2022) penalizes inconsistent segmentation outputs between teacherstudent models, with MixUp (Zhang et al., 2017) applied on image inputs of the teacher model and segmentation outputs of the student model.
Revised: (2022) penalizes inconsistent segmentation outputs between teacherstudent models, with MixUp (Zhang et al., 2017) applied on image inputs of the teacher model and segmentation outputs of the student model.

Original: Instead of enforcing consistency in the segmentation output space as above, our algorithm leverages the insensitivity of sparse labels to augmentations and encourages consistent encodings (in the latent space of encoder outputs) on label-sparse samples.
Revised: Instead of enforcing consistency in the segmentation output space as above, our algorithm leverages the insensitivity of sparse labels to augmentations and encourages consistent encodings (in the latent space of encoder outputs) on label-sparse samples.

Original:  P ROBLEM S ETUP
Revised:  P ROBLEM S ETUP

Original: Notations.
Revised: Notations.

Original: We denote [ K ] = { 1 , . . .
Revised: We denote [ K ] = { 1 , . . .

Original: , K } for any K ∈ N .
Revised: , K } for any K ∈ N .

Original: For an arbitrary tensor, we adapt the syntax for Python slicing on the subscript (except counting from 1 ) to represent its elements and subtensors.
Revised: For an arbitrary tensor, we adapt the syntax for Python slicing on the subscript (except counting from 1 ) to represent its elements and subtensors.

Original: For example, x [ i,j ] denotes the ( i, j ) -entry of the two-dimensional tensor x , and x [ i, :] denotes its i -th row.
Revised: For example, x [ i,j ] denotes the ( i, j ) -entry of the two-dimensional tensor x , and x [ i, :] denotes its i -th row.

Original: Let I be a function onto { 0 , 1 } such that, for any event e , I { e } = 1 if e is true and 0 otherwise.
Revised: Let I be a function onto { 0 , 1 } such that, for any event e , I { e } = 1 if e is true and 0 otherwise.

Original: For any distribution P and n ∈ N , let P n denote the joint distribution of n samples drawn i.i.d.
Revised: For any distribution P and n ∈ N , let P n denote the joint distribution of n samples drawn i.i.d.

Original: from P .
Revised: from P .

Original: We refer to an event as happening with high probability ( w.h.p. )
Revised: We refer to an event as happening with high probability ( w.h.p. )

Original: if it takes place with probability 1 − Ω (poly ( n )) − 1 .
Revised: if it takes place with probability 1 − Ω (poly ( n )) − 1 .

Original: 2. P IXEL - WISE C LASSIFICATION WITH S PARSE AND D ENSE L ABELS
Revised: 2. P IXEL - WISE C LASSIFICATION WITH S PARSE AND D ENSE L ABELS

Original: We consider the volumetric medical image segmentation as a pixel-wise multi-class classification problem where we aim to learn a pixel-wise classifier h :
Revised: We consider the volumetric medical image segmentation as a pixel-wise multi-class classification problem where we aim to learn a pixel-wise classifier h :

Original: X → [ K ] d that serves as a good approximation for the ground truth h ∗ : X → [ K ] d .
Revised: X → [ K ] d that serves as a good approximation for the ground truth h ∗ : X → [ K ] d .

Original: Recall the separation of cross-entropy losses between samples with different fractions of nonbackground labels during training from Figure 1.
Revised: Recall the separation of cross-entropy losses between samples with different fractions of nonbackground labels during training from Figure 1.

Original: We refer to a sample ( x , y ) ∈
Revised: We refer to a sample ( x , y ) ∈

Original: X ×
Revised: X ×

Original: [ K ] d as label-sparse if most pixels in y are labeled as the background such that the cross-entropy loss on ( x , y ) converges rapidly in the early stage of training 1 .
Revised: [ K ] d as label-sparse if most pixels in y are labeled as the background such that the cross-entropy loss on ( x , y ) converges rapidly in the early stage of training 1 .

Original: Otherwise, we say that ( x , y ) is label-dense .
Revised: Otherwise, we say that ( x , y ) is label-dense .

Original: Formally, we describe such variation as a subpopulation shift in the label distribution.
Revised: Formally, we describe such variation as a subpopulation shift in the label distribution.

Original: Definition 1 (Mixture of label-sparse and label-dense distributions) .
Revised: Definition 1 (Mixture of label-sparse and label-dense distributions) .

Original: Let P 0 and P 1 be the distributions of label-sparse and label-dense samples with distinct conditional distributions P 

Revised: Let P 0 and P 1 be the distributions of label-sparse and label-dense samples with distinct conditional distributions P

Original: ( y | x ) and P 1 ( y | x ) , respectively, but the same marginal distribution P ( x ) such that P i ( x , y ) =
Revised: ( y | x ) and P 1 ( y | x ) , respectively, but the same marginal distribution P ( x ) such that P i ( x , y ) =

Original: P i ( y | x ) P ( x ) ( i = 0 , 1 ).
Revised: P i ( y | x ) P ( x ) ( i = 0 , 1 ).

Original: For ξ ∈ [0 , 1] , we define a data distribution P ξ where each sample ( x , y ) ∼ P ξ is drawn either from P 1 with probability ξ or from P 0 with probability 1 − ξ .
Revised: For ξ ∈ [0 , 1] , we define a data distribution P ξ where each sample ( x , y ) ∼ P ξ is drawn either from P 1 with probability ξ or from P 0 with probability 1 − ξ .

Original: We aim to learn a pixel-wise classifier from a function class H ∋ h θ = argmax k ∈
Revised: We aim to learn a pixel-wise classifier from a function class H ∋ h θ = argmax k ∈

Original: [ K ] f θ ( x ) [ j, :] for all j ∈ [ d ] where the underlying function f θ ∈ F , parameterized by some θ ∈ F θ , admits an encoder-decoder structure:
Revised: [ K ] f θ ( x ) [ j, :] for all j ∈ [ d ] where the underlying function f θ ∈ F , parameterized by some θ ∈ F θ , admits an encoder-decoder structure:

Original: [Equation]
Revised: [Equation]

Original: Here ϕ θ , ψ θ correspond to the encoder and decoder functions, respectively; ( F θ , ⟨· , ·⟩ F ) denotes the inner product space of parameters with the induced norm ∥·∥ F and dual norm ∥·∥ F , ∗ ; ( Z , ϱ ) is a latent metric space.
Revised: Here ϕ θ , ψ θ correspond to the encoder and decoder functions, respectively; ( F θ , ⟨· , ·⟩ F ) denotes the inner product space of parameters with the induced norm ∥·∥ F and dual norm ∥·∥ F , ∗ ; ( Z , ϱ ) is a latent metric space.

Original: To learn from segmentation labels, we consider the averaged cross-entropy loss :
Revised: To learn from segmentation labels, we consider the averaged cross-entropy loss :

Original: [Equation]
Revised: [Equation]

Original: We assume proper learning where there exists θ ∗ ∈ (cid:84) ξ ∈
Revised: We assume proper learning where there exists θ ∗ ∈ (cid:84) ξ ∈

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ( θ ; ( x , y ))] that is invariant with respect to ξ .
Revised: x [Equation] y ))] that is invariant with respect to ξ .

Original: 2. A UGMENTATION C ONSISTENCY R EGULARIZATION
Revised: 2. A UGMENTATION C ONSISTENCY R EGULARIZATION

Original: Despite the invariance of f θ ∗ to P ξ on the population loss, in practice we have a finite number of training samples and the predominance of label-sparse samples in the training set introduces difficulties due to the class imbalances.
Revised: Despite the invariance of f θ ∗ to P ξ on the population loss, in practice we have a finite number of training samples and the predominance of label-sparse samples in the training set introduces difficulties due to the class imbalances.

Original: As a specific extreme scenario for the pixel-wise classifier with encoder-decoder structure (Equation (1)), when the label-sparse samples are predominant ( ξ ≪ 1 ), a decoder function ψ θ that predicts every pixel as background can achieve near-optimal crossentropy loss, regardless of what the encoder function ϕ θ is, considerably compromising the test performance ( cf.
Revised: As a specific extreme scenario for the pixel-wise classifier with encoder-decoder structure (Equation (1)), when the label-sparse samples are predominant ( ξ ≪ 1 ), a decoder function ψ θ that predicts every pixel as background can achieve near-optimal crossentropy loss, regardless of what the encoder function ϕ θ is, considerably compromising the test performance ( cf.

Original: Table 1).
Revised: Table 1).

Original: To encourage legit encoding even in absence of sufficient dense labels, we leverage the unsupervised consistency regularization on the encoder function ϕ θ based on data augmentations.
Revised: To encourage legit encoding even in absence of sufficient dense labels, we leverage the unsupervised consistency regularization on the encoder function ϕ θ based on data augmentations.

Original: Let A be a distribution over transformations on X where for any x ∈ X , each A ∼ A ( A :
Revised: Let A be a distribution over transformations on X where for any x ∈ X , each A ∼ A ( A :

Original: X → X ) induces an augmentation A ( x ) of x that perturbs low-level information in x .
Revised: X → X ) induces an augmentation A ( x ) of x that perturbs low-level information in x .

Original: We aim to learn an encoder function ϕ θ : X → Z that is capable of filtering out low-level information from x and therefore provides similar encodings for augmentations of the same sample.
Revised: We aim to learn an encoder function ϕ θ : X → Z that is capable of filtering out low-level information from x and therefore provides similar encodings for augmentations of the same sample.

Original: Recalling the metric ϱ on Z , for a given scaling hyperparameter λ AC > 0 , we measure the similarity between augmentations with a consistency regularization term on ϕ θ ( · ) : for any A 1 , A 2 ∼ A 2 ,
Revised: Recalling the metric ϱ on Z , for a given scaling hyperparameter λ AC > 0 , we measure the similarity between augmentations with a consistency regularization term on ϕ θ ( · ) : for any A 1 , A 2 ∼ A 2 ,

Original: [Equation]
Revised: [Equation]

Original: For the n training samples { ( x i , y i ) } i ∈ [ n ] ∼ P nξ , we consider n pairs of data augmentation transformations { ( A i, 1 , A i, 2 ) } i ∈ [ n ] ∼ A 2 n .
Revised: For the n training samples { ( x i , y i ) } i ∈ [ n ] ∼ P nξ , we consider n pairs of data augmentation transformations { ( A i, 1 , A i, 2 ) } i ∈ [ n ] ∼ A 2 n .

Original: In the basic version, we encourage the similar encoding ϕ
Revised: In the basic version, we encourage the similar encoding ϕ

Original: θ ( · ) of the augmentation pairs ( A i, 1 ( x i ) , A i, 2 ( x i ))
Revised: θ ( · ) of the augmentation pairs ( A i, 1 ( x i ) , A i, 2 ( x i ))

Original: for all i ∈ [ n ] via consistency regularization:
Revised: for all i ∈ [ n ] via consistency regularization:

Original: [Equation]
Revised: [Equation]

Original: We enforce consistency on ϕ θ ( · ) in light of the encoder-decoder architecture: the encoder is generally designed to abstract essential information and filters out low-level non-semantic perturbations ( e.g. , those introduced by augmentations), while the decoder recovers the low-level information for the pixel-wise classification.
Revised: We enforce consistency on ϕ θ ( · ) in light of the encoder-decoder architecture: the encoder is generally designed to abstract essential information and filters out low-level non-semantic perturbations ( e.g. , those introduced by augmentations), while the decoder recovers the low-level information for the pixel-wise classification.

Original: Therefore, with different A 1 , A 2 ∼ A , the encoder output ϕ θ ( · ) tends to be more consistent than the other intermediate layers, especially for label-dense samples.
Revised: Therefore, with different A 1 , A 2 ∼ A , the encoder output ϕ θ ( · ) tends to be more consistent than the other intermediate layers, especially for label-dense samples.

Original:  W EIGHTED A UGMENTATION C ONSISTENCY (WAC) R EGULARIZATION
Revised:  W EIGHTED A UGMENTATION C ONSISTENCY (WAC) R EGULARIZATION

Original: As the motivation, we begin with a key observation about the averaged cross-entropy:
Revised: As the motivation, we begin with a key observation about the averaged cross-entropy:

Original: Remark 1 (Separation of averaged cross-entropy loss on P 0 and P 1 ) .
Revised: Remark 1 (Separation of averaged cross-entropy loss on P 0 and P 1 ) .

Original: As demonstrated in Figure 1, the sparse labels from P 0 tend to be much easier to learn than the dense ones from P 1 , leading to considerable separation of averaged cross-entropy losses on the sparse and dense labels after a sufficient number of training epochs – ℓ CE ( θ ; ( x , y ))
Revised: As demonstrated in Figure 1, the sparse labels from P 0 tend to be much easier to learn than the dense ones from P 1 , leading to considerable separation of averaged cross-entropy losses on the sparse and dense labels after a sufficient number of training epochs – ℓ CE ( θ ; ( x , y ))

Original: ≪ ℓ CE ( θ ; ( x ′ , y ′ ))
Revised: ≪ ℓ CE ( θ ; ( x ′ , y ′ ))

Original: for most label-sparse samples ( x , y ) ∼ P 0 and label-dense samples ( x ′ , y ′ )
Revised: for most label-sparse samples ( x , y ) ∼ P 0 and label-dense samples ( x ′ , y ′ )

Original: ∼ P 1 .
Revised: ∼ P 1 .

Original: Although Equation (4) with consistency regularization alone can boost the segmentation accuracy during testing ( cf.
Revised: Although Equation (4) with consistency regularization alone can boost the segmentation accuracy during testing ( cf.

Original: Table 4), it does not take the separation between label-sparse and label-dense samples into account.
Revised: Table 4), it does not take the separation between label-sparse and label-dense samples into account.

Original: In Section 5, we will empirically demonstrate that proper exploitation of such separation, like the formulation introduced below, can bring compelling further improvements.
Revised: In Section 5, we will empirically demonstrate that proper exploitation of such separation, like the formulation introduced below, can bring compelling further improvements.

Original: Concretely, we formalized the notion of separation between P 0 and P 1 based on the consistency regularization term (Equation (3)) with the following assumption 2 .
Revised: Concretely, we formalized the notion of separation between P 0 and P 1 based on the consistency regularization term (Equation (3)) with the following assumption 2 .

Original: Assumption 1 ( N -separation between P 0 and P 1 ) .
Revised: Assumption 1 ( N -separation between P 0 and P 1 ) .

Original: Given a sufficiently small γ > 0 , let F θ ∗ ( γ )
Revised: Given a sufficiently small γ > 0 , let F θ ∗ ( γ )

Original: =
Revised: =

Original: { θ ∈ F θ | ∥ θ − θ ∗ ∥ F ≤ γ } be a compact and convex neighborhood of well-trained pixel-wise classifiers 3 .
Revised: { θ ∈ F θ | ∥ θ − θ ∗ ∥ F ≤ γ } be a compact and convex neighborhood of well-trained pixel-wise classifiers 3 .

Original: We say that P 0 and P 1 are N -separated over F θ ∗ ( γ ) if there exists ω > 0 such that both: (i) ℓ CE ( θ ; ( x , y )) <
Revised: We say that P 0 and P 1 are N -separated over F θ ∗ ( γ ) if there exists ω > 0 such that both: (i ) ℓ CE ( θ ; ( x , y )) <

Original: ℓ AC ( θ ; x , A1 , A2 ) for all θ ∈ F θ ∗ ( γ ) given ( x , y ) ∼ P(ii) 
Revised: ℓ AC ( θ ; x , A , A ) for all θ ∈ F θ ∗ ( γ ) given ( x , y ) ∼ P(ii) ℓ

Original: ℓ CE ( θ ; ( x , y )) > 
Revised: CE ( θ ; ( x , y )) > ℓ AC ( θ ; x A

Original: ℓ AC ( θ ; x , A1 , A 2 ) for all θ ∈ F θ ∗ ( γ ) given ( x , y ) ∼ Phold with probability 1 − Ω (cid:0) N 1+ ω (cid:1) − 1 over (( x , y ) , ( A 1 , A 2 )) ∼ P ξ ×
Revised: A ) for all θ ∈ F θ ∗ ( γ ) given ( x , y ) ∼ Phold with probability 1 − Ω (cid:0) N 1+ ω (cid:1) − 1 over (( x , y ) , ( A 1 , A 2 )) ∼ P ξ ×

Original: A 2 .
Revised: A 2 .

Original: This assumption is motivated by the empirical observation that the perturbation in ϕ θ ( · ) induced by
Revised: This assumption is motivated by the empirical observation that the perturbation in ϕ θ ( · ) induced by

Original: A is more uniform across P 0 and P 1 than the averaged cross-entropy, as instantiated in Figure 3.
Revised: A is more uniform across P 0 and P 1 than the averaged cross-entropy, as instantiated in Figure 3.

Original: Under Assumption 1, up to a proper scaling hyperparameter λ AC , the consistency regularization (Equation (3)) can separate the averaged cross-entropy loss (Equation (2)) on N label-sparse and label-dense samples with probability 1 − Ω ( N ω ) − 1 by the union bound (as explained formally in
Revised: Under Assumption 1, up to a proper scaling hyperparameter λ AC , the consistency regularization (Equation (3)) can separate the averaged cross-entropy loss (Equation (2)) on N label-sparse and label-dense samples with probability 1 − Ω ( N ω ) − 1 by the union bound (as explained formally in

Original: Appendix A).
Revised: Appendix A).

Original: In particular, the larger N correspond to the stronger separation between P 0 and P 1 .
Revised: In particular, the larger N correspond to the stronger separation between P 0 and P 1 .

Original: With Assumption 1, we introduce a minimax formulation that incentivizes the separation of labelsparse and label-dense samples automatically by introducing a flexible weight β
Revised: With Assumption 1, we introduce a minimax formulation that incentivizes the separation of labelsparse and label-dense samples automatically by introducing a flexible weight β

Original: [ i ] ∈ [0 , 1] that balances ℓ CE ( θ ; ( x i , y i )) and ℓ AC ( θ ; x i , A i, 1 ,
Revised: [ i ] ∈ [0 , 1] that balances ℓ CE ( θ ; ( x i , y i )) and ℓ AC ( θ ; x i , A i, 1 ,

Original: A i, 2 ) for each of the n samples.
Revised: A i, 2 ) for each of the n samples.

Original: [Equation]
Revised: [Equation]

Original: With convex and continuous loss and regularization terms (formally in Proposition 1), Equation (5) has a saddle point where (cid:98) β separates the label-sparse and label-dense samples under Assumption 1.
Revised: With convex and continuous loss and regularization terms (formally in Proposition 1), Equation (5) has a saddle point where (cid:98) β separates the label-sparse and label-dense samples under Assumption 1.

Original: Proposition 1 (Formal proof in Appendix A) .
Revised: Proposition 1 (Formal proof in Appendix A) .

Original: Assume that ℓ CE ( θ ; ( x , y )) and ℓ AC ( θ ; x , A 1 , A 2 ) are convex and continuous in θ for all ( x
Revised: Assume that ℓ CE ( θ ; ( x , y )) and ℓ AC ( θ ; x , A 1 , A 2 ) are convex and continuous in θ for all ( x

Original: , y ) ∈
Revised: , y ) ∈

Original: X ×
Revised: X ×

Original: [ K ] d and A 1 , A 2 ∼ A 2 ; F θ ∗ ( γ ) ⊂
Revised: [ K ] d and A 1 , A 2 ∼ A 2 ; F θ ∗ ( γ ) ⊂

Original: F θ is compact and convex.
Revised: F θ is compact and convex.

Original: If P 0 and P 1 are n -separated (Assumption 1), then there exists (cid:98) β ∈ { 0 , 1 } n and (cid:98) θ WAC ∈ argmin
Revised: If P 0 and P 1 are n -separated (Assumption 1), then there exists (cid:98) β ∈ { 0 , 1 } n and (cid:98) θ WAC ∈ argmin

Original: θ ∈F θ ∗ ( γ ) (cid:98) L WAC (cid:16) θ, (cid:98) β (cid:17) such that
Revised: θ ∈F θ ∗ ( γ ) (cid:98) L WAC (cid:16) θ, (cid:98) β (cid:17) such that

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Further, (cid:98) β separates the label-sparse and label-dense samples— β [ i ]
Revised: Further, (cid:98) β separates the label-sparse and label-dense samples— β [ i ]

Original: =
Revised: =

Original: I { ( x i , y i ) ∼ P 1 } — w.h.p. .
Revised: I { ( x i , y i ) ∼ P 1 } — w.h.p. .

Original: In other words, for n samples drawn from a mixture of the n -separated P 0 and P 1 , at the saddle point,
Revised: In other words, for n samples drawn from a mixture of the n -separated P 0 and P 1 , at the saddle point,

Original: Equation (5) automatically identifies the label-sparse samples with β [ i ] = 0 , learning more from the unsupervised consistency regularization, and the label-dense ones with β [ i ] = 1 , emphasizing more on the supervised averaged cross-entropy loss.
Revised: Equation (5) automatically identifies the label-sparse samples with β [ i ] = 0 , learning more from the unsupervised consistency regularization, and the label-dense ones with β [ i ] = 1 , emphasizing more on the supervised averaged cross-entropy loss.

Original:  A DAPTIVELY W EIGHTED A UGMENTATION C ONSISTENCY ( AdaWAC )
Revised:  A DAPTIVELY W EIGHTED A UGMENTATION C ONSISTENCY ( AdaWAC )

Original: Remark 2 (Connection to hard thresholding algorithms) .
Revised: Remark 2 (Connection to hard thresholding algorithms) .

Original: The saddle point of Equation (5) is closely related to hard thresholding algorithms like Ordered SGD (Kawaguchi & Lu, 2020) and iterative trimmed loss (Shen & Sanghavi, 2019).
Revised: The saddle point of Equation (5) is closely related to hard thresholding algorithms like Ordered SGD (Kawaguchi & Lu, 2020) and iterative trimmed loss (Shen & Sanghavi, 2019).

Original: In each iteration, these algorithms update the model only on a proper subset of training samples based on the (ranking of) current empirical risks.
Revised: In each iteration, these algorithms update the model only on a proper subset of training samples based on the (ranking of) current empirical risks.

Original: Compared to hard thresholding algorithms, (i) Equation (5) additionally leverages the unused samples ( e.g. , label-sparse samples) for unsupervised consistency regularization on data augmentations, which is known for improving generalization and feature learning even in supervised settings (Shen et al., 2022; Yang et al., 2022); (ii) meanwhile, it does not require prior knowledge of the sample subpopulations ( e.g. , ξ for P ξ ) which is essential for hard thresholding algorithms.
Revised: Compared to hard thresholding algorithms, (i) Equation (5) additionally leverages the unused samples ( e.g. , label-sparse samples) for unsupervised consistency regularization on data augmentations, which is known for improving generalization and feature learning even in supervised settings (Yang et al., 2022; Shen et al., 2022); (ii) meanwhile, it does not require prior knowledge of the sample subpopulations ( e.g. , ξ for P ξ ) which is essential for hard thresholding algorithms.

Original: Equation (5) further facilitates the more flexible optimization process.
Revised: Equation (5) further facilitates the more flexible optimization process.

Original: As we will empirically show in Table 2, despite the close relation between Equation (5) and the hard thresholding algorithms (Remark 2), such updating strategies may be suboptimal for solving Equation (5). 
Revised: As we will empirically show in Table 2, despite the close relation between Equation (5) and the hard thresholding algorithms (Remark 2), such updating strategies may be suboptimal for solving Equation (5). N pretrained initialization pixel wise classifier θ 0 ∈ F [Table] [Table]

Revised: [Table]

Original:  · ℓ CE (θ t − 1 ; ( x i t , y i t )))
Revised: [Table] ] ← b [1] · [Table] ( η β · ℓ CE (

Revised: [Table]

Original:  7: b [2] ← b [2] · exp ( η β · ℓ AC ( θ t − 1 ; x i t , A i t , 1 , A i t , 2 )) 8: β t ← β t − 1 , ( β t )
Revised: t − 1 ; ( [Table] [Table] i t ))) 7: b [2] ← b [2] · exp ( η β · ℓ AC ( θ t − 1 ; x i t , A i t , 1 , A i t , 2 )) 8: β t ← β t − 1 , ( β t )

Original: [ i t ] ← b [1] / ∥ b
Revised: [ i t ] ← b [1] / ∥ b

Original: ∥9: θ t ← θ t − 1 − η θ · (cid:16) ( β t )
Revised: ∥9: θ t ← θ t − 1 − η θ · (cid:16) ( β t )

Original: [ i t ] · ∇ θ ℓ CE ( θ t − 1 ; ( x i t , y i t ))
Revised: [ i t ] · ∇ θ ℓ CE ( θ t − 1 ; ( x i t , y i t [Equation]

Original: [Equation]
Revised: [Equation]

Original: 10: end for
Revised: 10: end for

Original: Inspired by the breakthrough made by Sagawa et al.
Revised: Inspired by the breakthrough made by Sagawa et al.

Original: (2020) in the distributionally-robust optimization (DRO) setting where gradient updating on weights is shown to enjoy better convergence guarantees than hard thresholding, in Algorithm 1, we introduce an adaptive weighting algorithm for solving Equation (5) based on online mirror descent.
Revised: (2020) in the distributionally-robust optimization (DRO) setting where gradient updating on weights is shown to enjoy better convergence guarantees than hard thresholding, in Algorithm 1, we introduce an adaptive weighting algorithm for solving Equation (5) based on online mirror descent.

Original: In contrast to the commonly used stochastic gradient descent (SGD), the flexibility of online mirror descent in choosing the associated norm space not only allows gradient updates on sample weights, but also grants distinct learning dynamics to sample weights β t and model parameters θ t , which leads to the following convergence guarantee.
Revised: In contrast to the commonly used stochastic gradient descent (SGD), the flexibility of online mirror descent in choosing the associated norm space not only allows gradient updates on sample weights, but also grants distinct learning dynamics to sample weights β t and model parameters θ t , which leads to the following convergence guarantee.

Original: Proposition 2 (Formally in Proposition 3, proof in Appendix B, assumptions instantiated in Example 1) .
Revised: Proposition 2 (Formally in Proposition 3, proof in Appendix B, assumptions instantiated in Example 1) .

Original: Assume that ℓ CE ( θ ; ( x , y )) and ℓ AC ( θ ; x , A 1 , A 2 ) are convex and continuous in θ for all ( x
Revised: Assume that ℓ CE ( θ ; ( x , y )) and ℓ AC ( θ ; x , A 1 , A 2 ) are convex and continuous in θ for all ( x

Original: , y ) ∈
Revised: , y ) ∈

Original: X ×
Revised: X ×

Original: [ K ] d and A 1 , A 2 ∼ A 2 ; F θ ∗ ( γ ) ⊂
Revised: [ K ] d and A 1 , A 2 ∼ A 2 ; F θ ∗ ( γ ) ⊂

Original: F θ is convex and compact.
Revised: F θ is convex and compact.

Original: If there exist 4 (i) C θ, ∗ > 0 such that 1 n (cid:80) ni =1 (cid:13)(cid:13)(cid:13) ∇ θ (cid:98) L WAC i ( θ, β ) (cid:13)(cid:13)(cid:13) 2 F , ∗ ≤ C 2 θ, ∗ and (ii) C β , ∗ > 0 such that 1 n 
Revised: If there exist 4 (i) C θ, ∗ > 0 such that 1 n (cid:80) ni =1 (cid:13)(cid:13)(cid:13) ∇ θ (cid:98) L WAC i ( θ, β ) (cid:13)(cid:13)(cid:13) 2 F , ∗ ≤ C 2 θ, ∗ and (ii) C β , ∗ > 0 such that 1 n (cid:80) ni

Revised: =1 max { ℓ CE ( θ ; ( x i , y i )) , ℓ AC ( θ ; x i , A i, 1 , A i, 2 ) } 2 ≤ C 2 β , ∗ for all θ ∈ F θ ∗ ( γ ) , β ∈ [0 , 1] n , then with η θ = η β = 2 (cid:113) 5 T ( γ 2 C 2 θ, ∗ +2 nC 2 β , ∗ ) , Algorithm 1 provides

Original: [Equation]
Revised: [Equation]

Original: ni =1 max { ℓ CE ( θ ; ( x i , y i )) , ℓ AC ( θ ; x i , A i, 1 , A i, 2 ) } 2 ≤ C 2 β, ∗ for all θ ∈ F θ ∗ ( γ ) , β ∈ [0 , 1] n , then with η θ = η β = 2 (cid:113) 5 T ( γ 2 C 2 θ, ∗ +2 nC 2 β , ∗ ) , Algorithm 1 provides
Revised: β

Original: [Equation]
Revised: [Equation]

Original: where θ T = 1 T (cid:80) Tt =1 θ t and β T = 1 T (cid:80)
Revised: where θ T = 1 T (cid:80) Tt =1 θ t and β T = 1 T (cid:80)

Original: Tt =1 β t .
Revised: Tt =1 β t .

Original: In addition to the convergence guarantee, Algorithm 1 also demonstrates superior performance over hard thresholding algorithms for segmentation problems in practice (Table 2).
Revised: In addition to the convergence guarantee, Algorithm 1 also demonstrates superior performance over hard thresholding algorithms for segmentation problems in practice (Table 2).

Original: An intuitive explanation is that instead of filtering out all the label-sparse samples via hard thresholding, the adaptive weighting allows the model to learn from some sparse labels at the early epochs, while smoothly down-weighting ℓ CE of these samples since learning sparse labels tends to be easier (Remark 1).
Revised: An intuitive explanation is that instead of filtering out all the label-sparse samples via hard thresholding, the adaptive weighting allows the model to learn from some sparse labels at the early epochs, while smoothly down-weighting ℓ CE of these samples since learning sparse labels tends to be easier (Remark 1).

Original: With the learned model tested on a mixture of label-sparse and label-dense samples, learning sparse labels at the early stage is crucial for accurate segmentation.
Revised: With the learned model tested on a mixture of label-sparse and label-dense samples, learning sparse labels at the early stage is crucial for accurate segmentation.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we investigate the proposed AdaWAC algorithm (Algorithm 1) on different medical image segmentation tasks with different UNet-like architectures.
Revised: In this section, we investigate the proposed AdaWAC algorithm (Algorithm 1) on different medical image segmentation tasks with different UNet-like architectures.

Original: We first demonstrate the performance improvements brought by AdaWAC in terms of sample efficiency and robustness to subpopulation shift (Table 1).
Revised: We first demonstrate the performance improvements brought by AdaWAC in terms of sample efficiency and robustness to subpopulation shift (Table 1).

Original: Then, we verify the empirical advantage of AdaWAC compared to the closely related hard thresholding algorithms as discussed in Remark 2 (Table 2).
Revised: Then, we verify the empirical advantage of AdaWAC compared to the closely related hard thresholding algorithms as discussed in Remark 2 (Table 2).

Original: Our ablation study (Table 4) further illustrates the indispensability of both sample reweighting and consistency regularization, the deliberate combination of which leads to the superior performance of AdaWAC 5 .
Revised: Our ablation study (Table 4) further illustrates the indispensability of both sample reweighting and consistency regularization, the deliberate combination of which leads to the superior performance of AdaWAC 5 .

Original: Experiment setup.
Revised: Experiment setup.

Original: We conduct experiments on two volumetric medical image segmentation tasks: abdominal CT segmentation for Synapse multi-organ dataset (Synapse) 6 and cine-MRI segmentation for Automated cardiac diagnosis challenge dataset (ACDC) 7 , with two UNet-like architectures: TransUNet (Chen et al., 2021) and UNet Ronneberger et al.
Revised: We conduct experiments on two volumetric medical image segmentation tasks: abdominal CT segmentation for Synapse multi-organ dataset (Synapse) 6 and cine-MRI segmentation for Automated cardiac diagnosis challenge dataset (ACDC) 7 , with two UNet-like architectures: TransUNet (Chen et al., 2021) and UNet Ronneberger et al.

Original: (2015) (deferred to Appendix E.1).
Revised: (2015) (deferred to Appendix E.1).

Original: For the main experiments with TransUNet in Section 5, we follow the official implementation in (Chen et al., 2021) and use ERM+SGD as the baseline.
Revised: For the main experiments with TransUNet in Section 5, we follow the official implementation in (Chen et al., 2021) and use ERM+SGD as the baseline.

Original: We evaluate segmentations with two standard metrics—the average Dice-similarity coefficient (DSC) and average 95 -percentile of
Revised: We evaluate segmentations with two standard metrics—the average Dice-similarity coefficient (DSC) and average 95 -percentile of

Original: Hausdorff distance (HD95).
Revised: Hausdorff distance (HD95).

Original: Dataset and implementation details are deferred to Appendix D. Given the sensitivity of medical image semantics to perturbations, our experiments only involve simple augmentations ( i.e. , rotation and mirroring) adapted from (Chen et al., 2021).
Revised: Dataset and implementation details are deferred to Appendix D. Given the sensitivity of medical image semantics to perturbations, our experiments only involve simple augmentations ( i.e. , rotation and mirroring) adapted from (Chen et al., 2021).

Original: 5. S EGMENTATION P ERFORMANCE OF AdaWAC WITH T RANS UN ET
Revised: 5. S EGMENTATION P ERFORMANCE OF AdaWAC WITH T RANS UN ET

Original: Segmentation on Synapse.
Revised: Segmentation on Synapse.

Original: Figure 2 visualizes the segmentation predictions on 6 Synapse test slices given by models trained via AdaWAC (ours) and via the baseline (ERM+SGD) with TransUNet (Chen et al., 2021).
Revised: Figure 2 visualizes the segmentation predictions on 6 Synapse test slices given by models trained via AdaWAC (ours) and via the baseline (ERM+SGD) with TransUNet (Chen et al., 2021).

Original: We observe that AdaWAC provides more accurate predictions on the segmentation boundaries and captures small organs better than the baseline.
Revised: We observe that AdaWAC provides more accurate predictions on the segmentation boundaries and captures small organs better than the baseline.

Original: Visualization of AdaWAC .
Revised: Visualization of AdaWAC .

Original: As shown in Figure 3, with ℓ CE ( θ t ; ( x i , y i ))
Revised: As shown in Figure 3, with ℓ CE ( θ t ; ( x i , y i ))

Original: (Equation (2)) of labelsparse versus label-dense slices weakly separated in the early epochs, the model further learns to distinguish ℓ CE ( θ t ; ( x i , y i )) of label-sparse/label-dense slices during training.
Revised: (Equation (2)) of labelsparse versus label-dense slices weakly separated in the early epochs, the model further learns to distinguish ℓ CE ( θ t ; ( x i , y i )) of label-sparse/label-dense slices during training.

Original: By contrast, ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) (Equation (3)) remains mixed for all the slices in the entire training process.
Revised: By contrast, ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) (Equation (3)) remains mixed for all the slices in the entire training process.

Original: As a result, the CE weights of label-sparse slices are much smaller than those of label-dense ones, pushing AdaWAC to learn more image representations but less pixel classification for slices with sparse labels and learn more pixel classification for slices with dense labels.
Revised: As a result, the CE weights of label-sparse slices are much smaller than those of label-dense ones, pushing AdaWAC to learn more image representations but less pixel classification for slices with sparse labels and learn more pixel classification for slices with dense labels.

Original: Sample efficiency and robustness.
Revised: Sample efficiency and robustness.

Original: We first demonstrate the sample efficiency of AdaWAC in comparison to the baseline (ERM+SGD) when training only on different subsets of the full Synapse training set (“ full ” in Table 1).
Revised: We first demonstrate the sample efficiency of AdaWAC in comparison to the baseline (ERM+SGD) when training only on different subsets of the full Synapse training set (“ full ” in Table 1).

Original: Specifically, (i) half-slice contains slices with even indices only in each volume 8 ; (ii) half-vol consists of 9 volumes uniformly sampled from the total 18 volumes in full where different volumes tend to have distinct ξ s ( i.e. , ratios of label-dense samples); (iii) half-sparse takes the first half slices in each volume, most of which tend to be label-sparse ( i.e. , ξ s are made to be small).
Revised: Specifically, (i) half-slice contains slices with even indices only in each volume 8 ; (ii) half-vol consists of 9 volumes uniformly sampled from the total 18 volumes in full where different volumes tend to have distinct ξ s ( i.e. , ratios of label-dense samples); (iii) half-sparse takes the first half slices in each volume, most of which tend to be label-sparse ( i.e. , ξ s are made to be small).

Original: As shown in Table 1, the model trained with AdaWAC on half-slice generalizes as well as a baseline model trained on full , if not better.
Revised: As shown in Table 1, the model trained with AdaWAC on half-slice generalizes as well as a baseline model trained on full , if not better.

Original: Moreover, the half-vol and half-sparse experiments illustrate the robustness of AdaWAC to subpopulation shift.
Revised: Moreover, the half-vol and half-sparse experiments illustrate the robustness of AdaWAC to subpopulation shift.

Original: Furthermore, such sample efficiency and distributional robustness of AdaWAC extend to the more widely used UNet architecture.
Revised: Furthermore, such sample efficiency and distributional robustness of AdaWAC extend to the more widely used UNet architecture.

Original: We defer the detailed results and discussions on UNet to Appendix E.
Revised: We defer the detailed results and discussions on UNet to Appendix E.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Comparison with hard thresholding algorithms.
Revised: Comparison with hard thresholding algorithms.

Original: Table 2 illustrates the empirical advantage of
Revised: Table 2 illustrates the empirical advantage of

Original: AdaWAC over the hard thresholding algorithms, as suggested in Remark 2.
Revised: AdaWAC over the hard thresholding algorithms, as suggested in Remark 2.

Original: In particular, we consider the following hard thresholding algorithms: (i) trim-train learns only from slices with at least one non-background pixel and trims the rest in each iteration on the fly; (ii) trim-ratio ranks the cross-entropy loss ℓ CE ( θ t ; ( x i , y i )) in each iteration (mini-batch) and trims samples with the lowest cross-entropy losses at a fixed ratio – the ratio of all-background slices in the full training set
Revised: In particular, we consider the following hard thresholding algorithms: (i) trim-train learns only from slices with at least one non-background pixel and trims the rest in each iteration on the fly; (ii) trim-ratio ranks the cross-entropy loss ℓ CE ( θ t ; ( x i , y i )) in each iteration (mini-batch) and trims samples with the lowest cross-entropy losses at a fixed ratio – the ratio of all-background slices in the full training set

Original: ( 1 − 12802211 ≈ 0 . 42 ); (iii) ACR further incorporates the augmentation consistency regularization directly via addition of ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) without reweighting; (iv) pseudo-
Revised: ( 1 − 12802211 ≈ 0 . 42 ); (iii) ACR further incorporates the augmentation consistency regularization directly via addition of ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) without reweighting; (iv) pseudo-

Original: AdaWAC simulates the sample weights β at the saddle point and learns via ℓ CE ( θ t ; ( x i , y i )) on slices with at least one non-background pixel while via ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) otherwise.
Revised: AdaWAC simulates the sample weights β at the saddle point and learns via ℓ CE ( θ t ; ( x i , y i )) on slices with at least one non-background pixel while via ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) otherwise.

Original: We notice that naive incorporation of ACR brings less observable boosts to the hard-thresholding methods.
Revised: We notice that naive incorporation of ACR brings less observable boosts to the hard-thresholding methods.

Original: Therefore, the deliberate combination via reweighting in AdaWAC is essential for performance improvement.
Revised: Therefore, the deliberate combination via reweighting in AdaWAC is essential for performance improvement.

Original: [Table]
Revised: [Table]

Original: Segmentation on ACDC.
Revised: Segmentation on ACDC.

Original: Performance improvements granted by AdaWAC are also observed on the ACDC dataset (Table 3).
Revised: Performance improvements granted by AdaWAC are also observed on the ACDC dataset (Table 3).

Original: We defer detailed visualization of ACDC segmentation to Appendix E.
Revised: We defer detailed visualization of ACDC segmentation to Appendix E.

Original: [Table]
Revised: [Table]

Original: 5. A BLATION S TUDY
Revised: 5. A BLATION S TUDY

Original: On the influence of consistency regularization.
Revised: On the influence of consistency regularization.

Original: To illustrate the role of consistency regularization in AdaWAC , we consider the reweight-only scenario with λ AC = 

Revised: To illustrate the role of consistency regularization in AdaWAC , we consider the reweight-only scenario with λ AC =

Original: such that ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) ≡ 0 and therefore b [2] (Algorithm 1 line 7) remains intact.
Revised: such that ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) ≡ 0 and therefore b [2] (Algorithm 1 line 7) remains intact.

Original: With zero consistency regularization in AdaWAC , reweighting alone brings little improvement (Table 4).
Revised: With zero consistency regularization in AdaWAC , reweighting alone brings little improvement (Table 4).

Original: On the influence of sample reweighting.
Revised: On the influence of sample reweighting.

Original: We then investigate the effect of sample reweighting under different reweighting learning rates η β (recall Algorithm 1): (i) ACR-only for η β =(equivalent to the naive addition of ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) ), (ii) AdaWAC -0.01 for η β = 0 .
Revised: We then investigate the effect of sample reweighting under different reweighting learning rates η β (recall Algorithm 1): (i) ACR-only for η β =(equivalent to the naive addition of ℓ AC ( θ t ; x i , A i, 1 , A i, 2 ) ), (ii) AdaWAC -0.01 for η β = 0 .

Original: 01 , and (iii) AdaWAC -1.0 for η β = 1 . 0 .
Revised: 01 , and (iii) AdaWAC -1.0 for η β = 1 . 0 .

Original: As Table 4 implies, when removing reweighting from AdaWAC , augmentation consistency regularization alone improves DSC slightly from 76 .
Revised: As Table 4 implies, when removing reweighting from AdaWAC , augmentation consistency regularization alone improves DSC slightly from 76 .

Original: 28 (baseline) to 77 .(ACR-only), whereas AdaWAC boosts DSC to 79 .
Revised: 28 (baseline) to 77 .(ACR-only), whereas AdaWAC boosts DSC to 79 .

Original: 12 ( AdaWAC -1.0) with a proper choice of η β .
Revised: 12 ( AdaWAC -1.0) with a proper choice of η β .

Original: [Table]
Revised: [Table]

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this paper, we exploit the non-uniformity in labels commonly observed in volumetric medical image segmentation via AdaWAC —a deliberate combination of adaptive weighting and augmentation consistency regularization.
Revised: In this paper, we exploit the non- uniformity in labels commonly observed in volumetric medical image segmentation via AdaWAC —a deliberate combination of adaptive weighting and augmentation consistency regularization.

Original: By casting the separation between sparse and dense segmentation labels as a subpopulation shift in the label distribution, we leverage the unsupervised consistency regularization on encoder layer outputs (of UNet architectures) as a natural reference to distinguish label-sparse and label-dense samples via comparisons against the supervised average cross-entropy losses.
Revised: By casting the separation between sparse and dense segmentation labels as a subpopulation shift in the label distribution, we leverage the unsupervised consistency regularization on encoder layer outputs (of UNet architectures) as a natural reference to distinguish label-sparse and label-dense samples via comparisons against the supervised average cross-entropy losses.

Original: We formulate such comparisons as a weighted augmentation consistency (WAC) regularization problem and propose an adaptive weighting scheme— AdaWAC —for iterative and smooth separation of samples from different subpopulations with a convergence guarantee.
Revised: We formulate such comparisons as a weighted augmentation consistency (WAC) regularization problem and propose an adaptive weighting scheme— AdaWAC —for iterative and smooth separation of samples from different subpopulations with a convergence guarantee.

Original: Our experiments demonstrate empirical effectiveness of AdaWAC not only in improving segmentation performance and sample efficiency but also in enhancing robustness to the subpopulation shift in labels.
Revised: Our experiments demonstrate empirical effectiveness of AdaWAC not only in improving segmentation performance and sample efficiency but also in enhancing robustness to the subpopulation shift in labels.