# Robust Risk-Sensitive Reinforcement Learning
with Conditional Value-at-Risk

1st Xinyi Ni  Electrical and Computer Engineering
University of California, Davis
Davis, USA
<EMAIL_ADDRESS>2nd Lifeng Lai  Electrical and Computer Engineering
University of California, Davis
Davis, USA
<EMAIL_ADDRESS>
###### Abstract

Robust Markov Decision Processes (RMDPs) have received significant research
interest, offering an alternative to standard Markov Decision Processes (MDPs)
that often assume fixed transition probabilities. RMDPs address this by
optimizing for the worst-case scenarios within ambiguity sets. While earlier
studies on RMDPs have largely centered on risk-neutral reinforcement learning
(RL), with the goal of minimizing expected total discounted costs, in this
paper, we analyze the robustness of CVaR-based risk-sensitive RL under RMDP.
Firstly, we consider predetermined ambiguity sets. Based on the coherency of
CVaR, we establish a connection between robustness and risk sensitivity, thus,
techniques in risk-sensitive RL can be adopted to solve the proposed problem.
Furthermore, motivated by the existence of decision-dependent uncertainty in
real-world problems, we study problems with state-action-dependent ambiguity
sets. To solve this, we define a new risk measure named NCVaR and build the
equivalence of NCVaR optimization and robust CVaR optimization. We further
propose value iteration algorithms and validate our approach in simulation
experiments.

###### Index Terms:

ambiguity sets, RMDP, risk-sensitive RL, CVaR

##  I Introduction

Markov Decision Processes (MDP) are foundational in Reinforcement Learning
(RL), typically premised on complete knowledge of model parameters.
Nevertheless, real-world applications frequently encounter uncertainties in
MDP elements, such as transition probabilities and reward/cost functions,
leading to estimation errors in RL algorithms and subsequent sensitivity to
model inaccuracies, thus impairing performance [1, 2, 3]. In light of these
challenges, Robust MDP (RMDP) has been developed to focus on optimal policies
that accommodate worst-case transition probabilities within an ambiguity set
[4], with most studies assuming known and rectangular ambiguity sets due to
computational considerations [5, 6, 7, 4, 8, 9].

The existing RMDP research has largely focused on risk-neutral objectives that
minimize the expected total discounted costs. This risk-neutral approach does
not take events that are rare but have high costs into consideration. To
counteract this, recently many risk-sensitive approaches where risk measures
critically evaluate and quantify associated risks have been developed. Within
risk-sensitive RL, the focus is on minimizing the risk of the total discounted
cost to ascertain optimal policies [10]. Coherent risk measures, conforming to
principles of monotonicity, translation invariance, subadditivity, and
positive homogeneity, offer a robust framework for such evaluations [11].
Notably, Conditional Value-at-Risk (CVaR) has gained popularity in RL, with
numerous studies proposing CVaR RL solutions for different setups [12, 13, 14,
15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28]. Although risk-
sensitive RL is widely popular, its robustness within the RMDP framework is
not clear. While Chow et al. (2015) [16] roughly mention how solving CVaR can
enhance the robustness of risk-neutral RL in certain uncertainty sets, there
is a noticeable gap in understanding how CVaR’s robustness fares against
various types of uncertainty sets.

This study presents a novel and comprehensive investigation into the
robustness of risk-sensitive RL within RMDP. The primary goal is to determine
an optimal policy that minimizes the robust CVaR value. This value is
characterized as the highest CVaR of the total discounted cost across
transition probabilities within a defined rectangular ambiguity set. We
initially explore scenarios where the uncertain budget is fixed, and utilize
the coherent properties of CVaR and the dual representation theorem to convert
the optimization challenge into a manageable risk-sensitive RL problem,
facilitating the use of existing algorithms.

Furthermore, considering that in many real-world applications, ambiguity sets
are often dynamic and influenced by decision-making processes [29], we delve
deep into a more challenging setup about designing robust CVaR optimization
under decision-dependent uncertainty. To tackle this problem, we introduce a
new coherent risk measure NCVaR and propose a crucial decomposition theorem.
We develop value iteration algorithms for NCVaR and validate our methods
through simulation experiments. Based on these results, the emergence of NCVaR
not only enhances the robustness of CVaR RL under decision-dependent
uncertainty but also brings insights to risk-sensitive RL. Adopting NCVaR as
the risk measure for risk-sensitive RL provides strong robustness compared to
risk-neutral RL while rationally capturing risk. This makes NCVaR promising
for potential future research and also shed lights on solving decision-
dependent uncertainty for RL.

The structure of this paper is as follows: In Section II, we outline
mathematical foundations and problem formulation. Section III discusses
solutions utilizing predetermined ambiguity sets and risk-sensitive RL
methods. Section IV focuses on undetermined ambiguity sets and corresponding
value iteration algorithms. Section V validates our approaches through
experimental simulations and presents the numerical results. Conclusions are
drawn in Section VI.

##  II Preliminaries

###  II-A RMDP and Ambiguity Set

We consider a MDP represented by the tuple
$(\mathcal{X},\mathcal{A},C,P,\gamma,x_{0})$, where $\mathcal{X}$ is the state
space, $\mathcal{A}$ denotes the action space, $C(x,a)$ specifies a bounded
deterministic cost for selecting action $a$ in state $x$, $P(\cdot|x,a)$
represents the transition probability distribution, $\gamma\in[0,1]$ is the
discount factor and $x_{0}$ denotes the given initial state. For each state
$x\in\mathcal{X}$, the corresponding set of actions is represented by
$\mathcal{A}(x)$. A policy $\pi$ is a mapping from the state space to the
action space. The history space up to time $t\geq 1$ is represented as
$H_{t}=H_{t-1}\times\mathcal{A}\times\mathcal{X}$, with $H_{0}=\mathcal{X}$,
where a history $h_{t}=(x_{0},a_{0},x_{1},\dots,a_{t-1},x_{t})$ is an element
of $H_{t}$. The policy at time $t$, $\pi_{t}$, maps $h_{t}$ to a distribution
over $\mathcal{A}$. The set of such policies at time $t$ is denoted as
$\Pi_{H,t}$, with $\Pi_{H}=\lim_{t\rightarrow\infty}\Pi_{H,t}$ encompassing
all history-dependent policies. Similarly, $\Pi_{M,t}$ and
$\Pi_{M}=\lim_{t\rightarrow\infty}\Pi_{M,t}$ denote the sets of all $t$-step
and overall Markovian policies, respectively.

Addressing robustness, the transition probability $P$ is known to belong to a
non-empty, compact set $\mathcal{P}$, with the uncertain transition
probability denoted as $\tilde{P}\in\mathcal{P}$. The robust policy evaluation
over non-rectangular ambiguity sets $\mathcal{P}$ is known to be NP-hard, even
with a fixed policy $\pi$ [3]. Therefore, robust RL research often focuses on
rectangular ambiguity sets. In this work, we examine a specific rectangular
ambiguity set:

$\mathcal{P}=\big{\\{}\tilde{P}:\sum_{x^{\prime}\in\mathcal{X}}\tilde{P}(x^{\prime}|x,a)=1,\hskip
5.69054ptD(\tilde{P},P)\leq K\big{\\}},$

where $K$ is the non-negative uncertain budget and the divergence measure
$D(\tilde{P},P)$ satisfies

$D(\tilde{P},P)=\sum_{x^{\prime}\in\mathcal{X}}P(x^{\prime}|x,a)\phi\big{(}\frac{\tilde{P}(x^{\prime}|x,a)}{P(x^{\prime}|x,a)}\big{)}\leq
K.$ (1)

In (1), $\phi:\mathbb{R}\rightarrow\mathbb{R}$ represents a convex function
with the constraint $\phi(1)=0$. This function represents the
$\phi$-divergence measure, a form of divergence extensively utilized in RL
[30].

###  II-B Risk Measures

In risk-sensitive RL, risk measures play a fundamental role in quantifying and
managing risk inherent in decision-making processes. We consider a probability
space $(\Omega,\mathcal{F},P)$, where $\Omega$ represents the sample space,
$\mathcal{F}$ is a $\sigma$-algebra over the sample space, and $\mathbb{P}$ is
a probability measure. $Z:\Omega\rightarrow\mathbb{R}$ is a bounded random
variable in the probability space.

Conditional Value-at-Risk(CVaR), which is also known as the expected shortfall
or tail conditional expectation. The CVaR at confidence level $\alpha\in(0,1]$
is defined as follows [31]:

$\text{CVaR}_{\alpha}(Z)=\inf_{t\in\mathbb{R}}\big{\\{}t+\frac{1}{\alpha}\mathbb{E}_{P}\left[(Z-t)^{+}\right]\big{\\}},$

where $(z)^{+}=\max(z,0)$. One important property of CVaR is coherency and the
corresponding dual representation, which serves as a crucial factor in
establishing the equivalence between risk-sensitive RL and the robustness of
risk-sensitive RL. The dual representation for CVaR is [32]:

$\text{CVaR}_{\alpha}(Z)=\sup_{Q\in\mathcal{U}_{\text{CVaR}}}\mathbb{E}_{Q}[Z],$

where $\mathcal{U}_{\text{CVaR}}=\left\\{Q\ll
P:D_{\text{RN}}(Q,P))\in\left[0,\frac{1}{\alpha}\right]\right\\}$ with
$D_{\text{RN}}(Q,P):=\frac{Q(\omega)}{P(\omega)}$.

We also introduce another significant risk measure, known as Entropic Value-
at-Risk (EVaR). Suppose that the moment generating function
$M_{Z}(t)=\mathbb{E}_{P}\left[e^{tZ}\right]$ exists for all
$t\in\mathbb{R}^{+}$ for the random variable $Z$. In such a case, the EVaR at
a given confidence level $\alpha$ is defined as follows [33]:

$\text{EVaR}_{\alpha}(Z)=\inf_{t>0}\left\\{t^{-1}\ln(M_{Z}(t))-t^{-1}\ln\alpha\right\\}.$

It is noteworthy that EVaR is also a coherent risk measure. The dual
representation theorem for EVaR, as outlined in [33], is as follows:

$\text{EVaR}_{\alpha}(Z)=\sup_{Q\in\mathcal{U}_{\text{EVaR}}}\mathbb{E}_{Q}[Z],$

where $\mathcal{U}_{\text{EVaR}}=\left\\{Q\ll
P:D_{\text{KL}}(Q,P)\leq-\ln\alpha\right\\}$ with
$D_{\text{KL}}(Q,P):=\sum_{\omega}Q(\omega)\log\frac{Q(\omega)}{P(\omega)}$.

##  III Robust CVaR Optimization with Predetermined Ambiguity Set

In this work, the robust CVaR value is defined as the worst-case CVaR value of
a policy $\pi$ when starting from the initial state $x_{0}$ and traversing
through transition probabilities specified in the ambiguity set. The objective
is to minimize this robust CVaR value across all history-dependent policies,
as expressed by the following optimization problem:

$\min_{\pi\in\Pi_{H}}\max_{\tilde{P}\in\mathcal{P}}\text{CVaR}_{\alpha}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]}.$ (2)

The sets $\Pi_{H}$ and $\mathcal{P}$ are both non-empty and compact.
Additionally, the objective function is finite due to $\gamma<1$. Thus, the
minimum and maximum values can be achieved, as guaranteed by the Weierstrass
theorem in optimization theory [9]. This theorem ensures that the optimization
problem is well-defined and can be effectively solved to obtain the desired
policy that minimizes the robust CVaR value under the given constraints.
Contrasting with the robustness analysis of CVaR in [16], our approach
evaluates the inner CVaR objective in Equation (2) across the entire set
$\mathcal{P}$, instead of limiting the analysis to the true transition
probabilities $P$ alone. This broader evaluation provides a more comprehensive
analysis of the robustness of CVaR in diverse uncertain environments.

Recalling the coherent nature of CVaR as a risk measure and leveraging the
dual representation theorem, the original optimization problem (2) can be
reformulated as follows:

$\min_{\pi\in\Pi_{H}}\max_{\tilde{P}\in\mathcal{P}}\max_{Q\in\mathcal{U}_{\text{CVaR}}}\mathbb{E}_{Q}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]}.$ (3)

where $\mathcal{U}_{\text{CVaR}}=\\{Q\ll\tilde{P}:0\leq
Q(x^{\prime}|x,a)/\tilde{P}(x^{\prime}|x,a)\leq\frac{1}{\alpha}\\}.$ Notice
that the ${}^{\prime}\sup^{\prime}$ has been replaced by
${}^{\prime}\max^{\prime}$ since $\mathcal{U}_{\text{CVaR}}$ is convex and
compact and the objective function is continuous in $Q$.

We first focus on solving problem (3) with a predetermined ambiguity set,
where the uncertain budget remains fixed for every state and action. Our
approach involves combining two inner maximization problems by analyzing the
divergence $D(Q,P)$. Under the assumption that the function $\phi$ in (1) is
chosen such that $D(Q,P)$ remains bounded, i.e., $D(Q,P)\leq\tilde{K}$ (a
condition satisfied by the divergence measure used in this paper), we show
that problem (3) can be reformulated to:

$\min_{\pi\in\Pi_{H}}\max_{Q\in\mathcal{Q}}\mathbb{E}_{Q}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]},$ (4)

where $\mathcal{Q}=\left\\{Q:D(Q,P)\leq\tilde{K}\right\\}$ represents the
uncertain transition problem set.

This approach effectively addresses robust CVaR across diverse uncertainty
sets by combining the set’s divergence measure with the Radon-Nikodym
derivative, forming a new envelope set for risk-sensitive RL. This strategy
not only links the robustness of risk-sensitive RL with its intrinsic
transformation but also provides a universal framework for evaluating CVaR’s
robustness. We further illustrate this approach by analyzing two specific
$\phi$-divergence measures.

###  III-A Radon-Nikodym Derivative

Firstly, we consider the scenario where $\phi$-divergence is Radon-Nikodym
derivative, subject to a fixed uncertain budget for all states and actions:
$D_{\text{RN}}(\tilde{P},P)=\frac{\tilde{P}(x^{\prime}|x,a)}{P(x^{\prime}|x,a)}\in[0,K]$,
where $K\geq 0$ is a predetermined constant.

Consequently, we obtain:
$D_{\text{RN}}(Q,P)\in\left[0,\frac{K}{\alpha}\right].$ In this context, the
original optimization problem (3) transforms into:

$\min_{\pi\in\Pi_{H}}\max_{Q\in\mathcal{U}_{\text{RN}}}\mathbb{E}_{Q}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]},$ (5)

where $\mathcal{U}_{\text{RN}}=\left\\{Q\ll
P:D_{\text{RN}}(Q,P)\in[0,\frac{K}{\alpha}]\right\\}.$

Notice that solving problem (5) is equivalent to solving the following CVaR
optimization problem with confidence level $\alpha^{\prime}=\frac{\alpha}{K}$:

$\min_{\pi\in\Pi_{H}}\text{CVaR}_{\alpha^{\prime}}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]},$

which can be solve by employing CVaR value iteration algorithms proposed in
[16].

###  III-B KL Divergence

In this scenario, we consider that the uncertain transition probability
$\tilde{P}$ is defined in the neighborhood of the true transition probability
$P$ using the KL divergence, given by:
$D_{\text{KL}}(\tilde{P},P)=\sum_{x^{\prime}\in\mathcal{X}}\tilde{P}(x^{\prime}|x,a)\log\big{(}\frac{\tilde{P}(x^{\prime}|x,a)}{P(x^{\prime}|x,a)}\big{)}\leq
K,$ where $K\geq 0$ is a fixed value. Without loss of generality, we set
$K=\ln\kappa$ with $\kappa\geq 1$. We can combine the two inner maximization
problems into one, as the KL divergence of $Q$ and $P$ satisfies:
$D_{\text{KL}}(Q,P)\leq-\ln\alpha+1/\alpha\ln\kappa=-\ln(\alpha/\kappa^{\frac{1}{\alpha}}).$
Then, the original optimization problem (3) is transformed into:

$\min_{\pi\in\Pi_{H}}\max_{Q\in\mathcal{U}_{\text{KL}}}\mathbb{E}_{Q}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]},$ (6)

where $\mathcal{U}_{\text{KL}}=\left\\{Q\ll
P:D_{\text{KL}}(Q,P)\leq-\ln\frac{\alpha}{\kappa^{\frac{1}{\alpha}}}\right\\}.$

Notice that solving problem (6) is equivalent to solving the following EVaR
optimization problem with confidence level
$\alpha^{\prime}=\alpha/\kappa^{\frac{1}{\alpha}}$:

$\min_{\pi\in\Pi_{H}}\text{EVaR}_{\alpha^{\prime}}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]}.$

The problem could be solved by existing EVaR RL works [34].

##  IV Robust CVaR Optimization with Decision-Dependent Uncertainty

In real-world scenarios, ambiguity sets can dynamically change due to
decisions made during optimization, introducing endogenous uncertainty [35].
This variability means that the uncertain budget can fluctuate over time,
adding complexity to robust CVaR optimization analysis. To tackle this
decision-dependent uncertainty, we focus on the Radon-Nikodym derivative,
i.e.,

$D_{\text{RN}}(\tilde{P},P)=\frac{\tilde{P}(x^{\prime}|x,a)}{P(x^{\prime}|x,a)}\in\left[0,\vec{\kappa}(x,a)\right],\forall(x,a)\in\mathcal{X}\times\mathcal{A},$

where $\vec{\kappa}:=\left\\{\vec{\kappa}(x,a),\forall
s\in\mathcal{S},a\in\mathcal{A}\right\\}$ is the decision-dependent
uncertainty budget vector.

By combining the dual representation theorem of CVaR, we obtain the following
expression:

$D_{\text{RN}}(Q,P)=\frac{Q(x^{\prime}|x,a)}{P(x^{\prime}|x,a)}\in\big{[}0,\frac{\vec{\kappa}(x,a)}{\alpha}\big{]},\forall(x,a)\in\mathcal{X}\times\mathcal{A}.$

The problem at hand cannot be straightforwardly addressed by treating it as a
fixed confidence level CVaR optimization. To overcome this challenge, we
introduce a novel risk measure called NCVaR, which incorporates both the
confidence level $\alpha$ and an undetermined uncertain budget vector
$\vec{\kappa}$. Before delving into its definition, we set forth an assumption
to ensure that both NCVaR and the uncertain budget are meaningful.

###### Assumption 1

The undetermined uncertain budget satisfies $1\leq\vec{\kappa}(x,a)\leq
K_{\text{max}},\forall x\in\mathcal{X}$ and $a\in\mathcal{A}$. Here
$K_{\max}\geq 1$ is a real value.

We now present the formal definition of NCVaR.

###### Definition 1

For a random variable $Z:\Omega\rightarrow\mathbb{R}$ with probability mass
function (p.m.f.) $P$, the NCVaR at a given confidence level $\alpha\in(0,1]$
with an undetermined uncertain budget $\vec{\kappa}$ is defined as follows:

$\text{NCVaR}_{\alpha,\vec{\kappa}}(Z)=\sup_{Q\in\mathcal{Q}}\mathbb{E}_{Q}[Z],$
(7)

where
$\mathcal{Q}=\left\\{Q:D_{\text{RN}}(Q,P)=\frac{Q(\omega)}{P(\omega)}\in\big{[}0,\frac{\vec{\kappa}(\omega)}{\alpha}\right],\forall\omega\in\Omega\big{\\}}.$

It’s easy to observe that when $P(\omega)=0$, it implies $Q(\omega)=0$,
indicating that $Q$ is absolutely continuous with respect to $P$ (i.e., $Q\ll
P$). By leveraging Theorem 3.2 in [33], we can demonstrate that NCVaR is a
coherent risk measure, which provides a solid theoretical foundation for
employing NCVaR in practical applications and risk-sensitive RL scenarios.

As a consequence of the coherency property, solving problem (4) with an
undetermined uncertain budget defined by the Radon-Nikodym derivative is
equivalently transformed into:

$\min_{\pi\in\Pi_{H}}\text{NCVaR}_{\alpha,\vec{\kappa}}\big{[}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})\mid
x_{0},\pi\big{]}.$ (8)

Given the computational challenges associated with directly computing NCVaR,
as it requires knowledge of the entire distribution of the total discounted
cost, we present a decomposition theorem for NCVaR, which is key to
simplifying NCVaR computation and the proof is detailed in Theorem 21 of [36].

###### Theorem 1

(NCVaR Decomposition) For any $\alpha\in(0,1]$ and $\vec{\kappa}$ satisfying
Assumption 1, the $\text{NCVaR}_{\alpha,\vec{\kappa}}$ has the following
decomposition

$\begin{split}\text{NCVaR}_{\alpha,\vec{\kappa}}(Z|H_{t},\pi)&=\max_{\xi\in\mathcal{U}_{\text{NCVaR}}(\alpha,\vec{\kappa}(x_{t},a_{t}),P(\cdot|x_{t},a_{t}))}\mathbb{E}_{P}[\xi_{x_{t+1}}\\\
&\cdot\text{NCVaR}_{\alpha\xi,\vec{\kappa}}(Z|H_{t+1},\pi)|H_{t},\pi],\end{split}$

where $\xi(x_{t+1})=\frac{Q(x^{\prime}|x,a)}{P(x^{\prime}|x,a)}\geq 0$ is in
the set

$\begin{split}&\mathcal{U}_{\text{NCVaR}}(\alpha,\vec{\kappa}(x_{t},a_{t}),P(\cdot|x_{t},a_{t}))\\\
&=\big{\\{}\xi:\xi(x_{t+1})\in\big{[}0,\frac{\vec{\kappa}(x_{t},a_{t})}{\alpha})\big{]},\\\
&\hskip
28.45274pt\sum_{x_{t+1}\in\mathcal{X}}\xi(x_{t+1})P(x_{t+1}|x_{t},a_{t})=1\big{\\}}.\end{split}$

This decomposition theorem provides a valuable insight to NCVaR computation,
effectively linking the risk measure between different states, and facilitates
a more tractable approach to handling the complexity of NCVaR evaluation
within risk-sensitive RL under the RMDP framework. In light of the distinct
confidence levels on both sides of equation (1), we introduce an augmented
continuous space $\mathcal{Y}=(0,1]$ to represent the domain of confidence
levels.

Accordingly, the value-function $V(x,y)$ for every
$(x,y)\in\mathcal{X}\times\mathcal{Y}$ is defined as:

$\begin{split}&V(x,y)\\\
&=\min_{\pi\in\Pi_{H}}\text{NCVaR}_{y,\vec{\kappa}}\big{(}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})|x_{0}=x,\pi\big{)}.\end{split}$

The Bellman operator
$\mathbf{T}:\mathcal{X}\times\mathcal{Y}\rightarrow\mathcal{X}\times\mathcal{Y}$
is defined as:

$\begin{split}\mathbf{T}[V](x,y)=&\min_{a\in\mathcal{A}}\big{[}C(x,a)+\gamma\max_{\xi\in\mathcal{U}_{\text{NCVaR}}(y,\vec{\kappa}(x,a),P(\cdot|x,a))}\\\
&\sum_{x^{\prime}\in\mathcal{X}}\xi(x^{\prime})V(x^{\prime},y\xi(x^{\prime}))P(x^{\prime}|x,a)\big{]}.\end{split}$

Lemma 1 introduces some important properties for the NCVaR Bellman operator.

###### Lemma 1

The Bellman operator $\mathbf{T}$ has the following properties: P1)
Monotonicity; P2) Transition invariance; P3) Contraction; P4) Concavity
preserving: Suppose $yV(x,y)$ is concave in $y\in\mathcal{Y},\forall
x\in\mathcal{X}$. Then the maximization problem in (8) is concave and
$y\mathbf{T}[V](x,y)$ is also concave in y.

Properties P1-P3 are similar to standard dynamic programming [37], and are key
to design a convergent value iteration method. P4 ensures that value-iteration
updates involve concave, and thus tractable, optimization problems.

Based on Lemma 1, we are able to propose the following theorem, which
demonstrates the existence of a unique fixed-point solution and outline a
method for deriving an optimal policy.

###### Theorem 2

The unique fixed-point solution $V^{*}(x,y)$ of $\mathbf{T}[V](x,y)=V(x,y)$
exists and equals to the optimal value of optimization problem (8), i.e.,

$\begin{split}&V^{*}(x,y)\\\
&=\min_{\pi\in\Pi_{H}}\text{NCVaR}_{y,\vec{\kappa}}\big{(}\lim_{T\rightarrow\infty}\sum_{t=0}^{T}\gamma^{t}C(x_{t},a_{t})|x_{0}=x,\pi\big{)}.\end{split}$

Although the problem is optimized over history-dependent policies, we
demonstrate that an optimal Markov policy exists, from which the optimal
history-dependent policy can be derived. Considering the easier implementation
of the Markov policy, we adopt the greedy policy w.r.t $V^{*}(x,y)$ as the
optimal policy.

We introduce Algorithm 1 to effectively solve the NCVaR optimization problem.
This solution is equivalent to addressing the original problem incorporating
an undetermined uncertain budget defined by the Radon-Nikodym derivative.

Algorithm 1 Value Iteration for NCVaR

1: for $x\in\mathcal{X}$ and $y\in\mathcal{Y}$ do

2: arbitrarily initialize $V_{0}(x,y)$

3: end for

4: for $t=0,1,2,\dots$ do

5: for $x\in\mathcal{X}$ and $y\in\mathcal{Y}$ do

6: $\hskip 22.76219ptV_{t+1}(x,y)=\mathbf{T}[V_{t}](x,y)$

7: end for

8: end for

9: set $V^{*}(x,y)=\lim_{t\rightarrow\infty}V_{t}(x,y)$, then construct
$\pi^{*}$ as the greedy policy w.r.t $V^{*}(x,y)$

However, implementing Algorithm 1 directly can be challenging due to the
continuous nature of the set $\mathcal{Y}$. To address this issue, we employ a
sampling approach, where we select multiple points in $\mathcal{Y}$ and
subsequently utilize linear interpolation to derive the value function $V$.
However, to guarantee convergence, we need to satisfy the following assumption
for the initial value function $V_{0}$.

###### Assumption 2

The initial value function $V_{0}(x,y)$ is continuous and bounded in
$y\in\mathcal{Y}$ for any $x\in\mathcal{X}$. Also, $yV_{0}(x,y)$ is concave in
$y\in\mathcal{Y}$.

Let $N(x)$ denote the number of sample points, and
$Y(x)={y_{1},y_{2},\dots,y_{N(x)}}\in[0,1]^{N(x)}$ be the corresponding
confidence level set. Notably, we have $y_{1}=0$ and $y_{N(x)}=1$. To perform
linear interpolation of $yV(x,y)$, we define the interpolation function
$\mathcal{I}_{x}V$ as follows:

$\begin{split}&\mathcal{I}_{x}[V](y)\\\
&=y_{i}V(x,y_{i})+\frac{y_{i+1}V(x,y_{i+1})-y_{i}V(x,y_{i})}{y_{i+1}-y_{i}}(y-y_{i}),\end{split}$
(9)

where $y_{i}$ and $y_{i+1}$ are the closest points such that
$y\in[y_{i},y_{i+1}]$. With this, we introduce the interpolated Bellman
operator for NCVaR, denoted as $\mathbf{T}_{\mathcal{I}}V$:

$\begin{split}\mathbf{T}_{\mathcal{I}}[V](x,y)=\min_{a\in\mathcal{A}}\big{[}&C(x,a)+\gamma\max_{\xi\in\mathcal{U}_{\text{NCVaR}}(y,P(\cdot|x,a))}\\\
&\sum_{x^{\prime}\in\mathcal{X}}\frac{\mathcal{I}_{x^{\prime}}[V](y\xi(x^{\prime}))}{y}P(x^{\prime}|x,a)\big{]}.\end{split}$
(10)

An essential observation regarding the interpolated Bellman operator is that
it also exhibits the properties stated in Lemma 1. This can be demonstrated by
employing a similar approach as used in the proof of Lemma 1. Moreover, we
present a more practical and applicable version of the value iteration
algorithm in Algorithm 2. This algorithm utilizes the interpolated Bellman
operator and leverages linear interpolation to achieve the near-optimal value
function and near-optimal policy.

Algorithm 2 NCVaR Value Iteration with Linear Interpolation

1: choose $Y(x)$, $V_{0}(x,y)$ satisfying Assumption 2

2: for $t=0,1,2,\dots$ do

3: for $x\in\mathcal{X}$ and $y\in\mathcal{Y}$ do

4: $\hskip 22.76219ptV_{t+1}(x,y)=\mathbf{T}_{\mathcal{I}}[V_{t}](x,y)$

5: end for

6: end for

7: set $V^{*}(x,y)=\lim_{t\rightarrow\infty}V_{t}(x,y)$, then construct
$\pi^{*}$ as the greedy policy w.r.t $V^{*}(x,y)$

##  V Experiment

In this study, we adopt an experimental setup that aligns with previous works
[16, 34], ensuring comparability and consistency in our results. We use a
$64\times 53$ grid world RL environment with a state space representing all
positions. The agent starts at $(60,50)$, aiming to reach $(60,2)$. It can
move east, south, west, or north, transitioning to adjacent states with a
probability of $0.95$, or to any other neighboring state with a probability of
$0.05/3$. The environment has $80$ obstacles; colliding with one incurs a $40$
cost, while safe movements cost 1. The agent’s goal is to find a secure and
cost-effective path. For value iteration with linear interpolation, we use
$21$ sample points, following the rule $y_{i+1}=\theta y_{i}$ for
$i=1,2,\dots,20$.

(a) $\alpha=0.48$, no uncertainty

(b) $\alpha=0.48$, $K_{\text{RN}}=2$

(c) $\alpha=0.48$, $K_{\text{KL}}=2$

(d) $\alpha=0.48$, $K_{\text{unfix}}\in[1,2]$

Figure 1: Optimal value function and path in robust CVaR optimization across
various uncertainty sets.

We first validate our approach for a fixed uncertain budget using Radon-
Nikodym derivative and KL divergence. This involves visualizing the optimal
value function with color variations (a bluer color indicates a lower risk
while a yellower color indicates a higher risk) and tracing the optimal path
as a red line (Figure 1(a)). In Figure 1(a), 1(b) and 1(c), we select a
confidence level of $\alpha=0.48$ and an uncertain budget of $K=2$ for both RN
derivative and KL divergence. Consequently, we obtain
$\alpha^{\prime}_{\text{CVaR}}=0.24$ and $\alpha^{\prime}_{\text{EVaR}}=0.03$,
which indicates that the new optimal policy will exhibit a more risk-averse
behavior compared to the original one. Accordingly, the optimal path becomes
longer and is positioned closer to obstacles, aligning with the result that
the value function is larger. We further assess Algorithm 2 for decision-
dependent cases, setting the uncertain budget range to $[1,2]$. As a result,
for a fixed current state $x$, the new confidence level on the right side of
the decomposition theorem significantly deviates from the fixed case. This
increased deviation leads to the agent becoming more risk-averse as shown in
Figure 1(d). In conclusion, our algorithms effectively induce risk-averse
policies, equipping agents to navigate more cautiously in uncertain
environments. The experiments validate our methodology’s efficacy in guiding
agents towards safer decision-making strategies.

##  VI Conclusion and Future Direction

In this study, we have conducted a comprehensive and novel analysis of robust
CVaR-based risk-sensitive RL within the framework of RMDP. We have
successfully addressed robust CVaR optimization in the presence of fixed
uncertain budgets while adopting a rectangular ambiguity set. We have
introduced a novel risk measure NCVaR and devised NCVaR value iteration
algorithms to solve the challenges associated with state-action dependent
uncertainty. Furthermore, we have demonstrated the convergence of our
algorithms through theoretical analysis. We have validated the proposed
approaches through simulation experiments, and the results showcased the
effectiveness and practicality of our methods. This paper leaves several
interesting directions for future works, including the extension of robustness
analysis to a broader spectrum of uncertainty sets and a deeper exploration of
NCVaR within risk-sensitive RL.

## References

  * [1] K. Zhang, T. Sun, Y. Tao, S. Genc, S. Mallya, and T. Basar, “Robust Multi-Agent Reinforcement Learning with Model Uncertainty,” _Advances in Neural Information Processing Systems_ , vol. 33, pp. 10 571–10 583, 2020.
  * [2] Y. Le Tallec, “Robust, Risk-Snsitive, and Data-Driven Control of Markov Decision Processes,” Ph.D. dissertation, Massachusetts Institute of Technology, 2007.
  * [3] W. Wiesemann, D. Kuhn, and B. Rustem, “Robust Markov Decision Processes,” _Mathematics of Operations Research_ , vol. 38, no. 1, pp. 153–183, 2013.
  * [4] C. Ho, M. Petrik, and W. Wiesemann, “Fast Bellman Updates for Robust MDPs,” in _Proc. International Conference on Machine Learning_ , Stockholm, Sweden, July. 2018, pp. 1979–1988.
  * [5] A. Nilim and L. Ghaoui, “Robustness in Markov Decision Problems with Uncertain Transition Matrices,” _Advances in Neural Information Processing Systems_ , vol. 16, 2003.
  * [6] G. Iyengar, “Robust Dynamic Programming,” _Mathematics of Operations Research_ , vol. 30, no. 2, pp. 257–280, 2005.
  * [7] A. Ben-Tal, D. Den Hertog, A. De Waegenaere, B. Melenberg, and G. Rennen, “Robust Solutions of Optimization Problems Affected by Uncertain Probabilities,” _Management Science_ , vol. 59, no. 2, pp. 341–357, 2013\.
  * [8] Y. Wang and S. Zou, “Online Robust Reinforcement Learning with Model Uncertainty,” _Advances in Neural Information Processing Systems_ , vol. 34, pp. 7193–7206, 2021.
  * [9] C. Ho, M. Petrik, and W. Wiesemann, “Robust Phi-Divergence MDPs,” _arXiv preprint arXiv:2205.14202_ , 2022.
  * [10] O. Mihatsch and R. Neuneier, “Risk-Sensitive Reinforcement Learning,” _Machine Learning_ , vol. 49, pp. 267–290, 2002.
  * [11] P. Artzner, F. Delbaen, J. Eber, and D. Heath, “Coherent Measures of Risk,” _Mathematical Finance_ , vol. 9, no. 3, pp. 203–228, Jul. 1999\.
  * [12] H. Kashima, “Risk-Sensitive Learning via Minimization of Empirical Conditional Value-at-Risk,” _IEICE Transactions on Information and Systems_ , vol. 90, no. 12, pp. 2043–2052, Dec. 2007.
  * [13] A. Tamar, Y. Glassner, and S. Mannor, “Policy Gradients beyond Expectations: Conditional Value-at-Risk,” _arXiv preprint arXiv:1404.3862_ , 2014.
  * [14] Y. Chow and M. Ghavamzadeh, “Algorithms for CVaR Optimization in MDPs,” _Advances in Neural Information Processing Systems_ , vol. 27, 2014\.
  * [15] L. Prashanth, “Policy Gradients for CVaR-Constrained MDPs,” in _International Conference on Algorithmic Learning Theory_. Springer, 2014, pp. 155–169.
  * [16] Y. Chow, A. Tamar, S. Mannor, and M. Pavone, “Risk-Sensitive and Robust Decision-Making: A CVaR Optimization Approach,” _arXiv preprint arXiv:1506.02188_ , 2015.
  * [17] A. Tamar, Y. Glassner, and S. Mannor, “Optimizing the CVaR via Sampling,” in _Proc. AAAI Conference on Artificial Intelligence_ , Austin, TX, Feb. 2015, pp. 2993–2999.
  * [18] Y. Chow, M. Ghavamzadeh, L. Janson, and M. Pavone, “Risk-Constrained Reinforcement Learning with Percentile Risk Criteria,” _The Journal of Machine Learning Research_ , vol. 18, no. 1, pp. 6070–6120, Jan. 2017\.
  * [19] S. Stanko and K. Macek, “Risk-Averse Distributional Reinforcement Learning: A CVaR Optimization Approach.” in _International Journal of Child-Computer Interaction_ , 2019, pp. 412–423.
  * [20] X. Ma, L. Xia, Z. Zhou, J. Yang, and Q. Zhao, “Dsac: Distributional Soft Actor Critic for Risk-Sensitive Reinforcement Learning,” _arXiv preprint arXiv:2004.14547_ , 2020.
  * [21] R. Keramati, C. Dann, A. Tamkin, and E. Brunskill, “Being Optimistic to Be Conservative: Quickly Learning a CVaR Policy,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 34, no. 04, 2020, pp. 4436–4443.
  * [22] M. Godbout, M. Heuillet, S. Chandra, R. Bhati, and A. Durand, “CARL: Conditional-Value-at-Risk Adversarial Reinforcement Learning,” _arXiv preprint arXiv:2109.09470_ , Sep. 2021.
  * [23] D. Kim and S. Oh, “Trc: Trust Region Conditional Value at Risk for Safe Reinforcement Learning,” _IEEE Robotics and Automation Letters_ , vol. 7, no. 2, pp. 2621–2628, 2022.
  * [24] C. Ying, X. Zhou, H. Su, D. Yan, N. Chen, and J. Zhu, “Towards Safe Reinforcement Learning via Constraining Conditional Value-at-Risk,” _arXiv preprint arXiv:2206.04436_ , 2022.
  * [25] S. H. Lim and I. Malik, “Distributional Reinforcement Learning for Risk-Sensitive Policies,” _Advances in Neural Information Processing Systems_ , vol. 35, pp. 30 977–30 989, 2022.
  * [26] I. Greenberg, Y. Chow, M. Ghavamzadeh, and S. Mannor, “Efficient Risk-Averse Reinforcement Learning,” _Advances in Neural Information Processing Systems_ , vol. 35, pp. 32 639–32 652, 2022.
  * [27] K. Wang, N. Kallus, and W. Sun, “Near-Minimax-Optimal Risk-Sensitive Reinforcement Learning with CVaR,” in _International Conference on Machine Learning_. PMLR, 2023, pp. 35 864–35 907.
  * [28] Q. Zhang, S. Leng, X. Ma, Q. Liu, X. Wang, B. Liang, Y. Liu, and J. Yang, “CVaR-Constrained Policy Optimization for Safe Reinforcement Learning,” _IEEE Transactions on Neural Networks and Learning Systems_ , 2024.
  * [29] O. Nohadani and K. Sharma, “Optimization Under Decision-Dependent Uncertainty,” _SIAM Journal on Optimization_ , vol. 28, no. 2, pp. 1773–1795, 2018.
  * [30] C. Gong, Q. He, Y. Bai, Z. Yang, X. Chen, X. Hou, X. Zhang, Y. Liu, and G. Fan, “The $f$-Divergence Reinforcement Learning Framework,” _arXiv preprint arXiv:2109.11867_ , 2021.
  * [31] R. T. Rockafellar and S. Uryasev, “Optimization of Conditional Value-at-Risk,” _Journal of Risk_ , vol. 2, pp. 21–42, Apr. 2000.
  * [32] M. Ang, J. Sun, and Q. Yao, “On the Dual Representation of Coherent Risk Measures,” _Annals of Operations Research_ , vol. 262, no. 1, pp. 29–46, Mar. 2018.
  * [33] A. Ahmadi-Javid, “Entropic Value-at-Risk: A New Coherent Risk Measure,” _Journal of Optimization Theory and Applications_ , vol. 155, no. 3, pp. 1105–1123, 2012.
  * [34] X. Ni and L. Lai, “Risk-Sensitive Reinforcement Learning via Entropic-VaR Optimization,” in _Proc. Asilomar Conference on Signals, Systems, and Computers_ , Pacific Grove, CA, Oct. 2022, pp. 953–959.
  * [35] F. Luo and S. Mehrotra, “Distributionally Robust Optimization with Decision Dependent Ambiguity Sets,” _Optimization Letters_ , vol. 14, pp. 2565–2594, 2020.
  * [36] G. C. Pflug and A. Pichler, “Time-Consistent Decisions and Temporal Decomposition of Coherent Risk Functionals,” _Mathematics of Operations Research_ , vol. 41, no. 2, pp. 682–699, May. 2016.
  * [37] D. Bertsekas, _Dynamic Programming and Optimal Control: Volume I_. Nashua, NA: Athena scientific, 2012, vol. 1.


# Progressive Feature Self-Reinforcement for
Weakly Supervised Semantic Segmentation

Jingxuan He1, Lechao Cheng1, Chaowei Fang3, Zunlei Feng2, Tingting Mu4, Mingli
Song2 corresponding author.

###### Abstract

Compared to conventional semantic segmentation with pixel-level supervision,
weakly supervised semantic segmentation (WSSS) with image-level labels poses
the challenge that it always focuses on the most discriminative regions,
resulting in a disparity between fully supervised conditions. A typical
manifestation is the diminished precision on the object boundaries, leading to
a deteriorated accuracy of WSSS. To alleviate this issue, we propose to
adaptively partition the image content into certain regions (e.g., confident
foreground and background) and uncertain regions (e.g., object boundaries and
misclassified categories) for separate processing. For uncertain cues, we
propose an adaptive masking strategy and seek to recover the local information
with self-distilled knowledge. We further assume that the unmasked confident
regions should be robust enough to preserve the global semantics. Building
upon this, we introduce a complementary self-enhancement method that
constrains the semantic consistency between these confident regions and an
augmented image with the same class labels. Extensive experiments conducted on
PASCAL VOC 2012 and MS COCO 2014 demonstrate that our proposed single-stage
approach for WSSS not only outperforms state-of-the-art benchmarks remarkably
but also surpasses multi-stage methodologies that trade complexity for
accuracy. The code can be found at https://github.com/Jessie459/feature-self-
reinforcement.

## Introduction

Weakly supervised semantic segmentation (WSSS) reduces the cost of annotating
“strong” pixel-level labels by using “weak” labels such as bounding boxes
(Dai, He, and Sun 2015; Song et al. 2019), scribbles (Lin et al. 2016; Vernaza
and Chandraker 2017), points (Bearman et al. 2016; Su et al. 2022) and image-
level class labels (Araslanov and Roth 2020; Ru et al. 2022; Wu et al. 2023;
Ru et al. 2023). Among these, image-level class labels are the most
affordable, but challenging to exploit. A commonly used WSSS approach based on
image-level class labels typically includes the following steps: (1) to train
a neural network for image classification; (2) to use the network to generate
class activation maps (CAMs) (Zhou et al. 2016) as seed regions; (3) to refine
the CAMs to pseudo segmentation labels that will be used as the ground truth
for supervising a segmentation network. These steps can either be implemented
as separate stages or as a single collaborative stage, and single-stage
frameworks are usually more efficient as they streamline the training
pipeline. In general, high-quality pseudo labels lead to superior semantic
segmentation performance. In this work, we focus on developing an effective
single-stage approach to generate more accurate pseudo labels from image-level
class labels.

Figure 1: Our main idea. The flawed CAM only identifies discriminative
regions. To solve this, we propose to partition the image into uncertain
regions (e.g., object boundaries) and confident regions (e.g., the main body
of an object) and reinforce features of these regions in a complementary way.

Unfortunately, CAMs are essentially flawed because they are intended for
classification, i.e., they strive to identify the most discriminative regions
of an object aiming at improved classification accuracy. To tackle this, one
can improve the initial seeds (Lee et al. 2019; Wang et al. 2020) or refine
pseudo labels (Ahn, Cho, and Kwak 2019; Chen et al. 2020), by expanding
activations or labels to semantically consistent pixels in the neighborhood.
Recent studies have found that the restricted receptive field of convolution
negatively affects the recognition of integral objects (Ru et al. 2022, 2023)
and use vision transformer (Dosovitskiy et al. 2020) to model the global
relationships for improvement. But this does not resolve the issue of CAM
seeds or pseudo labels, and we still observe empirically high uncertainty in
(1) boundary regions between foreground objects and background, and (2)
misclassified regions within multiple semantically-different objects. In the
example of Figure 1, the generated CAM is uncertain about the two arms of the
person on the chair, also the boundary between the foreground (person and
chair) and the background is unclear. These uncertain regions are easily
confused by obscure semantic clues due to the absence of pixel-level
supervision.

Our goal is to clarify the visual semantics of uncertain regions mentioned
above. We emphasize that the local visual patterns should be explicitly
modeled and captured. As can be seen from Figure 1, head and upper thighs are
well recognized, while the recognition of arms and lower legs needs
improvement. A better understanding of that arms and lower legs surely belong
to a person should be established using local visual context. Although some
methods can deal with noisy object boundaries by employing off-the-shelf
saliency detection models for rich object contours (Lee et al. 2021b; Li, Fan,
and Zhang 2022), they overlook uncertain regions caused by low confidence
within objects. Alternatively, it has been proposed to attain the training
objective using knowledge gathered from the past training iterations, i.e.,
self-distillation (Caron et al. 2021). Encouraged by the success of self-
distillation, we discard saliency detection models, but take advantage of the
strategy of self-distillation in our model training.

To this end, to explore and strengthen semantics over uncertain regions, we
propose a progressive self-reinforcement method. To distinguish uncertain
regions from confident ones, we define those with intermediate CAM scores as
uncertain regions, since a very low/high score strongly indicates the
background/foreground. Specifically, we propose to mask uncertain features
(equivalent to image patch tokens) and learn to recover the original
information with the help of an online momentum teacher. This masking strategy
aligns with a state-of-the-art pre-training paradigm called masked image
modeling (MIM) that brings locality inductive bias to the model (Xie et al.
2023). We upgrade its random masking strategy with semantic uncertainty so
that the network can focus on uncertain features controlled by the masking
ratio. This design is beneficial to facilitate features in both object
boundaries and misclassified regions. Assuming that confident features should
be robust enough to present global semantics, we also introduce a
complementary method that constrains semantic consistency between two
augmented views with the same class labels. Our proposal can be seamlessly
integrated into a vision transformer based single-stage WSSS framework. We
summarize our main contributions as follows :

  * •

We propose a novel WSSS approach, _progressive feature self-reinforcement_ ,
to effectively enhance the semantics of uncertain regions. The investigation
of uncertain regions, including both object boundaries and misclassified
categories, significantly improves WSSS performance.

  * •

We design an adaptive masking strategy to identify uncertain regions. Unlike
most previous works that adopt additional saliency detection models, we locate
uncertain regions with the guidance of semantic-aware CAMs.

  * •

Exhaustive experiments on PASCAL VOC 2012 (Everingham et al. 2010) and MS COCO
2014 (Lin et al. 2014) show that our method outperforms SOTA single-stage
competitors, even better than existing sophisticated multi-stage methods.

## Related Work

Figure 2: Overview of our framework. For the student pipeline, we forward one
view through the encoder, and the encoded patch tokens are fed into the
classifier for classification and the decoder for semantic segmentation,
separately. The other view is masked and sequentially forwarded through the
encoder, the aggregation module, and the projector. For the teacher pipeline,
both views are propagated through the encoder, the aggregation module, and the
projector. The teacher network requires no gradient and is an exponential
moving average (EMA) of the student network.

### Weakly Supervised Semantic Segmentation

Multi-stage WSSS methods adopt a classification model to generate CAMs as
pseudo labels, then train a segmentation model for evaluating the final
performance. To overcome the commonly acknowledged weakness that CAMs can only
focus on discriminative regions, several works aim at improving the training
dynamic by erasing and seeking (Hou et al. 2018) or adversarial learning (Yoon
et al. 2022). Some recent approaches also adopt vision transformer
(Dosovitskiy et al. 2020) for the WSSS task, considering its favorable long-
range modeling capability. TS-CAM (Gao et al. 2021) proposes to couple class-
agnostic attention maps with semantic-aware patch tokens to promote object
localization. MCTformer (Xu et al. 2022) introduces multiple class tokens so
that class-specific attention maps can be generated. Other approaches
incorporate extra data into training or post-processing, e.g., saliency maps
(Lee et al. 2021b) or contrastive language-image pre-training (CLIP) models
(Lin et al. 2023). Our solution aims at improving pseudo labels as well, but
it is integrated into a single-stage framework for simplicity, and it requires
neither extra data nor off-the-shelf saliency detection models.

Single-stage WSSS methods treat multiple stages such as classification, pseudo
label refinement, segmentation as a whole and perform joint training. 1Stage
(Araslanov and Roth 2020) achieves comparable performance with dominant multi-
stage approaches by ensuring local consistency, semantic fidelity and mask
completeness. AFA (Ru et al. 2022) explores the intrinsic architecture of ViT
and derives reliable semantic affinity from multi-head self-attention for
pseudo label refinement. ToCo (Ru et al. 2023) tackles the issue of over-
smoothing observed in ViT by supervising the final patch tokens with
intermediate knowledge. Despite the simplified and streamlined training
procedure, single-stage methods often suffer from inferior performance
compared with multi-stage ones. In this work, we achieve superior semantic
segmentation results using a single-stage framework by discovering and
reinforcing underlying semantic layouts.

### Self-Distillation

Self-distillation associates self-supervised learning (He et al. 2020) with
knowledge distillation (Hinton, Vinyals, and Dean 2015), where knowledge is
transferred and learned without resorting to any labels. It is primarily
designed to compress large networks, and is hoping to promote performance on
downstream tasks via mimicking the output of a frozen teacher (Noroozi et al.
2018; Zhang et al. 2023; Wang et al. 2023). Recently, some approaches (Caron
et al. 2021; Zhou et al. 2021) build the teacher dynamically during training,
where the teacher adopts the same architecture as that of the student and is
updated with the knowledge of past iterations. The resulting framework
simplifies the training process and achieves compelling results compared with
other self-training frameworks. This motivates us to adapt the core idea of
self-distillation to the WSSS task for the purpose of rectifying inaccurate
object boundaries as well as improving discriminative object features.

## Methodology

### A Single-Stage Framework for WSSS

The proposed single-stage framework for WSSS is illustrated in Figure 2. We
use an encoder-decoder architecture to accomplish semantic segmentation with
image-level supervision. The encoder is a vision transformer supervised by
image-level class labels. We adopt patch token contrast (PTC) (Ru et al. 2023)
for affinity learning as it is crucial to constrain affinities between patch
tokens of the last layer against over-smoothing (Gong et al. 2021). As for
semantic segmentation, we borrow a lightweight convolutional decoder from
DeepLab (Chen et al. 2017), which is supervised by pseudo segmentation labels
that are generated from CAMs. An aggregation module is used to summarize patch
tokens into one class token and an MLP-based projector to transform all tokens
into an appropriate feature space for feature learning. To improve model
training, we enable a student and a teacher pipeline to achieve self-
distillation.

Formally, let $\mathcal{F}$ be the transformer encoder with its output
embedding dimension denoted by $D$, $\mathcal{P}$ the projector, $\mathcal{M}$
the masking operator, and $\mathcal{A}$ the aggregating operator. We start
from explaining the student pipeline. An input image is randomly augmented to
two distorted views: $x_{1}$ and $x_{2}$. Each view is subsequently divided
into $HW$ non-overlapping patch tokens, denoted as
$T_{1}=\left\\{t_{1}^{(i)}\right\\}_{i=1}^{HW}$ and
$T_{2}=\left\\{t_{2}^{(i)}\right\\}_{i=1}^{HW}$, respectively. We forward
$T_{1}$ into the encoder to obtain the logits
$Z_{1}=\mathcal{F}(T_{1})\in\mathbb{R}^{HW\times D}$, which are then fed into
the classifier for classification, and also the decoder for segmentation,
following the standard image classification and segmentation setup. To
reinforce features, we divide $T_{2}$ into uncertain and confident tokens and
mask the uncertain ones with learnable parameters, for which the uncertain
token selection and masking approaches will be explained later. The resulting
masked view, denoted as $\hat{T}_{2}=\mathcal{M}(T_{2})$, is also fed into the
encoder to obtain $\hat{Z}_{2}=\mathcal{F}\left(\hat{T}_{2}\right)$.
Embeddings of the unmasked confident tokens in $\hat{Z}_{2}$ are summarized
into a class token by an aggregation module, denoted by
$\mathcal{A}\left(\hat{Z}_{2}\right)\in\mathbb{R}^{1\times D}$. This class
token is concatenated with $\hat{Z}_{2}$, and further projected and normalized
to resemble probabilities distributions in
$\hat{P}_{2}\in\mathbb{R}^{(1+HW)\times D}$, as

$\hat{P}_{2}=\sigma\left(\mathcal{P}\left(\left[\mathcal{A}\left(\hat{Z}_{2}\right);\hat{Z}_{2}\right]\right)\right),$
(1)

where $\sigma$ is the row-wise softmax function, and $[;]$ the concatenation.
We will explain the aggregation design later.

The teacher shares the same architecture as the student’s encoder and
projector, and has a similar pipeline described by Eq. (1), except it takes
the unmasked inputs $T_{1}$ and $T_{2}$, and returns two distributions $P_{1}$
and $P_{2}$ for the two views, respectively. The student output $\hat{P}_{2}$
and the teacher outputs $P_{1}$ and $P_{2}$ are used for feature reinforcement
training.

### Uncertain Patch Token Selection

We select uncertain patch tokens under the guidance of semantic-aware CAMs,
generated using the logits computed earlier with the first view, i.e.,
$Z_{1}=\mathcal{F}(T_{1})$. We linearly project $Z_{1}$ using the weights
$W\in\mathbb{R}^{C\times D}$ of the classifier for image classification, where
$C$ is the class number, and then normalize it by the $\operatorname{ReLU}$
function and $\operatorname{min-max}$ normalization. The normalized CAM,
denoted as $M_{c}\in\mathbb{R}^{HW\times C}(0\leq M_{C}\leq 1)$, is defined by

$M_{c}:=\operatorname{min-
max}\left(\operatorname{ReLU}\left(ZW^{\top}\right)\right).$ (2)

It encodes the semantic uncertainty for each patch driven by CAM scores
$ZW^{\top}$.

Next, we identify the uncertain regions based on the normalized CAM and mask
the uncertain patches, following an adaptive masking strategy. Features in
non-reliable regions are considered as uncertain features. However, some
reliable regions can be wrongly labeled, and their corresponding features can
also be uncertain. To remedy this, we propose an adaptive masking strategy,
resulting in a soft masking vector $M_{s}\in\mathbb{R}^{HW}$ with each element
given as

$M_{s}^{(i)}=\left\\{\begin{array}[]{lr}u_{i}+1,&\text{if
}\beta_{l}<\max\left(M_{c}^{(i,:)}\right)<\beta_{h},\\\
u_{i},&\text{otherwise},\end{array}\right.$ (3)

where $u_{i}\sim\text{U}(0,1)$ draws from a standard uniform distribution and
enables a stochastic selection process. The above use of two background
thresholds $0<\beta_{l}<\beta_{h}<1$ for dividing patches into reliable and
non-reliable ones is inspired by Zhang et al. (2020) and Ru et al. (2022),
which suggests an intermediate score to be a sign of uncertainty. As a result,
elements in $M_{s}$ with larger values suggest uncertain patches.

We use a masking ratio $0<r<1$ to control the amount of selected uncertain
patches, and defines the following binary selection mask
$M_{b}\in\mathbb{R}^{HW}$ with each element as

$M_{b}^{(i)}=\left\\{\begin{array}[]{lr}1,&\text{if
}i\in\operatorname*{arg\,max_{top(k)}}(M_{s}),k:=\lfloor HW*r\rfloor,\\\
0,&\text{otherwise},\end{array}\right.$ (4)

where $\lfloor\cdot\rfloor$ denotes the floor function. The selected uncertain
patches, flagged by 1 in $M_{b}$, correspond to those top-$k$ large-valued
elements in $M_{s}$. Our masking strategy is designed to relax the hard
foreground-background thresholds by the masking ratio $r$. When more patches
are flagged as uncertain by
$\beta_{l}<\max\left(M_{c}^{(i,:)}\right)<\beta_{h}$, the selection is
randomly conducted within them through $u_{i}$. When less uncertain patches
are flagged, part of confident patches are also selected. The original
features of the selected tokens to mask are replaced by learnable parameters
with the same feature dimension.

Figure 3: Illustration of the aggregation module. This module is composed of
several aggregation blocks, where each block alternates in turn a cross-
attention layer and a feed-forward layer. The cross-attention layer computes
attention between a class token and a sequence of unmasked patch tokens.

### Certain Feature Aggregation

We design an attentive aggregation module to compress the embeddings of a
sequence of $N=HW$ patch tokens, stored in $\hat{Z}\in\mathbb{R}^{N\times D}$,
into one class token embedding $\bar{Z}\in\mathbb{R}^{1\times D}$. As shown in
Figure 3, this module contains several aggregation blocks, where each block
contains a masked cross-attention (MCA) layer and a feed-forward (FF) layer,
given as

$\displaystyle\bar{Z}^{(l)}_{(o)}$
$\displaystyle=\bar{Z}^{(l)}+\text{MCA}\left(\eta\left(\left[\bar{Z}^{(l)};\hat{Z}^{(l)}\right]\right)\right),$
(5) $\displaystyle\bar{Z}^{(l+1)}$
$\displaystyle=\bar{Z}^{(l)}_{(o)}+\text{FF}\left(\eta\left(\bar{Z}^{(l)}_{(o)}\right)\right),$

where $l$ denotes the layer index and $\eta$ is the LayerNorm (Ba, Kiros, and
Hinton 2016).

MCA is analogous to self-attention (Vaswani et al. 2017), except that it
computes attention between the class token and a set of unmasked patch tokens.
We parameterize MCA with projection weights
$W_{Q},W_{K},W_{V},W_{O}\in\mathbb{R}^{D\times D}$, and calculate the queries
$Q\in\mathbb{R}^{1\times D}$, keys $K\in\mathbb{R}^{N\times D}$ and values
$V\in\mathbb{R}^{N\times D}$ by projection, so that

$Q=\eta\left(\bar{Z}\right)W_{Q}^{\top},K=\eta\left(\hat{Z}\right)W_{K}^{\top},V=\eta\left(\hat{Z}\right)W_{V}^{\top}.$
(6)

Note that queries are derived from the class token, while keys and values are
calculated on patch tokens. The masked cross-attention
$A\in\mathbb{R}^{1\times N}$ is then formulated as

$A=\sigma\left(\frac{\left(1-M_{b}\right)\left(QK^{\top}\right)}{\sqrt{D}}\right).$
(7)

The output of MCA is computed as a weighted sum of values, i.e.,
$\left(AV\right)W_{O}^{\top}$.

### Feature Self-reinforcement

We adopt self-distillation (Caron et al. 2021; Zhou et al. 2021; Oquab et al.
2023) to improve the model training for feature reinforcement. As explained
earlier, given two distorted views of the same image, we compute one student
output $\hat{P}_{2}$ and two teacher outputs $P_{1}$ and $P_{2}$, where their
first element stores the aggregated token information, while the remaining the
individual token content. We propose a self-reinforcement loss
$\mathcal{L}_{u}$ for the uncertain tokens, as the cross-entropy loss between
each student’s patch token and its corresponding teacher’s patch token:

$\mathcal{L}_{u}=-\sum_{i=2}^{1+N}M_{b}^{(i)}P_{2}^{(i)}\log\hat{P}_{2}^{(i)},$
(8)

where $M_{b}$ is the mask in Eq. (4) to help select masked patch tokens. We
also conduct self-reinforcement for the confident tokens, formulated as the
cross-entropy loss on the two aggregated class tokens of the two views, given
as

$\mathcal{L}_{c}=-P_{1}^{(1)}\log\hat{P}_{2}^{(1)}.$ (9)

Following a common practice, we adopt the multi-label soft margin loss
$\mathcal{L}_{cls}$ for classification, the pixel-wise cross-entropy loss
$\mathcal{L}_{seg}$ for segmentation, and the cosine similarity loss
$\mathcal{L}_{aff}$ for affinity regularization. Denote the weighting factors
as $\\{\lambda_{i}\\}_{i=1}^{5}$, the overall training objective is

$\mathcal{L}=\lambda_{1}\mathcal{L}_{cls}+\lambda_{2}\mathcal{L}_{aff}+\lambda_{3}\mathcal{L}_{seg}+\lambda_{4}\mathcal{L}_{u}+\lambda_{5}\mathcal{L}_{c}.$
(10)

It consolidates classification, segmentation and feature self-reinforcement
within a single-stage framework.

## Experiments

Method | Sup. | Net. | Val | Test
---|---|---|---|---
Multi-stage WSSS methods.
RIB (Lee et al. 2021a) | $\mathcal{I}+\mathcal{S}$ | RN101 | 70.2 | 70.0
EDAM (Wu et al. 2021) | $\mathcal{I}+\mathcal{S}$ | RN101 | 70.9 | 70.6
EPS (Lee et al. 2021b) | $\mathcal{I}+\mathcal{S}$ | RN101 | 71.0 | 71.8
SANCE (Li, Fan, and Zhang 2022) | $\mathcal{I}+\mathcal{S}$ | RN101 | 72.0 | 72.9
L2G (Jiang et al. 2022) | $\mathcal{I}+\mathcal{S}$ | RN101 | 72.1 | 71.7
RCA (Zhou et al. 2022) | $\mathcal{I}+\mathcal{S}$ | RN38 | 72.2 | 72.8
SEAM (Wang et al. 2020) | $\mathcal{I}$ | RN38 | 64.5 | 65.7
BES (Chen et al. 2020) | $\mathcal{I}$ | RN101 | 65.7 | 66.6
CPN (Zhang et al. 2021) | $\mathcal{I}$ | RN38 | 67.8 | 68.5
CDA (Su et al. 2021) | $\mathcal{I}$ | RN38 | 66.1 | 66.8
ReCAM (Chen et al. 2022) | $\mathcal{I}$ | RN101 | 68.5 | 68.4
URN (Li et al. 2022b) | $\mathcal{I}$ | RN101 | 69.5 | 69.7
ESOL (Li et al. 2022a) | $\mathcal{I}$ | RN101 | 69.9 | 69.3
$\dagger$ViT-PCM (Rossetti et al. 2022) | $\mathcal{I}$ | RN101 | 70.3 | 70.9
$\dagger$MCTformer (Xu et al. 2022) | $\mathcal{I}$ | RN38 | 71.9 | 71.6
$\dagger$OCR (Cheng et al. 2023) | $\mathcal{I}$ | RN38 | 72.7 | 72.0
$\dagger$BECO (Rong et al. 2023) | $\mathcal{I}$ | MiT-B2 | 73.7 | 73.5
$\dagger$MCTformer+ (Xu et al. 2023) | $\mathcal{I}$ | RN38 | 74.0 | 73.6
Single-stage WSSS methods.
RRM (Zhang et al. 2020) | $\mathcal{I}$ | RN38 | 62.6 | 62.9
1Stage (Araslanov and Roth 2020) | $\mathcal{I}$ | RN38 | 62.7 | 64.3
$\dagger$AFA (Ru et al. 2022) | $\mathcal{I}$ | MiT-B1 | 66.0 | 66.3
$\dagger$ToCo (Ru et al. 2023) | $\mathcal{I}$ | ViT-B | 71.1 | 72.2
$\dagger$Ours | $\mathcal{I}$ | ViT-B | 75.7 | 75.0
Table 1: Performance comparison of semantic segmentation on PASCAL VOC 2012 in terms of mIoU(%). Sup. denotes the supervision type. $\mathcal{I}$: image-level class labels. $\mathcal{S}$: off-the-shelf saliency maps. Net. denotes the segmentation network for multi-stage methods or the backbone for single-stage methods. RN38: Wide ResNet38 (Wu, Shen, and Van Den Hengel 2019), RN101: ResNet101 (He et al. 2016), MiT: Mix Transformer (Xie et al. 2021). $\dagger$ flags transformer based classification network or backbone. Method | Sup. | Net. | Val
---|---|---|---
Multi-stage WSSS methods.
EPS (Lee et al. 2021b) | $\mathcal{I}+\mathcal{S}$ | RN101 | 35.7
RIB (Lee et al. 2021a) | $\mathcal{I}+\mathcal{S}$ | RN101 | 43.8
L2G (Jiang et al. 2022) | $\mathcal{I}+\mathcal{S}$ | RN101 | 44.2
CDA (Su et al. 2021) | $\mathcal{I}$ | RN38 | 33.2
URN (Li et al. 2022b) | $\mathcal{I}$ | RN101 | 40.7
ESOL (Li et al. 2022a) | $\mathcal{I}$ | RN101 | 42.6
$\dagger$MCTformer (Xu et al. 2022) | $\mathcal{I}$ | RN38 | 42.0
$\dagger$ViT-PCM (Rossetti et al. 2022) | $\mathcal{I}$ | RN101 | 45.0
$\dagger$OCR (Cheng et al. 2023) | $\mathcal{I}$ | RN38 | 42.5
BECO (Rong et al. 2023) | $\mathcal{I}$ | RN101 | 45.1
$\dagger$MCTformer+ (Xu et al. 2023) | $\mathcal{I}$ | RN38 | 45.2
Single-stage WSSS methods.
$\dagger$AFA (Ru et al. 2022) | $\mathcal{I}$ | MiT-B1 | 38.9
$\dagger$ToCo (Ru et al. 2023) | $\mathcal{I}$ | ViT-B | 42.3
$\dagger$Ours | $\mathcal{I}$ | ViT-B | 45.4
Table 2: Performance comparison of semantic segmentation on MS COCO 2014 in
terms of mIoU(%). We use the same notations as in Table 1.

### Experimental Settings

#### Datasets

We evaluate our method on two benchmarks: PASCAL VOC 2012 (Everingham et al.
2010) and MS COCO 2014 (Lin et al. 2014). PASCAL VOC contains 20 object
classes and one background class. Following the common practice of previous
works (Zhang et al. 2020; Araslanov and Roth 2020; Ru et al. 2022, 2023), it
is augmented with data from the SBD dataset (Hariharan et al. 2011), resulting
in $10,582$, $1,449$ and $1,456$ images for training, validation and testing,
respectively. MS COCO contains 80 object classes and one background class. It
has $82,081$ images for training and $40,137$ images for validation. Note that
we only adopt image-level labels during the training phase. We report mean
Intersection-over-Union (mIoU) as the evaluation metric.

#### Implementation Details

We adopt ViT-B (Dosovitskiy et al. 2020) pretrained on ImageNet (Deng et al.
2009) as the transformer encoder. The convolutional decoder refers to DeepLab-
LargeFOV (Chen et al. 2017). We use two aggregation blocks in the aggregation
module. The projector comprises a 3-layer perceptron and a weight-normalized
fully connected layer (Caron et al. 2021). Parameters in the aggregation
module and the projector are randomly initialized. We use a light data
augmentation: random resized cropping to $448\times 448$ with the scale
$[0.32,1.0]$ and the ratio $[3/4,4/3]$, random horizontal flipping, and random
color jittering. The student network is optimized with AdamW (Loshchilov and
Hutter 2017). The base learning rate is warmed up to $6e-5$ at the first 1,500
iterations and decayed with a cosine schedule. The weighting factors
$(\lambda_{1},\lambda_{2},\lambda_{3},\lambda_{4},\lambda_{5})$ are set to
$(1.0,0.2,0.1,0.1,0.1)$. The teacher network requires no gradient and is
updated with the EMA momentum. Experimentally, we find that synchronizing the
teacher encoder with the student (i.e., momentum is $0.0$) works better. The
momentum for the teacher projector is $0.996$ and increases to $1.0$ with a
cosine schedule during training. We embrace the centering and sharpening
technique suggested in (Caron et al. 2021) to avoid collapsed solutions. The
masking ratio $r$ is $0.4$ for adaptive uncertain feature selection. The
background scores $(\beta_{l},\beta_{h})$ introduced to determine uncertain
regions are $(0.2,0.7)$. Training iterations are 20,000 for PASCAL VOC 2012
and 80,000 for MS COCO 2014. We use multi-scale testing and dense CRF (Chen et
al. 2014) at test time following (Ru et al. 2022, 2023).

### Comparison with State-of-the-arts

##### PASCAL VOC 2012

Table 1 shows comparison results of our proposed Feature Self-Reinforcement
(FSR) with other state-of-the-art methods on PASCAL VOC 2012. As can be seen
from this table, FSR significantly outperforms other single-stage approaches,
achieving $75.7\%$ and $75.0\%$ mIoU on the validation and test sets,
respectively. It is noticeable that our method achieves even higher mIoU than
several sophisticated multi-stage methods, e.g., FSR surpasses BECO (Rong et
al. 2023) by margins of $2.0\%$ and $1.5\%$. Compared with multi-stage methods
using both image-level labels and off-the-shelf saliency maps, e.g., L2G
(Jiang et al. 2022) and RCA (Zhou et al. 2022), our method still achieves
superior performance. We assume although saliency maps are effective in
providing additional background clues, our method can strengthen both
confident regions (mostly the main body of objects or the background) and
uncertain regions (mostly object boundaries), so that semantically distinct
objects can be better differentiated. Moreover, it shows that recent methods
with transformer-based networks (with $\dagger$) generally outperform those
with convolutional networks (without $\dagger)$. Nevertheless, due to the
difficulty of end-to-end optimization, single-stage transformer-based methods
(e.g., ToCo reports $71.1\%$ and $72.2\%$) can only achieve comparable
performance with multi-stage ones (e.g., BECO reports $73.7\%$ and $73.5\%$).
Our method proves the efficacy of transformer-based single-stage training by
attaining even better results.

##### MS COCO 2014

Table 2 gives comparison results of semantic segmentation on a more
challenging benchmark MS COCO 2014. We achieve $45.5\%$ mIoU on the validation
set, which outperforms previous single-stage solutions and is slightly better
than multi-stage MCTformer+ (Xu et al. 2023) by $0.2\%$. This further
demonstrates the superiority of our proposed method.

Figure 4: Visualization results of CAMs and predicted segmentation labels with
SOTA single-stage frameworks (i.e., AFA and ToCo). (left) Ground truth.
(middle) Comparison results of CAMs. (right) Comparison results of predicted
segmentation labels.

##### Visualization Results

In Figure 4, we visualize CAMs derived from the classifier and semantic
segmentation labels predicted by the decoder of three single-stage methods,
i.e., AFA (Ru et al. 2022), ToCo (Ru et al. 2023) and our proposed FSR.
Compared with AFA, ToCo and FSR can generate more integral and deterministic
CAMs. For instance, the wheels of “motorbike” are mildly activated by AFA
while strongly confirmed by ToCo and FSR. This proves the effectiveness of FSR
for uncertain features. However, AFA only focuses on boosting uncertain
features, whereas our method enhances both uncertain and certain ones. For
instance, AFA mistakes “drawer” as “chair”, while FSR successfully recognizes
the different semantics. This shows the importance of FSR for seemingly
certain features.

### Ablation Studies

In this section, we present extensive ablation studies to verify the
effectiveness of our proposed FSR. We report segmentation performance of
pseudo labels (Pseu.) derived from CAMs as well as predicted labels (Pred.)
generated by the decoder. All results are evaluated on PASCAL VOC 2012 val
set. Dense CRF is not applied with ablations.

| Edge | CAM | CAM
---|---|---|---
mask ratio | (strict) | (strict) | 0.1 | 0.2 | 0.3 | 0.4 | 0.5
Pseu. label results (%)
random | - | - | 73.1 | 73.6 | 74.1 | 74.2 | 73.2
uncertain | 73.3 | 74.0 | 74.1 | 74.2 | 73.9 | 74.4 | 73.7
Pred. label results (%)
random | - | - | 71.7 | 72.3 | 71.3 | 72.3 | 71.2
uncertain | 71.6 | 71.8 | 72.2 | 72.3 | 72.0 | 72.5 | 72.1

Table 3: Ablation results of uncertain feature selection methods. “random” means random masking, “uncertain” means our adaptive masking strategy that gives priority to masking uncertain regions.  Masking | unc.FSR | cer.FSR | Pseu. (%) | Pred. (%)
---|---|---|---|---
- |  |  | 71.1 | 67.9
CAM | ✓ |  | 74.4${}_{\textbf{{\color[rgb]{1,0,0}+3.3}}}$ | 72.5${}_{\textbf{{\color[rgb]{1,0,0}+4.6}}}$
| ✓(GAP) | 72.3${}_{\textbf{{\color[rgb]{1,0,0}+1.2}}}$ | 70.9${}_{\textbf{{\color[rgb]{1,0,0}+3.0}}}$
| ✓(GMP) | 71.8${}_{\textbf{{\color[rgb]{1,0,0}+0.7}}}$ | 70.0${}_{\textbf{{\color[rgb]{1,0,0}+2.1}}}$
| ✓(MCA) | 75.2${{}_{\textbf{{\color[rgb]{1,0,0}+4.1}}}}$ | 73.3${}_{\textbf{{\color[rgb]{1,0,0}+5.4}}}$
✓ | ✓(MCA) | 75.7${}_{\textbf{{\color[rgb]{1,0,0}+4.6}}}$ | 73.6${}_{\textbf{{\color[rgb]{1,0,0}+5.7}}}$
Table 4: Ablation results of unc.FSR and cer.FSR. “GAP”, “GMP”, and “MCA” are
aggregation methods of cer.FSR.

#### Analysis of Uncertain Feature Selection

In Table 3, we compare two _strict_ selection methods for uncertain features:
edge-based selection and CAM-based selection. For edge-based selection, we
choose the conventional Canny edge detector to extract edges in an image and
generate exact masks of these edges. Activation thresholds for CAM-based
selection are $(0.2,0.7)$. CAM-based selection is marginally better than edge-
based selection; the improvement continues when CAM-based selection is
relaxed, i.e., uncertain features are not strictly but preferentially masked.
Empirically, we find that $r=0.4$ gives the best result. In addition,
uncertain feature masking achieves higher performance than random feature
masking in most cases, showing it is important to reinforce uncertain features
for semantics clarification.

Figure 5: FSR optimizes the boundary regions (e.g., dashed red box) through
the adaptive masking of regions characterized by uncertainty and the
integration of unc.FSR and cer.FSR.

#### Analysis of Feature Self-reinforcement

Table 4 shows the ablation results of FSR on uncertain regions (unc.FSR) and
on certain regions (cer.FSR). The masking ratio is set to $0.4$ for
comparison. It demonstrates the advancement of unc.FSR by achieving $74.4\%$
(compared to $71.1\%$) on pseudo labels and $72.5\%$ (compared to $67.9\%$) on
predicted labels. This proves that reinforcing uncertain features, which
mainly contain ambiguous object boundaries and misclassified categories, is
fairly effective. When combining unc.FSR with cer.FSR, the quality of pseudo
labels can be further improved, from $74.4\%$ to $75.7\%$; predicted labels
directly supervised by pseudo labels are promoted as well, from $72.5\%$ to
$73.6\%$. This indicates that reinforcing confident features is complementary
to unc.FSR with enhanced global understanding. We showcase examples of our FSR
strategy and its effect on object boundaries in Figure 5.

Figure 6: Average attention entropy of different attention heads (dots) across
different layers.

##### (a) Analysis of unc.FSR

To gain a deep understanding of unc.FSR, we investigate the training process
by analyzing the attention mechanism. Specifically, we compute average
attention entropy (Attanasio et al. 2022) for each attention head across
transformer layers. As shown in Figure 6, the entropy at shallow layers (e.g.,
layer 0 to 6) holds similar without unc.FSR; however, it becomes higher and
tighter at deep layers (e.g., layer 7 to 11) when unc.FSR is applied. A large
entropy for a specific token indicates that a broad context contributes to
this token, while a small entropy tells the opposite. From this point of view,
we assume that unc.FSR benefits semantic segmentation by improving the degree
of contextualization at deep layers.

Figure 7: Class-to-patch attention maps derived from the aggregation module.
Class labels are displayed below.

##### (b) Analysis of cer.FSR

We compare our attentive aggregation of certain features (MCA) with two
conventional methods: Global Average Pooling (GAP) and Global Maximum Pooling
(GMP). GAP assigns an equal weight to each unmasked patch token, while GMP
picks up the dominant one along each dimension. Table 4 shows that GAP
performs better than GMP, as GMP tends to intensify the most discriminative
features, which may have an adverse effect in recognizing an integral object.
It is noticeable that MCA outperforms GAP by a large margin, indicating an
attentive weighting mechanism is superior to average weighting. We visualize
class-to-patch attention maps in Figure 7, which illustrates that the class
token can adaptively learn to pay attention to object regions. Note that the
class token is not directly supervised by classification in our design; it
interacts with unmasked patch tokens and learns to summarize effective
information from them.

| Ours | +GaussianBlur | +Solarization | AutoAugment
---|---|---|---|---
Pseu. (%) | 75.7 |  75.9 $\pm$ 0.05 | 75.3 $\pm$ 0.12 | 74.8 $\pm$ 0.09
Pred. (%) | 73.6 |  73.6 $\pm$ 0.02 | 73.2 $\pm$ 0.06 | 72.8 $\pm$ 0.04

Table 5: 10-trial experimental results of data augmentations. “Ours” is our
default data augmentation setting.

#### Data Augmentation

We present comparison results with other data augmentations in Table 5, which
reveals that data augmentations have limited impacts on the performance. For
example, the performances display variations within the anticipated range when
incorporating GaussianBlur or Solarization. Even when we substitute the data
augmentation with the robust AutoAugmentation (Cubuk et al. 2018), the results
witness a slight decline as a strong augmentation may interfere with the
segmentation objective.

## Conclusion

In this work, we propose to estimate boundaries with the guidance of semantic
uncertainty identified by CAM. To achieve this, we design an activation-based
masking strategy and seek to recover local information with self-distilled
knowledge. We further introduce a self-distillation method to reinforce
semantic consistency with another augmented view. We integrate our method into
the single-stage WSSS framework and validate its effectiveness on PASCAL VOC
2012 and MS COCO 2014 benchmarks.

## Acknowledgement

This work was supported by the National Natural Science Foundation of China
under Grant No. (62106235, 62202015, 62376206, 62003256), and in part by the
Exploratory Research Project of Zhejiang Lab under Grant 2022PG0AN01.

## References

  * Ahn, Cho, and Kwak (2019) Ahn, J.; Cho, S.; and Kwak, S. 2019.  Weakly supervised learning of instance segmentation with inter-pixel relations.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2209–2218.
  * Araslanov and Roth (2020) Araslanov, N.; and Roth, S. 2020.  Single-stage semantic segmentation from image labels.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 4253–4262.
  * Attanasio et al. (2022) Attanasio, G.; Nozza, D.; Hovy, D.; and Baralis, E. 2022.  Entropy-based attention regularization frees unintended bias mitigation from lists.  _arXiv preprint arXiv:2203.09192_.
  * Ba, Kiros, and Hinton (2016) Ba, J. L.; Kiros, J. R.; and Hinton, G. E. 2016.  Layer normalization.  _arXiv preprint arXiv:1607.06450_.
  * Bearman et al. (2016) Bearman, A.; Russakovsky, O.; Ferrari, V.; and Fei-Fei, L. 2016.  What’s the point: Semantic segmentation with point supervision.  In _European conference on computer vision_ , 549–565. Springer.
  * Caron et al. (2021) Caron, M.; Touvron, H.; Misra, I.; Jégou, H.; Mairal, J.; Bojanowski, P.; and Joulin, A. 2021.  Emerging properties in self-supervised vision transformers.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 9650–9660.
  * Chen et al. (2020) Chen, L.; Wu, W.; Fu, C.; Han, X.; and Zhang, Y. 2020.  Weakly supervised semantic segmentation with boundary exploration.  In _Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XXVI 16_ , 347–362. Springer.
  * Chen et al. (2014) Chen, L.-C.; Papandreou, G.; Kokkinos, I.; Murphy, K.; and Yuille, A. L. 2014.  Semantic image segmentation with deep convolutional nets and fully connected crfs.  _arXiv preprint arXiv:1412.7062_.
  * Chen et al. (2017) Chen, L.-C.; Papandreou, G.; Kokkinos, I.; Murphy, K.; and Yuille, A. L. 2017.  Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs.  _IEEE transactions on pattern analysis and machine intelligence_ , 40(4): 834–848.
  * Chen et al. (2022) Chen, Z.; Wang, T.; Wu, X.; Hua, X.-S.; Zhang, H.; and Sun, Q. 2022.  Class re-activation maps for weakly-supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 969–978.
  * Cheng et al. (2023) Cheng, Z.; Qiao, P.; Li, K.; Li, S.; Wei, P.; Ji, X.; Yuan, L.; Liu, C.; and Chen, J. 2023.  Out-of-candidate rectification for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 23673–23684.
  * Cubuk et al. (2018) Cubuk, E. D.; Zoph, B.; Mane, D.; Vasudevan, V.; and Le, Q. V. 2018.  Autoaugment: Learning augmentation policies from data.  _arXiv preprint arXiv:1805.09501_.
  * Dai, He, and Sun (2015) Dai, J.; He, K.; and Sun, J. 2015.  Boxsup: Exploiting bounding boxes to supervise convolutional networks for semantic segmentation.  In _Proceedings of the IEEE international conference on computer vision_ , 1635–1643.
  * Deng et al. (2009) Deng, J.; Dong, W.; Socher, R.; Li, L.-J.; Li, K.; and Fei-Fei, L. 2009.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , 248–255. Ieee.
  * Dosovitskiy et al. (2020) Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.; et al. 2020\.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_.
  * Everingham et al. (2010) Everingham, M.; Van Gool, L.; Williams, C. K.; Winn, J.; and Zisserman, A. 2010\.  The pascal visual object classes (voc) challenge.  _International journal of computer vision_ , 88: 303–338.
  * Gao et al. (2021) Gao, W.; Wan, F.; Pan, X.; Peng, Z.; Tian, Q.; Han, Z.; Zhou, B.; and Ye, Q. 2021\.  Ts-cam: Token semantic coupled attention map for weakly supervised object localization.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2886–2895.
  * Gong et al. (2021) Gong, C.; Wang, D.; Li, M.; Chandra, V.; and Liu, Q. 2021.  Vision transformers with patch diversification.  _arXiv preprint arXiv:2104.12753_.
  * Hariharan et al. (2011) Hariharan, B.; Arbeláez, P.; Bourdev, L.; Maji, S.; and Malik, J. 2011.  Semantic contours from inverse detectors.  In _2011 international conference on computer vision_ , 991–998. IEEE.
  * He et al. (2020) He, K.; Fan, H.; Wu, Y.; Xie, S.; and Girshick, R. 2020.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 9729–9738.
  * He et al. (2016) He, K.; Zhang, X.; Ren, S.; and Sun, J. 2016.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 770–778.
  * Hinton, Vinyals, and Dean (2015) Hinton, G.; Vinyals, O.; and Dean, J. 2015.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.02531_.
  * Hou et al. (2018) Hou, Q.; Jiang, P.; Wei, Y.; and Cheng, M.-M. 2018.  Self-erasing network for integral object attention.  _Advances in Neural Information Processing Systems_ , 31.
  * Jiang et al. (2022) Jiang, P.-T.; Yang, Y.; Hou, Q.; and Wei, Y. 2022.  L2g: A simple local-to-global knowledge transfer framework for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 16886–16896.
  * Lee et al. (2021a) Lee, J.; Choi, J.; Mok, J.; and Yoon, S. 2021a.  Reducing information bottleneck for weakly supervised semantic segmentation.  _Advances in Neural Information Processing Systems_ , 34: 27408–27421.
  * Lee et al. (2019) Lee, J.; Kim, E.; Lee, S.; Lee, J.; and Yoon, S. 2019.  Ficklenet: Weakly and semi-supervised semantic image segmentation using stochastic inference.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 5267–5276.
  * Lee et al. (2021b) Lee, S.; Lee, M.; Lee, J.; and Shim, H. 2021b.  Railroad is not a train: Saliency as pseudo-pixel supervision for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 5495–5505.
  * Li, Fan, and Zhang (2022) Li, J.; Fan, J.; and Zhang, Z. 2022.  Towards noiseless object contours for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 16856–16865.
  * Li et al. (2022a) Li, J.; Jie, Z.; Wang, X.; Wei, X.; and Ma, L. 2022a.  Expansion and shrinkage of localization for weakly-supervised semantic segmentation.  _Advances in Neural Information Processing Systems_ , 35: 16037–16051.
  * Li et al. (2022b) Li, Y.; Duan, Y.; Kuang, Z.; Chen, Y.; Zhang, W.; and Li, X. 2022b.  Uncertainty estimation via response scaling for pseudo-mask noise mitigation in weakly-supervised semantic segmentation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, 1447–1455.
  * Lin et al. (2016) Lin, D.; Dai, J.; Jia, J.; He, K.; and Sun, J. 2016.  Scribblesup: Scribble-supervised convolutional networks for semantic segmentation.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 3159–3167.
  * Lin et al. (2014) Lin, T.-Y.; Maire, M.; Belongie, S.; Hays, J.; Perona, P.; Ramanan, D.; Dollár, P.; and Zitnick, C. L. 2014.  Microsoft coco: Common objects in context.  In _Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13_ , 740–755. Springer.
  * Lin et al. (2023) Lin, Y.; Chen, M.; Wang, W.; Wu, B.; Li, K.; Lin, B.; Liu, H.; and He, X. 2023.  Clip is also an efficient segmenter: A text-driven approach for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 15305–15314.
  * Loshchilov and Hutter (2017) Loshchilov, I.; and Hutter, F. 2017.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_.
  * Noroozi et al. (2018) Noroozi, M.; Vinjimoor, A.; Favaro, P.; and Pirsiavash, H. 2018.  Boosting self-supervised learning via knowledge transfer.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 9359–9367.
  * Oquab et al. (2023) Oquab, M.; Darcet, T.; Moutakanni, T.; Vo, H.; Szafraniec, M.; Khalidov, V.; Fernandez, P.; Haziza, D.; Massa, F.; El-Nouby, A.; et al. 2023.  Dinov2: Learning robust visual features without supervision.  _arXiv preprint arXiv:2304.07193_.
  * Rong et al. (2023) Rong, S.; Tu, B.; Wang, Z.; and Li, J. 2023.  Boundary-Enhanced Co-Training for Weakly Supervised Semantic Segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 19574–19584.
  * Rossetti et al. (2022) Rossetti, S.; Zappia, D.; Sanzari, M.; Schaerf, M.; and Pirri, F. 2022.  Max pooling with vision transformers reconciles class and shape in weakly supervised semantic segmentation.  In _European Conference on Computer Vision_ , 446–463. Springer.
  * Ru et al. (2022) Ru, L.; Zhan, Y.; Yu, B.; and Du, B. 2022.  Learning affinity from attention: End-to-end weakly-supervised semantic segmentation with transformers.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 16846–16855.
  * Ru et al. (2023) Ru, L.; Zheng, H.; Zhan, Y.; and Du, B. 2023.  Token contrast for weakly-supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 3093–3102.
  * Song et al. (2019) Song, C.; Huang, Y.; Ouyang, W.; and Wang, L. 2019.  Box-driven class-wise region masking and filling rate guided loss for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 3136–3145.
  * Su et al. (2022) Su, H.; Ye, Y.; Hua, W.; Cheng*, L.; and Song, M. 2022.  SASFormer: Transformers for Sparsely Annotated Semantic Segmentation.  In _IEEE International Conference on Multimedia and Expo (ICME), 2023_.
  * Su et al. (2021) Su, Y.; Sun, R.; Lin, G.; and Wu, Q. 2021.  Context decoupling augmentation for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 7004–7014.
  * Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017.  Attention is all you need.  _Advances in neural information processing systems_ , 30.
  * Vernaza and Chandraker (2017) Vernaza, P.; and Chandraker, M. 2017.  Learning random-walk label propagation for weakly-supervised semantic segmentation.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 7158–7166.
  * Wang et al. (2023) Wang, Y.; Cheng*, L.; Duan, M.; Wang, Y.; Feng, Z.; and Kong, S. 2023.  Improving Knowledge Distillation via Regularizing Feature Norm and Direction.  _arXiv preprint arXiv:2305.17007_.
  * Wang et al. (2020) Wang, Y.; Zhang, J.; Kan, M.; Shan, S.; and Chen, X. 2020.  Self-supervised equivariant attention mechanism for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 12275–12284.
  * Wu et al. (2023) Wu, F.; He, J.; Cheng*, L.; Yin, Y.; Hao, Y.; and Huang, G. 2023.  Masked Collaborative Contrast for Weakly Supervised Semantic Segmentation.  _arXiv preprint arXiv:2305.08491_.
  * Wu et al. (2021) Wu, T.; Huang, J.; Gao, G.; Wei, X.; Wei, X.; Luo, X.; and Liu, C. H. 2021.  Embedded discriminative attention mechanism for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 16765–16774.
  * Wu, Shen, and Van Den Hengel (2019) Wu, Z.; Shen, C.; and Van Den Hengel, A. 2019.  Wider or deeper: Revisiting the resnet model for visual recognition.  _Pattern Recognition_ , 90: 119–133.
  * Xie et al. (2021) Xie, E.; Wang, W.; Yu, Z.; Anandkumar, A.; Alvarez, J. M.; and Luo, P. 2021.  SegFormer: Simple and efficient design for semantic segmentation with transformers.  _Advances in Neural Information Processing Systems_ , 34: 12077–12090.
  * Xie et al. (2023) Xie, Z.; Geng, Z.; Hu, J.; Zhang, Z.; Hu, H.; and Cao, Y. 2023.  Revealing the dark secrets of masked image modeling.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 14475–14485.
  * Xu et al. (2023) Xu, L.; Bennamoun, M.; Boussaid, F.; Laga, H.; Ouyang, W.; and Xu, D. 2023.  MCTformer+: Multi-Class Token Transformer for Weakly Supervised Semantic Segmentation.  _arXiv preprint arXiv:2308.03005_.
  * Xu et al. (2022) Xu, L.; Ouyang, W.; Bennamoun, M.; Boussaid, F.; and Xu, D. 2022.  Multi-class token transformer for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 4310–4319.
  * Yoon et al. (2022) Yoon, S.-H.; Kweon, H.; Cho, J.; Kim, S.; and Yoon, K.-J. 2022.  Adversarial erasing framework via triplet with gated pyramid pooling layer for weakly supervised semantic segmentation.  In _European Conference on Computer Vision_ , 326–344. Springer.
  * Zhang et al. (2020) Zhang, B.; Xiao, J.; Wei, Y.; Sun, M.; and Huang, K. 2020.  Reliability does matter: An end-to-end weakly supervised semantic segmentation approach.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, 12765–12772.
  * Zhang et al. (2021) Zhang, F.; Gu, C.; Zhang, C.; and Dai, Y. 2021.  Complementary patch for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 7242–7251.
  * Zhang et al. (2023) Zhang, T.; Xue, M.; Zhang, J.; Zhang, H.; Wang, Y.; Cheng, L.; Song, J.; and Song, M. 2023.  Generalization Matters: Loss Minima Flattening via Parameter Hybridization for Efficient Online Knowledge Distillation.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2023_.
  * Zhou et al. (2016) Zhou, B.; Khosla, A.; Lapedriza, A.; Oliva, A.; and Torralba, A. 2016.  Learning deep features for discriminative localization.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2921–2929.
  * Zhou et al. (2021) Zhou, J.; Wei, C.; Wang, H.; Shen, W.; Xie, C.; Yuille, A.; and Kong, T. 2021.  ibot: Image bert pre-training with online tokenizer.  _arXiv preprint arXiv:2111.07832_.
  * Zhou et al. (2022) Zhou, T.; Zhang, M.; Zhao, F.; and Li, J. 2022.  Regional semantic contrast and aggregation for weakly supervised semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 4299–4309.