$|\mathcal{S}|M\|\mathbf{A}_{n\mathcal{S}}\|_{\infty}+(|\mathcal{S}|M)^{2}\|\mathbf{B}_{n\mathcal{S}}\|_{\infty}$,
which establishes the first equation of the display. The second equation of
the display in the Lemma statement can be deduced similarly.

By taking the change of variables
$\bm{\theta}_{\mathcal{S}}-\bm{\theta}_{0\mathcal{S}}=n^{-1/2}\mathbf{u}_{\mathcal{S}}$,
for $\|\mathbf{u}_{\mathcal{S}}\|_{\infty}\leq C$, we can let $M=n^{-1/2}C$ in
the above derivation, so that the LHS of the third equation (26) is bounded by

$|\mathcal{S}|Cn^{-1/2}\|\mathbf{A}_{n\mathcal{S}}\|_{\infty}+(|\mathcal{S}|Cn^{-1/2})^{2}\|\mathbf{B}_{n\mathcal{S}}\|_{\infty}=o_{p}(n^{-1})+o_{p}(n^{-3/2})$

which yields the final result.

∎

The outcome of Lemma C.5 is that, heuristically, the empirical risk function
and its gradient with nuisance parameters known is approximated using cross-
fitting on compacts of $\bm{\theta}_{\mathcal{S}}$. Furthermore, the rate of
uniform approximation improves to $o_{p}(n^{-1})$ for the first quantity in
any $n^{-1/2}$-neighborhood of $\bm{\theta}_{0\mathcal{S}}$. This proves
useful, as examining the behavior of the estimated risk on these shrinking
neighborhoods will allow us to establish asymptotic properties of our proposed
adaptive lasso. The following Lemma establishes similar results for the
perturbation bootstrap when applied to the squared-error risk.

######  Lemma C.6.

Let $G_{1},\ldots,G_{n}\overset{i.i.d.}{\sim}P_{G}$ satisfying
$\mathbb{E}G_{1}^{2}<\infty$ be drawn independently of the data. Then under
the setup of Lemma C.5,

$\displaystyle\sup_{\bm{\theta}_{\mathcal{S}}\in\Theta_{\mathcal{S}}}\bigg{|}\frac{1}{n}\sum_{k=1}^{K}\sum_{i\in\mathbf{I}_{k}}G_{i}\big{\\{}\bar{L}_{\mathcal{S}}\big{(}\bm{\theta}_{\mathcal{S}};\hat{\bm{\mu}}_{k},\mathbf{O}\big{)}-\bar{L}_{\mathcal{S}}\big{(}\bm{\theta}_{\mathcal{S}};\bm{\mu}_{0},\mathbf{O}\big{)}\big{\\}}\bigg{|}=o_{p^{*}}(n^{-1/2})$
$\displaystyle\sup_{\bm{\theta}_{\mathcal{S}}\in\Theta_{\mathcal{S}}}\bigg{\|}\frac{1}{n}\sum_{k=1}^{K}\sum_{i\in\mathbf{I}_{k}}G_{i}\big{\\{}\nabla
L_{\mathcal{S}}\big{(}\bm{\theta}_{\mathcal{S}};\hat{\bm{\mu}}_{k},\mathbf{O}\big{)}-\nabla
L_{\mathcal{S}}\big{(}\bm{\theta}_{\mathcal{S}};\bm{\mu}_{0},\mathbf{O}\big{)}\big{\\}}\bigg{\|}_{\infty}=o_{p^{*}}(n^{-1/2})$
$\displaystyle\sup_{\|\mathbf{u}_{\mathcal{S}}\|_{\infty}\leq
C}\bigg{|}\frac{1}{n}\sum_{k=1}^{K}\sum_{i\in\mathbf{I}_{k}}G_{i}\big{\\{}\bar{L}_{\mathcal{S}}\big{(}\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}};\hat{\bm{\mu}}_{k},\mathbf{O}\big{)}-\bar{L}_{\mathcal{S}}\big{(}\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}};\bm{\mu}_{0},\mathbf{O}\big{)}\big{\\}}\bigg{|}=o_{p^{*}}(n^{-1}).$

where $o_{p^{*}}(1)$ is to be understood as measuring probability under the
product measure $P_{0}\times P_{G}$.

###### Proof.

Since Lemmas C.1 and C.2 also apply when using perturbations $G$ in the sum,
nearly-identical arguments to those in the proof of Lemma C.5 establish the
desired results. ∎

###  C.3 Proof of Theorem 4.1

###### Proof of Theorem 4.1.

The estimators $\hat{\bm{\theta}}_{\mathcal{S}}$ and
$\tilde{\bm{\theta}}_{\mathcal{S}}$ minimize $\hat{R}_{n\mathcal{S}}(\cdot)$
and $R_{n\mathcal{S}}(\cdot)$, respectively. Both
$\hat{R}_{n\mathcal{S}}(\cdot)$ and $R_{n\mathcal{S}}(\cdot)$ are convex in
their arguments, which satisfy (1)-(2) of Theorem C.3. To examine (3), we make
the change of variables
$\bm{\theta}_{\mathcal{S}}=\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}}$
and study:

$\displaystyle
n\\{R_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}})-R_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}})\\}=$
$\displaystyle\mathbf{u}_{\mathcal{S}}^{\top}\Big{\\{}-2n^{-1/2}\sum_{i=1}^{n}\epsilon_{\mathcal{S}i}\big{(}\mathbf{Z}_{i}-\bm{\mu}_{Z0i}\big{)}(\mathcal{S})\Big{\\}}+\mathbf{u}_{\mathcal{S}}^{\top}\mathbf{H}_{n}(\mathcal{S})\mathbf{u}_{\mathcal{S}}$
$\displaystyle=$
$\displaystyle-2\mathbf{u}_{\mathcal{S}}^{\top}\bar{\mathbf{G}}_{n\mathcal{S}}+\mathbf{u}_{\mathcal{S}}^{\top}\mathbf{H}_{n}(\mathcal{S})\mathbf{u}_{\mathcal{S}}$
$\displaystyle=$
$\displaystyle-2\mathbf{u}_{\mathcal{S}}^{\top}\bar{\mathbf{G}}_{n\mathcal{S}}+\mathbf{u}_{\mathcal{S}}^{\top}\mathbf{H}_{0}(\mathcal{S})\mathbf{u}_{\mathcal{S}}+o_{p}(1),$

where the first line follows from (18) along with the representation of the
gradient and Hessian of $L$, the second results from the substitutions
$\bar{\mathbf{G}}_{n\mathcal{S}}=n^{-1/2}\sum_{i=1}^{n}\epsilon_{i}(\mathbf{Z}_{i}-\bm{\mu}_{Z0i})(\mathcal{S})$
and $\mathbf{H}_{n}$, and the final from
$\|\mathbf{H}_{n}-\mathbf{H}_{0}\|_{\infty}=o_{p}(1)$ via the LLN. This
establishes (3). Since $\bar{\mathbf{G}}_{n\mathcal{S}}$ is a normalized sum
with mean $\bm{0}$ by (20) of Proposition C.4, CLT arguments imply
$\bar{\mathbf{G}}_{n\mathcal{S}}\overset{d}{\longrightarrow}\mathbb{G}_{1}\sim
N(\bm{0},\mathbf{V}_{1\mathcal{S}})$ for

$\mathbf{V}_{1\mathcal{S}}=\mathbb{E}\left[\epsilon_{1\mathcal{S}}^{2}\left\\{\mathbf{Z}_{1}(\mathcal{S})-\bm{\mu}_{Z01}(\mathcal{S})\right\\}^{\otimes
2}\right],$ (33)

and hence (4) follows from $\bar{\mathbf{G}}_{n\mathcal{S}}=O_{p}(1)$.
Finally, (5) follows from the eigenvalue restrictions of Assumption C.3. Then
Theorem C.3 implies
$\|\sqrt{n}(\tilde{\bm{\theta}}_{\mathcal{S}}-\bm{\theta}_{0\mathcal{S}})-\mathbf{H}_{0}^{-1}\bar{\mathbf{G}}_{n\mathcal{S}}\|_{\infty}=o_{p}(1)$.
We may infer from the Continuous Mapping Theorem and Slutsky’s Theorem

$\sqrt{n}(\tilde{\bm{\theta}}_{\mathcal{S}}-\bm{\theta}_{0\mathcal{S}})\overset{d}{\longrightarrow}\mathbb{G}_{2}\sim
N\left(\bm{0},\left\\{\mathbf{H}_{0}(\mathcal{S})\right\\}^{-1}\mathbf{V}_{1\mathcal{S}}\left\\{\mathbf{H}_{0}(\mathcal{S})\right\\}^{-1}\right).$

Next, we use a similar argument to establish the distribution of
$\sqrt{n}(\hat{\bm{\theta}}_{\mathcal{S}}-\bm{\theta}_{0\mathcal{S}})$.
Conditions (1)-(2) of Theorem C.3 again follow from convexity arguments. We
examine the function
$n\\{\hat{R}_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}})-\hat{R}_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}})\\}$
for condition (3) by adding and subtracting the terms
$n\\{R_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}})-R_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}})\\}$.
For each $\mathbf{u}_{\mathcal{S}}\in\mathbb{R}^{|\mathcal{S}|}$,

$n\\{\hat{R}_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}})-\hat{R}_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}})\\}=-2\mathbf{u}_{\mathcal{S}}^{\top}\bar{\mathbf{G}}_{n\mathcal{S}}+\mathbf{u}_{\mathcal{S}}^{\top}\mathbf{H}_{0}(\mathcal{S})\mathbf{u}_{\mathcal{S}}+o_{p}(1)\\\
+\underbrace{n\\{(\hat{R}_{n\mathcal{S}}-R_{n\mathcal{S}})(\bm{\theta}_{0\mathcal{S}}+n^{-1/2}\mathbf{u}_{\mathcal{S}})-(\hat{R}_{n\mathcal{S}}-R_{n\mathcal{S}})(\bm{\theta}_{0\mathcal{S}})\\}}_{o_{p}(1)}.$
(34)

The first line on the RHS represents the component examined in the
known-$\bm{\mu}_{0}$ case while the second line examines the error induced
from estimation. This error term is equivalent to the expression (24) and is
$o_{p}(1)$ by Lemma C.5. This establishes (3); the remaining conditions follow
as before. Then we may conclude by this theorem that
$\|\sqrt{n}(\hat{\bm{\theta}}_{\mathcal{S}}-\bm{\theta}_{0\mathcal{S}})-\mathbf{H}_{0}^{-1}\bar{\mathbf{G}}_{n\mathcal{S}}\|_{\infty}=o_{p}(1)$.
Recalling the previous known-$\bm{\mu}_{0}$ case, this implies both (i)
$\sqrt{n}(\hat{\bm{\theta}}_{\mathcal{S}}-\tilde{\bm{\theta}}_{\mathcal{S}})=o_{p}(1)$
and (ii)
$\sqrt{n}(\hat{\bm{\theta}}_{\mathcal{S}}-\bm{\theta}_{0\mathcal{S}})\overset{d}{\longrightarrow}\mathbb{G}_{2}$.

For examining $\hat{\bm{\alpha}}$, we notice that a risk similar to
$\hat{R}_{n}(\cdot)$ is involved. We examine the suitably normalized risk for
$M_{j}$:

$H_{nj}(\bm{\alpha}_{0}(\\{j\\})+n^{-1/2}u_{j})-H_{nj}(\bm{\alpha}_{0}(\\{j\\}))=-2u_{j}\left\\{n^{-1/2}\sum_{i=1}^{n}\eta_{ij}(D_{i}-\mu_{D0i})\right\\}+u_{j}^{2}\left\\{n^{-1}\sum_{i=1}^{n}(D_{i}-\mu_{D0i})^{2}\right\\}.$

The total risk corresponding to all $p$ variables in $\mathbf{M}$ may be
derived as the sum of these risks. Define the $p\times 1$ vector
$\bar{\mathbf{G}}_{nM}:=\left\\{n^{-1/2}\sum_{i=1}^{n}\bm{\eta}_{i}(D_{i}-\mu_{D0i})\right\\}$
and write

$\displaystyle
H_{n}(\bm{\alpha}_{0}+n^{-1/2}\mathbf{u})-H_{n}(\bm{\alpha}_{0})=-2\mathbf{u}^{\top}\bar{\mathbf{G}}_{nM}+\|\mathbf{u}\|_{2}^{2}\mathbf{H}_{0M}+o_{p}(1).$

From here, a similar argument establishes the result. Notably, the eigenvalue
restrictions on $\mathbf{V}_{2}$ ensure
$\|\bar{\mathbf{G}}_{nM}\|_{\infty}=O_{p}(1)$ and the positivity assumption
ensures $0<\mathbf{H}_{0M}<\infty$.

∎

###  C.4 Proof of Lemma C.8

Next we move on to establish properties of the adaptive lasso estimator. The
main lemma of this section will simplify establishing the properties of
various weight functions under different scenarios. Below, we establish some
notation.

Represent the known-$\bm{\mu}_{0}$ version of (9) by

$\mathcal{L}_{n}(\bm{\theta};\hat{\mathbf{w}},\lambda_{n}):=\frac{1}{n}\sum_{i=1}^{n}{\Big{\\{}Y_{i}-{\mu}_{Y0i}-\big{(}\mathbf{Z}_{i}-{\bm{\mu}}_{Z0i}\big{)}^{\top}\bm{\theta}\Big{\\}}^{2}}+\frac{\lambda_{n}}{n}\sum_{j=1}^{p+1}{\hat{\mathbf{w}}(\\{j\\})|\theta_{j}|},$
(35)

with associated minimizer $\tilde{\bm{\theta}}^{L}$. In the sequel, we assume
the weights $\hat{\mathbf{w}}=(0,\hat{w}_{2},\ldots,\hat{w}_{p+1})$ depend on
pilot estimators $\hat{\bm{\alpha}},\hat{\bm{\beta}}$ which are
$\sqrt{n}$-consistent, as described in Section 3.2. Let us define two sets
$\mathcal{S}$ and $\mathcal{S}^{c}$, depending on the weights
$\hat{\mathbf{w}}$, which satisfy

$\displaystyle\max_{j\in\mathcal{S}}n^{-1/2}\lambda_{n}\hat{\mathbf{w}}(\\{j\\})\overset{p}{\longrightarrow}0$
(36)
$\displaystyle\min_{j\in\mathcal{S}^{c}}n^{-1/2}\lambda_{n}\hat{\mathbf{w}}(\\{j\\})\overset{p}{\longrightarrow}\infty.$

We shall assume that these sets partition $\\{1,\ldots,p+1\\}$. That is,
$\mathcal{S}$ represents the subset of $\mathbf{Z}_{1}$ for which the scaled
variable-specific penalty will become asympotically negligible, and
$\mathcal{S}^{c}$ represents the remaining subset for which the corresponding
penalties diverge. We draw attention to our definition of $\hat{w}_{1}=0$, so
that we always have $1\in\mathcal{S}$.

Let us also define the random matrix
$C_{n\mathcal{S}}=\mathbb{P}_{n}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})^{\otimes
2}$, along with its nonrandom target
$C_{0\mathcal{S}}=\mathbb{E}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})^{\otimes
2}$. Let
$G_{n\mathcal{S}}=\sqrt{n}\mathbb{P}_{n}\big{(}\epsilon_{\mathcal{S}}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})\big{)}$
where

$\displaystyle\epsilon_{\mathcal{S}}=$ $\displaystyle
Y-\mu_{Y0}-\bm{\theta}_{0\mathcal{S}}^{\top}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})$
$\displaystyle=$
$\displaystyle\epsilon+\bm{\theta}_{0}^{\top}(\mathbf{Z}-\bm{\mu}_{Z0})-\bm{\theta}_{0\mathcal{S}}^{\top}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}}).$

The vector $\mathbf{Z}_{\mathcal{S}}$ represents the sub-vector of
$\mathbf{Z}$ corresponding to $\mathcal{S}$, and $\bm{\mu}_{Z0\mathcal{S}}$ is
defined similarly. Recall that
$\bm{\theta}_{0\mathcal{S}}=C_{0\mathcal{S}}^{-1}\mathbb{E}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})(Y-\mu_{Y0})$,
ensuring that
$\mathbb{E}(\epsilon_{\mathcal{S}}|\mathbf{Z}_{\mathcal{S}},\mathbf{X})=0$.
Based on this relationship, we can see that
$G_{n\mathcal{S}}\overset{d}{\longrightarrow}\mathbb{G}_{1\mathcal{S}}\sim
N(\bm{0},V_{1\mathcal{S}})$, where
$V_{1\mathcal{S}}=\mathbb{E}\big{(}\epsilon_{\mathcal{S}}^{2}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})^{\otimes
2}\big{)}$.

Finally, we will state a useful Proposition before stating the Lemma.

######  Proposition C.7.

Let $U_{n}$, $V_{n}$ be real-valued random variables defined on the same
probability space for each $n\geq 1$. Suppose $U_{n}=O_{p}(1)$ and
$V_{n}\overset{p}{\longrightarrow}\infty$. Then $Pr\left(U_{n}\geq
V_{n}\right)\rightarrow 0$.

###### Proof.

By assumption, for any $\delta_{1}>0$, there exists $M_{1}$, $N_{1}$ such that
for all $n\geq N_{1}$, $Pr(U_{n}\geq M_{1})\leq\delta_{1}$, and for any
$\delta_{2}>0$, $0<M_{2}<\infty$, there exists $N_{2}$ such that for all
$n\geq N_{2}$, $Pr(V_{n}\leq M_{2})\leq\delta_{2}$. Decompose the original
probability statement and apply inequalities to obtain

$Pr\Big{(}U_{n}\geq V_{n}\Big{)}=Pr\Big{(}U_{n}\geq V_{n},V_{n}\leq
M_{2}\Big{)}+Pr\Big{(}U_{n}\geq V_{n},V_{n}>M_{2}\Big{)}\\\ \leq
Pr\Big{(}V_{n}\leq M_{2}\Big{)}+Pr\Big{(}U_{n}\geq M_{2},V_{n}>M_{2}\Big{)}\\\
\leq\delta_{2}+\min\Big{(}Pr(U_{n}\geq M_{2}),1-\delta_{2}\Big{)},$

where the last line follows from Frechet’s inequality. Since $M_{2},\
\delta_{2},$ and $\delta_{1}$ are arbitrary, we may choose $\delta_{1}$ and
subsequently set $M_{2}$ such that this last expression is bounded by the
arbitrary $\delta=\delta_{2}+\delta_{1}$ for all $n\geq\max(N_{1},N_{2})$. ∎

######  Lemma C.8.

Assume the mediation model (1) holds, and use the previously defined notation.
Let $\tilde{\bm{\theta}}^{L}$ and $\hat{\bm{\theta}}^{L}$ minimize (35) and
(9), respectively, with weights $\hat{\mathbf{w}}$. Let the sets $\mathcal{S}$
and $\mathcal{S}^{c}$ satisfy the asymptotic conditions of (36). Under the
assumptions of Theorem 4.1,

$\displaystyle Pr\big{\\{}\tilde{\bm{\theta}}^{L}(\\{j\\})\neq 0\text{ for any
}j\in\mathcal{S}^{c}\big{\\}}\longrightarrow 0$ (37) $\displaystyle
Pr\big{\\{}\hat{\bm{\theta}}^{L}(\\{j\\})\neq 0\text{ for any
}j\in\mathcal{S}^{c}\big{\\}}\longrightarrow 0$ (38)

and the coefficients $\hat{\bm{\theta}}^{L}(\mathcal{S})$ corresponding to
$\mathcal{S}$ follow the oracle asymptotic normal distribution

$\displaystyle\big{\|}\sqrt{n}\big{\\{}\hat{\bm{\theta}}^{L}(\mathcal{S})-\bm{\theta}_{0\mathcal{S}}\big{\\}}-\sqrt{n}\big{\\{}\tilde{\bm{\theta}}(\mathcal{S})-\bm{\theta}_{0\mathcal{S}}\big{\\}}\big{\|}_{\infty}=o_{p}(1)$
(39)
$\displaystyle\sqrt{n}\big{\\{}\tilde{\bm{\theta}}(\mathcal{S})-\bm{\theta}_{0\mathcal{S}}\big{\\}}\overset{d}{\longrightarrow}\mathbb{G}_{2\mathcal{S}}\sim
N(\bm{0},C_{0\mathcal{S}}^{-1}V_{1\mathcal{S}}C_{0\mathcal{S}}^{-1}),$

with
$C_{0\mathcal{S}}=\mathbb{E}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})^{\otimes
2}$, and
$V_{1\mathcal{S}}=\mathbb{E}\big{(}\epsilon_{\mathcal{S}}^{2}(\mathbf{Z}_{\mathcal{S}}-\bm{\mu}_{Z0\mathcal{S}})^{\otimes
2}\big{)}$. Finally, we have

$\displaystyle
Pr\big{\\{}\tilde{\bm{\theta}}^{L}(\mathcal{S}^{c})=\hat{\bm{\theta}}^{L}(\mathcal{S}^{c})=\bm{0}\big{\\}}\longrightarrow
1$ (40)
$\displaystyle\|\hat{\bm{\theta}}^{L}-\tilde{\bm{\theta}}^{L}\|_{\infty}=o_{p}(n^{-1/2}).$

###### Proof of Lemma C.8.

We follow the general method of proof adopted in Zou and Zhang (2009) and Das
et al. (2019). First, we will prove that
$\tilde{\bm{\theta}}^{L}(\mathcal{S}^{c})=\hat{\bm{\theta}}^{L}(\mathcal{S}^{c})=\bm{0}$
in probability for large $n$, demonstrating the variable selection consistency
results (37)-(38) from the KKT conditions. Then we will show the oracle normal
distribution result for the nonzero parameters (39). The final statement then
follows.

First, we show variable selection consistency. As discussed in Huang and Zhang
(2012), since the functions $\hat{R}_{n}(\bm{\theta})$ and
$R_{n}(\bm{\theta})$ are convex in $\bm{\theta}$, the KKT conditions are both
necessary and sufficient. These authors also give the form of these
conditions, which we present for the estimator with $\bm{\mu}_{0}$ known.
These KKT conditions for $\tilde{\bm{\theta}}^{L}$, the minimizer of
$\mathcal{L}_{n}$, are given by

$-\mathbb{P}_{n}\big{\\{}\nabla
L(\tilde{\bm{\theta}}^{L};\bm{\mu}_{0},\mathbf{O})\big{\\}}=\mathbf{g},$

where $\mathbb{P}_{n}f=n^{-1}\sum_{i=1}^{n}f(\mathbf{O}_{i})$ is the empirical
measure. We can write the LHS as

$\displaystyle-\mathbb{P}_{n}\big{\\{}\nabla
L(\tilde{\bm{\theta}}^{L};\bm{\mu}_{0},\mathbf{O})\big{\\}}=2\mathbb{P}_{n}\Big{\\{}\big{(}\epsilon-(\tilde{\bm{\theta}}^{L}-\bm{\theta}_{0})^{\top}(\mathbf{Z}-\bm{\mu}_{Z0})\big{)}(\mathbf{Z}-\bm{\mu}_{Z0})\Big{\\}}$
$\displaystyle=2\Big{\\{}n^{-1/2}\bar{\mathbf{G}}_{n}-\mathbf{H}_{n}(\tilde{\bm{\theta}}^{L}-\bm{\theta}_{0})\Big{\\}}$

which results from the representation of the gradient in the previous section,
and where $\bar{\mathbf{G}}_{n}\equiv\bar{\mathbf{G}}_{n\mathcal{S}^{F}}$ as
defined in the proof of Theorem 4.1, and
$\mathbf{H}_{n}=n^{-1}\sum_{i=1}^{n}(\mathbf{Z}_{i}-\bm{\mu}_{Z0i})^{\otimes
2}$ . Of particular interest is the form of the vector $\mathbf{g}$. When
$\tilde{\bm{\theta}}^{L}(\\{j\\})\neq 0$,
$\mathbf{g}(\\{j\\})=n^{-1}\lambda_{n}\hat{\mathbf{w}}(\\{j\\})sgn\\{\tilde{\bm{\theta}}^{L}(\\{j\\})\\}$.
By multiplying $sgn(\tilde{\bm{\theta}}^{L}(\\{j\\}))\sqrt{n}$ through both
sides of KKT equality, we may examine the probability of non-sparsity in
$\mathcal{S}_{0}^{c}$ by comparing bounded variables with unbounded ones. Let
$\bm{e}_{j}$ be a conformable column vector of 0 with 1 in the $j^{th}$
position,
$U_{nj}=\bm{e}_{j}^{\top}\\{\bar{\mathbf{G}}_{n}-\sqrt{n}\mathbf{H}_{n}(\tilde{\bm{\theta}}^{L}-\bm{\theta}_{0})\\}$,
$v_{nj}=n^{-1/2}\lambda_{n}\hat{\mathbf{w}}(\\{j\\})$, and
$\bar{v}_{n}=\min_{j\in\mathcal{S}_{0}^{c}}v_{nj}$. Then,

$\displaystyle Pr\big{\\{}\tilde{\bm{\theta}}^{L}(\\{j\\})\neq 0\text{ for any
}j\in\mathcal{S}_{0}^{c}\big{\\}}\leq$ $\displaystyle
Pr\Big{(}\cup_{j\in\mathcal{S}_{0}^{c}}\left\\{2sgn(\tilde{\bm{\theta}}^{L}(\\{j\\}))U_{nj}\geq
v_{nj}\right\\}\Big{)}$ (41) $\displaystyle\leq$
$\displaystyle\sum_{j\in\mathcal{S}_{0}^{c}}Pr\left\\{2sgn(\tilde{\bm{\theta}}^{L}(\\{j\\}))U_{nj}\geq
v_{nj}\right\\}$ $\displaystyle\leq$
$\displaystyle\sum_{j\in\mathcal{S}_{0}^{c}}Pr\left(2|U_{nj}|\geq
v_{nj}\right)$ $\displaystyle\leq$
$\displaystyle\sum_{j\in\mathcal{S}_{0}^{c}}Pr\left(2U_{n}\geq v_{nj}\right)$
$\displaystyle\leq$
$\displaystyle|\mathcal{S}_{0}^{c}|Pr\left(2U_{n}\geq\bar{v}_{n}\right),$

where the first inequality follows from the necessary KKT equality, the second
inequality follows from subadditivity, the third from the trivial inequality
$A_{n}\leq|A_{n}|$ for any random variable $A_{n}$, the fourth from defining
the stochastically larger random variable
$U_{n}=\|\bar{\mathbf{G}}_{n}\|_{\infty}+\|\sqrt{n}\mathbf{H}_{n}\|_{\infty}\|\tilde{\bm{\theta}}^{L}-\bm{\theta}_{0}\|_{\infty}$,
and the final inequality follows from the definition of $\bar{v}_{n}$. First,
notice that $\|\tilde{\bm{\theta}}^{L}-\bm{\theta}_{0}\|_{\infty}=O_{p}(1)$ by
a similar argument as in the proof of Theorem 1 in Knight and Fu (2000). Next,
notice that under our modeling assumptions,
$\|\bar{\mathbf{G}}_{n}\|_{\infty}$ and $\|\sqrt{n}\mathbf{H}_{n}\|_{\infty}$
are both $O_{p}(1)$. Then $U_{n}=O_{p}(1)$, and since
$\bar{v}_{n}\overset{p}{\longrightarrow}\infty$, we apply Proposition C.7
along with $|\mathcal{S}_{0}^{c}|<\infty$ to arrive at (37).

The result for $\hat{\bm{\theta}}^{L}$, the minimizer of
$\hat{\mathcal{L}}_{n}$ follows similarly. First, note that an identical
argument as before establishes that
$\|\hat{\bm{\theta}}^{L}-\bm{\theta}_{0}\|_{\infty}=O_{p}(1)$. The LHS of the
KKT conditions related to the estimator $\hat{\bm{\theta}}^{L}$ may be written
as

$-\frac{1}{n}\sum_{k=1}^{K}\sum_{i\in\mathbf{I}_{k}}\nabla
L(\hat{\bm{\theta}}^{L};\hat{\bm{\mu}}_{k},\mathbf{O}_{i})$

which by Lemma C.5 is well-approximated by

$-\frac{1}{n}\sum_{k=1}^{K}\sum_{i\in\mathbf{I}_{k}}\nabla
L(\hat{\bm{\theta}}^{L};\bm{\mu}_{0},\mathbf{O}_{i}),$

up to $o_{p}(n^{-1/2})$. This last display is equal to
$2\\{n^{-1/2}\bar{\mathbf{G}}_{n}-(\hat{\bm{\theta}}^{L}-\bm{\theta}_{0})\mathbf{H}_{n}\\}$
which only differs from the known-$\bm{\mu}_{0}$ case by the presence of
$\hat{\bm{\theta}}^{L}$ rather than $\tilde{\bm{\theta}}^{L}$. The arguments
from here follow similarly with
$\hat{U}_{n}=\|\bar{\mathbf{G}}_{n}\|_{\infty}+\|\sqrt{n}\mathbf{H}_{n}\|_{\infty}\|\hat{\bm{\theta}}^{L}-\bm{\theta}_{0}\|_{\infty}+o_{p}(1)$
to establish (38). This completes the proof of selection consistency.

Now we move on to the asymptotic normality. We have shown that with
arbitrarily high probability and for sufficiently large $n$, the solutions to
$\mathcal{L}_{n}$ and $\hat{\mathcal{L}}_{n}$ satisfy
$\tilde{\bm{\theta}}^{L}(\mathcal{S}_{0}^{c})=\bm{0}$ and
$\hat{\bm{\theta}}^{L}(\mathcal{S}_{0}^{c})=\bm{0}$, respectively.
Consequently, $\tilde{\bm{\theta}}(\mathcal{S}_{0})^{L}$ and
$\hat{\bm{\theta}}(\mathcal{S}_{0})^{L}\in\mathbb{R}^{|\mathcal{S}_{0}|}$ are
the respective minimizers, in probability, of

$\displaystyle\mathcal{L}_{n1}(\bm{\theta}_{\mathcal{S}_{0}};\lambda_{n},\hat{\mathbf{w}})=R_{n\mathcal{S}_{0}}(\bm{\theta}_{\mathcal{S}_{0}})+n^{-1}\lambda_{n}\sum_{j\in\mathcal{S}_{0}}\hat{\mathbf{w}}(\\{j\\})|\bm{\theta}_{\mathcal{S}_{0}}(\\{j\\})|$
(42)
$\displaystyle\hat{\mathcal{L}}_{n1}(\bm{\theta}_{\mathcal{S}_{0}};\lambda_{n},\hat{\mathbf{w}})=\hat{R}_{n\mathcal{S}_{0}}(\bm{\theta}_{\mathcal{S}_{0}})+n^{-1}\lambda_{n}\sum_{j\in\mathcal{S}_{0}}\hat{\mathbf{w}}(\\{j\\})|\bm{\theta}_{\mathcal{S}_{0}}(\\{j\\})|,$
(43)

where the submodel risks $R_{n\mathcal{S}}(\bm{\theta}_{\mathcal{S}_{0}})$ and
$\hat{R}_{n\mathcal{S}}(\bm{\theta}_{\mathcal{S}_{0}})$ are defined as before.

From this point, we follow the proof of Theorem 4.1. We present an argument
for $\hat{\bm{\theta}}^{L}_{\mathcal{S}_{0}}$. The argument establishing the
result for $\tilde{\bm{\theta}}^{L}_{\mathcal{S}_{0}}$ is nearly identical.
Note that both (42) and (43) are convex functions of
$\bm{\theta}_{\mathcal{S}_{0}}\in\mathbb{R}^{|\mathcal{S}_{0}|}$, which
ensures conditions (1)-(2) of Theorem C.3. To establish (3), we examine

$\displaystyle
n\\{\hat{\mathcal{L}}_{n1}(\bm{\theta}_{0\mathcal{S}_{0}}+n^{-1/2}\mathbf{u}_{\mathcal{S}_{0}};\hat{\mathbf{w}})-\hat{\mathcal{L}}_{n1}(\bm{\theta}_{0\mathcal{S}_{0}};\hat{\mathbf{w}})\\}=n\left\\{\hat{R}_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}_{0}}+n^{-1/2}\mathbf{u}_{\mathcal{S}_{0}})-\hat{R}_{n\mathcal{S}}(\bm{\theta}_{\mathcal{S}_{0}})\right\\}$
(44)
$\displaystyle+\sum_{j\in\mathcal{S}_{0}}v_{nj}\Big{[}n^{1/2}\big{\\{}|\bm{\theta}_{\mathcal{S}_{0}}(\\{j\\})+n^{-1/2}u_{\mathcal{S}_{0}}(\\{j\\})|-|\bm{\theta}_{\mathcal{S}_{0}}(\\{j\\})|\big{\\}}\Big{]},$
(45)

where we recall $v_{nj}=n^{-1/2}\lambda_{n}\hat{\mathbf{w}}(\\{j\\})$. The RHS
of (44) equals
$\mathbf{u}_{\mathcal{S}_{0}}^{\top}\mathbf{H}_{0}(\mathcal{S}_{0})\mathbf{u}_{\mathcal{S}_{0}}-2\mathbf{u}_{\mathcal{S}_{0}}^{\top}\bar{\mathbf{G}}_{n\mathcal{S}_{0}}+o_{p}(1)$
by (34). The term in square brackets within (45) is upper- and lower-bounded
by $|u_{\mathcal{S}_{0}}(\\{j\\})|$ and $-|u_{\mathcal{S}_{0}}(\\{j\\})|$,
respectively, for each $u_{\mathcal{S}_{0}}(\\{j\\})$. This follows by taking
the absolute value and applying the reverse triangle inequality. Now, since
$v_{nj}\overset{p}{\longrightarrow}0$ for all $j\in\mathcal{S}_{0}$, (45)
converges to zero in probability for each
$\mathbf{u}_{\mathcal{S}_{0}}\in\mathbb{R}^{|\mathcal{S}_{0}|}$. This
establishes condition (3) of Theorem C.3:

$n\\{\hat{\mathcal{L}}_{n1}(\bm{\theta}_{0\mathcal{S}_{0}}+n^{-1/2}\mathbf{u}_{\mathcal{S}_{0}};\hat{\mathbf{w}})-\hat{\mathcal{L}}_{n1}(\bm{\theta}_{0\mathcal{S}_{0}};\hat{\mathbf{w}})\\}=\mathbf{u}_{\mathcal{S}_{0}}^{\top}\mathbf{H}_{0}(\mathcal{S}_{0})\mathbf{u}_{\mathcal{S}_{0}}-2\mathbf{u}_{\mathcal{S}_{0}}^{\top}\bar{\mathbf{G}}_{n\mathcal{S}_{0}}+o_{p}(1).$

Condition (4) follows from
$\bar{\mathbf{G}}_{n\mathcal{S}_{0}}\overset{d}{\longrightarrow}\mathbb{G}_{1\mathcal{S}_{0}}\sim
N(\bm{0},\mathbf{V}_{1\mathcal{S}_{0}})$, where
$\mathbf{V}_{1\mathcal{S}_{0}}$ is defined in (33). The final condition (5)
follows from the eigenvalue restrictions on $\mathbf{H}_{0}(\mathcal{S}_{0})$.
Then we may conclude

$\|\sqrt{n}\\{\hat{\bm{\theta}}(\mathcal{S}_{0})^{L}-\bm{\theta}_{0\mathcal{S}_{0}}\\}-[\mathbf{H}_{0}(\mathcal{S}_{0})]^{-1}\bar{\mathbf{G}}_{n\mathcal{S}_{0}}\|_{\infty}=o_{p}(1).$

We combine this result with the Triangle inequality and Theorem 4.1 to
establish the first part of (39). The second part of (39) follows from the
Continuous Mapping Theorem.

The argument of the preceding paragraph follows similarly in the case where
$\bm{\mu}_{0}$ is known. The main change is made to line (44), instead written
as

$n\\{\mathcal{L}_{n1}(\bm{\theta}_{0\mathcal{S}_{0}}+n^{-1/2}\mathbf{u}_{\mathcal{S}_{0}};\hat{\mathbf{w}})-\mathcal{L}_{n1}(\bm{\theta}_{0\mathcal{S}_{0}};\hat{\mathbf{w}})\\}=n\left\\{{R}_{n\mathcal{S}}(\bm{\theta}_{0\mathcal{S}_{0}}+n^{-1/2}\mathbf{u}_{\mathcal{S}_{0}})-{R}_{n\mathcal{S}}(\bm{\theta}_{\mathcal{S}_{0}})\right\\}+p_{n}(\mathbf{u}_{\mathcal{S}_{0}})$

with $p_{n}(\mathbf{u}_{\mathcal{S}_{0}})$ defined by (45). The remaining
arguments follow as before.

Finally, we show (40). For the first statement, apply the bound $Pr(A\cap
B)\leq Pr(A)+Pr(B)$ for events $A$, $B$ defined in (37) and (38),
respectively. Thus
$\tilde{\bm{\theta}}^{L}_{\mathcal{S}_{0}^{c}}=\hat{\bm{\theta}}^{L}_{\mathcal{S}_{0}^{c}}=\bm{0}$
with probability converging to one. Using the Triangle Inequality argument
above, we have
$\|\sqrt{n}(\tilde{\bm{\theta}}(\mathcal{S}_{0})^{L}-\hat{\bm{\theta}}(\mathcal{S}_{0})^{L})\|_{\infty}=o_{p}(1)$,
which in combination with the previous probability statement completes the
proof of (40). ∎

###  C.5 Proof of Theorem 4.2

Lemma C.8 is flexible and will prove useful in demonstrating asymptotic
results in several interesting cases. These demonstrations essentially require
straightforward checks on the behavior of
$v_{nj}=n^{-1/2}\lambda_{n}\hat{\mathbf{w}}(\\{j\\})$ for $j=2,\ldots,p+1$. As
long as these variables either converge to zero or diverge in probability, we
may construct $\mathcal{S}$ and $\mathcal{S}^{c}$ and satisfy the above
conditions. For the remainder of this section, let $j=1,\ldots,p$ index the
elements of $\bm{\beta}$ rather than components of $\bm{\theta}$. This also
extends to the weights $\hat{\mathbf{w}}(\\{j\\})$, which range from 1 to $p$
rather than to $p+1$ as previously defined. Since we previously defined
$\hat{w}_{1}\equiv 0$, implying $1\in\mathcal{S}_{Z}$ always holds, this re-
indexing does not induce any fundamental changes to the proceeding arguments.

###### Proof of Theorem 4.2.

Recall $v_{nj}=n^{-1/2}\lambda_{n}\hat{\mathbf{w}}(\\{j\\})$. Based on Lemma
C.8, we see that it is sufficient to determine appropriate rates for
$\lambda_{n}$ to ensure the $v_{nj}$ converge to either 0 or $\infty$,
ensuring that each variable either belongs to $\mathcal{S}$ or
$\mathcal{S}^{c}$, respectively.

Suppose
$\hat{\mathbf{w}}(\\{j\\})=|\hat{\bm{\alpha}}(\\{j\\})\hat{\bm{\beta}}(\\{j\\})|^{-\kappa}$
for fixed $\kappa>0$.

  * •

Case 1: $\bm{\alpha}_{0}(\\{j\\})\bm{\beta}_{0}(\\{j\\})\neq 0$. Since
$\hat{\mathbf{w}}(\\{j\\})\overset{p}{\longrightarrow}|\bm{\alpha}_{0}(\\{j\\})\bm{\beta}_{0}(\\{j\\})|^{-\kappa}<\infty$,
it follows that $v_{nj}=n^{-1/2}\lambda_{n}O_{p}(1)$.

  * •

Case 2:  $\bm{\alpha}_{0}(\\{j\\})=0$, $\bm{\beta}_{0}(\\{j\\})\neq 0$.

$n^{-\kappa/2}\hat{\mathbf{w}}(\\{j\\})=\Big{[}\big{|}\sqrt{n}\hat{\bm{\alpha}}(\\{j\\})\big{\\{}\bm{\beta}_{0}(\\{j\\})+o_{p}(1)\big{\\}}\big{|}\Big{]}^{-\kappa}=O_{p}(1)$

which implies $v_{nj}=n^{(\kappa-1)/2}\lambda_{n}O_{p}(1)$.

  * •

Case 3:  $\bm{\alpha}_{0}(\\{j\\})\neq 0$, $\bm{\beta}_{0}(\\{j\\})=0$. By the
symmetry of the weights, a similar argument as above implies
$v_{nj}=n^{(\kappa-1)/2}\lambda_{n}O_{p}(1)$.

  * •

Case 4:  $\bm{\alpha}_{0}(\\{j\\})=0$, $\bm{\beta}_{0}(\\{j\\})=0$. By the
above arguments, we have

$v_{nj}=n^{(2\kappa-1)/2}\lambda_{n}/\Big{\\{}|n\hat{\bm{\alpha}}(\\{j\\})\hat{\bm{\beta}}(\\{j\\})|\Big{\\}}^{\kappa}=n^{(2\kappa-1)/2}\lambda_{n}O_{p}(1).$

Recall
$\mathcal{S}^{*}=\\{j:\bm{\alpha}_{0}(\\{j\\})\bm{\beta}_{0}(\\{j\\})\neq
0\\}$ represent the set of true mediators. To ensure asymptotic normality for
$\mathcal{S}^{*}$, we require $n^{-1/2}\lambda_{n}\rightarrow 0$. To remove
the complete noise variables in Case 4, we require
$n^{(2\kappa-1)/2}\lambda_{n}\rightarrow\infty$. Whether the middle two cases
result in asymptotic normality or elimination with high probability is
controlled by the behavior of $n^{(\kappa-1)/2}\lambda_{n}\rightarrow 0$ or
$\infty$, respectively.

Now, consider the ADP weights
$\hat{\mathbf{w}}(\\{j\\})=|\hat{\bm{\beta}}(\\{j\\})|^{-\kappa}$ for fixed
$\kappa>0$.

  * •

Cases 1&2: $\bm{\beta}_{0}(\\{j\\})\neq 0$. Since
$\hat{\mathbf{w}}(\\{j\\})\overset{p}{\longrightarrow}|\bm{\beta}_{0}(\\{j\\})|^{-\kappa}<\infty$,
it follows that $v_{nj}=n^{-1/2}\lambda_{n}O_{p}(1)$.

  * •

Cases 3&4:  $\bm{\beta}_{0}(\\{j\\})=0$.

$n^{-\kappa/2}\hat{\mathbf{w}}(\\{j\\})=\big{|}\sqrt{n}\hat{\beta}_{j}\big{|}^{-\kappa}=O_{p}(1),$

which implies $v_{nj}=n^{(\kappa-1)/2}\lambda_{n}O_{p}(1)$.

Note that there are only two classes of rates available to distinguish between
the cases. To ensure asymptotic normality for $\mathcal{S}^{*}$, we require
$n^{-1/2}\lambda_{n}\rightarrow 0$. On the other hand, we would require
$n^{(\kappa-1)/2}\lambda_{n}\rightarrow\infty$ to ensure that the Case 4
variables are deselected with high probability. This implies the ADP weights
asymptotically select the mediators $\mathcal{S}^{\ddagger}$ under these rate
assumptions. ∎

###  C.6 Proof of Theorem 4.5

Next we will prove the local asymptotic version of the previous Theorem. We
will use Lemma C.8, as the local parameters do not affect its development.
Recall that the true parameters vary with $n$, with
$\bm{\alpha}_{0n}=\bm{\alpha}_{0}+\mathbf{h}_{1n}$,
$\bm{\beta}_{0n}=\bm{\beta}_{0}+\mathbf{h}_{2n}$, and
$\gamma_{0n}=\gamma_{0}+\mathbf{h}_{3n}$. For $\ell=1,2$, we have that
$\mathbf{h}_{\ell n}(\\{j\\})=0$ if $j\in\tilde{\mathcal{S}}_{\ell}^{c}$. For
the purposes of this proof, we use the convention that $c_{\ell j}=r_{\ell
nj}=0$ if $j\in\tilde{\mathcal{S}}_{\ell}^{c}$, $\ell=1,2$. We will appeal to
decompositions of the pilot estimates:

$\displaystyle\hat{\bm{\alpha}}(\\{j\\})\hat{\bm{\beta}}(\\{j\\})=$
$\displaystyle{\bm{\alpha}_{0n}(\\{j\\})\bm{\beta}_{0n}(\\{j\\})}+(\hat{\bm{\alpha}}(\\{j\\})-\bm{\alpha}_{0n}(\\{j\\}))(\hat{\bm{\beta}}(\\{j\\})-\bm{\beta}_{0n}(\\{j\\}))$
$\displaystyle+\big{\\{}(\hat{\bm{\alpha}}(\\{j\\})-\bm{\alpha}_{0n}(\\{j\\}))\bm{\beta}_{0n}(\\{j\\})+(\hat{\bm{\beta}}(\\{j\\})-\bm{\beta}_{0n}(\\{j\\}))\bm{\alpha}_{0n}(\\{j\\})\big{\\}}$
$\displaystyle=$
$\displaystyle{\bm{\alpha}_{0}(\\{j\\})\bm{\beta}_{0}(\\{j\\})+\mathbf{h}_{1n}(\\{j\\})\bm{\beta}_{0}(\\{j\\})+\mathbf{h}_{2n}(\\{j\\})\bm{\alpha}_{0}(\\{j\\})+\mathbf{h}_{2n}(\\{j\\})\mathbf{h}_{1n}(\\{j\\})}$
$\displaystyle+\big{\\{}O_{p}(n^{-1/2})[\bm{\alpha}_{0n}(\\{j\\})+\bm{\beta}_{0n}(\\{j\\})]\big{\\}}+O_{p}(n^{-1})$
$\displaystyle=$
$\displaystyle\bm{\alpha}_{0}(\\{j\\})\bm{\beta}_{0}(\\{j\\})$
$\displaystyle+O_{p}\Big{(}\max\Big{\\{}\mathrm{1}(\bm{\beta}_{0}(\\{j\\})\neq
0)r_{1nj},\mathrm{1}(\bm{\alpha}_{0}(\\{j\\})\neq 0)r_{2nj},r_{1nj}r_{2nj},$
$\displaystyle\hphantom{+O_{p}\Big{(}\max\Big{\\{}\mathrm{1}}\mathrm{1}(\bm{\alpha}_{0n}(\\{j\\}),\bm{\beta}_{0n}(\\{j\\})\neq
0)n^{-1/2},n^{-1}\Big{\\}}\Big{)}$

for the product of the pilot estimates, as well as

$\displaystyle\hat{\bm{\beta}}(\\{j\\})=$
$\displaystyle\bm{\beta}_{0}(\\{j\\})+\mathbf{h}_{2n}(\\{j\\})+(\hat{\bm{\beta}}(\\{j\\})-\bm{\beta}_{0n}(\\{j\\}))$
$\displaystyle=$
$\displaystyle\bm{\beta}_{0}(\\{j\\})+O_{p}\Big{(}\max\Big{\\{}r_{2nj},n^{-1/2}\Big{\\}}\Big{)},$

for the $\hat{\bm{\beta}}$ pilot estimate.

Since both $r_{1nj},r_{2nj}\rightarrow 0$, the product $r_{1nj}r_{2nj}$ will
always be smaller than either rate individually. Hence, as long as one of
$\bm{\alpha}_{0}(\\{j\\}),\bm{\beta}_{0}(\\{j\\})$ are nonzero, we may safely
remove this product term from the maximum above.

###### Proof of Theorem 4.5.

As before, we will apply Lemma C.8. Hence, we simply check the behavior of the
variable-specific penalties.

Suppose
$\hat{\mathbf{w}}(\\{j\\})=|\hat{\bm{\alpha}}(\\{j\\})\hat{\bm{\beta}}(\\{j\\})|^{-\kappa}$
for fixed $\kappa>0$.

  * •

Case 1: $\bm{\alpha}_{0}(\\{j\\})\bm{\beta}_{0}(\\{j\\})\neq 0$. Little
changes from the fixed-coefficient case. $v_{nj}=n^{-1/2}\lambda_{n}O_{p}(1)$.

  * •

Case 2:  $\bm{\alpha}_{0}(\\{j\\})=0$, $\bm{\beta}_{0}(\\{j\\})\neq 0$. Let
$R^{(2)}_{nj}=\max\Big{\\{}r_{1nj},n^{-1/2}\Big{\\}}^{-1}$. Since
$R^{(2)}_{nj}\hat{\bm{\alpha}}(\\{j\\})\hat{\bm{\beta}}(\\{j\\})=O_{p}(1)$, we
have $v_{nj}=(R^{(2)}_{nj})^{\kappa}n^{-1/2}\lambda_{n}O_{p}(1)$.

  * •

Case 3:  $\bm{\alpha}_{0}(\\{j\\})\neq 0$, $\bm{\beta}_{0}(\\{j\\})=0$.
Similarly define $R^{(3)}_{nj}=\max\Big{\\{}r_{2nj},n^{-1/2}\Big{\\}}^{-1}$.
In this case $v_{nj}=(R^{(3)}_{nj})^{\kappa}n^{-1/2}\lambda_{n}O_{p}(1)$.

  * •

Case 4:  $\bm{\alpha}_{0}(\\{j\\})=0$, $\bm{\beta}_{0}(\\{j\\})=0$. Define
$R^{(4)}_{nj}=\max\Big{\\{}r_{1nj}r_{2nj},n^{-1/2}\Big{\\}}^{-1}$. Then

$v_{nj}=(R^{(4)}_{nj})^{\kappa}n^{-1/2}\lambda_{n}O_{p}(1).$

By construction, we have $R^{(2)}_{nj}=n^{c_{1j}},\ R^{(3)}_{nj}=n^{c_{2j}},\
\text{ and }R^{(4)}_{nj}=\sqrt{n}$ for $j=1,\ldots,p$. Each of these rates is
no larger than $\sqrt{n}$, hence allowing
$n^{(\kappa-1)/2}\lambda_{n}\rightarrow 0$ ensures that the variables in all 4
of the cases above will be included in the model asymptotically.

For the ADP weights, we can restrict to two cases:

  * •

Cases 1&2: $\bm{\beta}_{0}(\\{j\\})\neq 0$.
$v_{nj}=n^{-1/2}\lambda_{n}O_{p}(1)$.

  * •

Cases 3&4: $\bm{\beta}_{0}(\\{j\\})=0$. Use $R^{(3)}_{nj}$ as defined above
and note $R^{(3)}_{nj}\hat{\bm{\beta}}(\\{j\\})=O_{p}(1)$. Then
$v_{nj}=(R^{(3)}_{nj})^{\kappa}n^{-1/2}\lambda_{n}O_{p}(1)$. When
$\beta_{nj}=0$, this corresponds to $r_{2nj}=0$ and hence
$R^{(3)}_{nj}=\sqrt{n}$. Otherwise, if $c_{2j}=1/2$, we also have
$R^{(3)}_{nj}=\sqrt{n}$.

∎

###  C.7 Proof of Theorem 4.4

###### Proof of Theorem 4.4.

The proof of this Theorem follows nearly identically to the proof in Minnier
et al. (2011). In Section A.1 of that reference, the authors essentially study
the behavior of the following quantity in order to establish the perturbation-
bootstrap distribution:

$n\\{\mathcal{L}_{n}^{b}(\bm{\theta}_{0}+n^{-1/2}\mathbf{u};\hat{\mathbf{w}}^{b},\lambda_{n})-\mathcal{L}_{n}^{b}(\bm{\theta}_{0};\hat{\mathbf{w}}^{b},\lambda_{n})\\},$

where $\mathcal{L}_{n}^{b}$ is defined similarly to $\mathcal{L}_{n}$ with
multipliers in the risk. Define $\tilde{\bm{\theta}}^{Lb}$ and
$\tilde{\bm{\alpha}}^{b}$ as in (12) but using the true $\bm{\mu}_{0}$ rather
than the estimates. Following along these arguments, we would establish

  * (i’)

$Pr\big{\\{}\tilde{\bm{\theta}}^{Lb}(\\{j\\})\neq 0\text{ for any
}j\in\mathcal{S}^{*c}\big{\\}}\rightarrow 0$ and
$\sqrt{n}\big{\\{}\tilde{\bm{\theta}}^{Lb}({\mathcal{S}^{*}})-\tilde{\bm{\theta}}_{\mathcal{S}^{*}}\big{\\}}\overset{d}{\longrightarrow}N(\bm{0},\mathcal{J}_{1\mathcal{S}^{*}})$.

  * (ii’)

$\sqrt{n}(\tilde{\bm{\alpha}}^{b}-\tilde{\bm{\alpha}})({\mathcal{S}^{*}})\overset{d}{\longrightarrow}N(\bm{0},\mathcal{J}_{2\mathcal{S}^{*}})$.

One implication of these results is that
$\tilde{\bm{\theta}}_{\mathcal{S}^{*}}^{L*}$ and $\tilde{\bm{\alpha}}^{b}$ are
both $O_{p^{*}}(1)$, where we use the subscript of $p^{*}$ to indicate that
the implied probability statements should be taken with respect to the product
measure of the observed data and the multipliers. By Lemma C.6, and under our
assumptions,

$\sup_{\|\mathbf{u}\|_{\infty}\leq
C}\Big{|}\mathbb{P}_{n}^{*}\left\\{\bar{L}(\bm{\theta}_{0}+n^{-1/2}\mathbf{u};\hat{\bm{\mu}},\mathbf{O})-\bar{L}(\bm{\theta}_{0}+n^{-1/2}\mathbf{u};\bm{\mu}_{0},\mathbf{O})\right\\}\Big{|}=o_{p^{*}}(n^{-1}).$

This implies that
$n[\hat{\mathcal{L}}_{n}^{b}(\bm{\theta}_{0}+n^{-1/2}\mathbf{u};\hat{\mathbf{w}}^{b},\lambda_{n})-\hat{\mathcal{L}}_{n}^{b}(\bm{\theta}_{0};\hat{\mathbf{w}}^{b},\lambda_{n})]$
is a uniformly $o_{p^{*}}(1)$ approximation to
$n[\mathcal{L}_{n}^{b}(\bm{\theta}_{0}+n^{-1/2}\mathbf{u};\hat{\mathbf{w}}^{b},\lambda_{n})-\mathcal{L}_{n}^{b}(\bm{\theta}_{0};\hat{\mathbf{w}}^{b},\lambda_{n})]$
over compacts of $\mathbf{u}$. Hence the result (i) in the Theorem statement
can be inferred from (i’). We can also apply this same logic to the un-
penalized risks derived from the loss functions
$L_{j}(\alpha_{j};\bm{\mu},\mathbf{O}_{i})=\Big{(}\mathbf{M}_{i}(\\{j\\})-\mu_{Mj}(\mathbf{X}_{i})-\bm{\alpha}(\\{j\\})\\{D_{i}-\mu_{D}(\mathbf{X}_{i})\\}\Big{)}^{2}$
for $j=1,\ldots,p$ in order to establish (ii) from (ii’). The final result
(iii) follows from the previous results along with the Delta Method. ∎

##  Appendix D Appendix: Natural Effects for Submodels

Here, we prove the natural effect representations for submodels
$\mathcal{S}\subseteq\mathcal{S}^{F}$. Under the causal assumptions 2.1(ii),
2.1(iv), 2.1(iii) and 2.1(v) for submodel $\mathcal{S}$, we may follow the
development of VanderWeele and Vansteelandt (2014) and use Pearl’s mediation
formula (Pearl, 2001):

$\mathbb{E}(Y_{d,\mathbf{M}_{\mathcal{S},d^{*}}}~{}|~{}\mathbf{X})=\int_{\mathbf{m}_{\mathcal{S}}}\mathbb{E}\left(Y~{}|~{}\mathbf{X},D=d,\mathbf{M}_{\mathcal{S}}=\mathbf{m}_{\mathcal{S}}\right)dP_{0}(\mathbf{m}_{\mathcal{S}}~{}|~{}\mathbf{X},D=d^{*})\\\
=\mu_{Y}(\mathbf{X})+d\gamma_{0\mathcal{S}}+\left[\int_{\mathbf{m}_{\mathcal{S}}}\\{\mathbf{M}_{\mathcal{S}}-\bm{\mu}_{M\mathcal{S}}(\mathbf{X})\\}^{\top}dP_{0}(\mathbf{m}_{\mathcal{S}}~{}|~{}\mathbf{X},D=d^{*})\right]\bm{\beta}_{0\mathcal{S}}\\\
=\mu_{Y}(\mathbf{X})+d\gamma_{0\mathcal{S}}+\mathbb{E}\left[\\{\mathbf{M}_{\mathcal{S}}-\bm{\mu}_{M\mathcal{S}}(\mathbf{X})\\}^{\top}~{}|~{}\mathbf{X},D=d^{*}\right]\bm{\beta}_{0\mathcal{S}}\\\
=\mu_{Y}(\mathbf{X})+d\gamma_{0\mathcal{S}}+d^{*}\bm{\alpha}_{0\mathcal{S}}^{\top}\bm{\beta}_{0\mathcal{S}}.$
(46)

The second line follows from (5) along with the requirement that
$\mathbb{E}(\epsilon_{\mathcal{S}^{*}}~{}|~{}\mathbf{X},D,\mathbf{M}_{\mathcal{S}^{*}})=0$.
The NDE fixes $d^{*}$ and contrasts $d=1$ vs $d=0$. The NIE fixes $d$ and
varies $d^{*}=1$ vs $d^{*}=0$. The representations (6)-(7) follow directly.

##  Appendix E Remarks on Assumptions for Mediator Selection

######  Remark E.1.

Assumptions 2.2 and 2.1 are sufficient to achieve the second and third goals
in the case where $\mathcal{S}^{*}$ is known exactly. However, we are using
the data in the full model $\mathcal{S}^{F}$ to search for associations and
establish an appropriate set of mediators. Without Assumption 2.3, we have not
ruled out the possibility that some of the collected variables indexed by
$\mathcal{S}^{*c}$ have associations attributable to confounding. This may
lead to spurious members of $\mathcal{S}^{*}$.

######  Remark E.2.

Assumptions 2.1(ii), 2.1(iv), 2.1(iii) and 2.1(v) for $\mathcal{S}^{*}$
generally follow if Assumptions 2.1(ii), 2.1(iv), 2.1(iii) and 2.1(v) hold for
$\mathcal{S}^{F}$ and none of $\mathbf{M}(\mathcal{S}^{*c})$ confound the
relationship between $\mathbf{M}(\mathcal{S}^{*})$ and either of $D$ or $Y$,
given pre-treatment covariates $\mathbf{X}$.

######  Remark E.3.

We might also relax the requirement in Assumption 2.3 that (1) is correctly-
specified. Since we take the view that $\mathrm{NIE}_{\mathcal{S}^{F}}$ is not
directly interesting, there is no need to provide a valid causal estimate
using the observed data for this quantity. However, we still require
Assumption 2.2 in order to estimate and define the $\mathcal{S}^{*}$-based
estimands.

As an example, suppose $\alpha_{0p}=0$ and the outcome-model term
$\beta_{0p}M_{p}$ was replaced by some possibly-nonlinear function
$\phi(M_{p})\in L_{2}(P_{0})(\mathbb{R})$. In this case, we would require
$\phi(M_{p})$ to effectively be included in the submodel errors in order to
satisfy Assumption 2.2. Such a situation might occur if
$M_{p}\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\\{\mathbf{M}(\mathcal{S}^{*}),D\\}~{}|~{}\mathbf{X}$.

######  Remark E.4.

We remarked in Section 2.2 that a slightly different estimand may obviate
Assumption 2.1(v). Even in this case, Assumption 2.1(iv) is not trivially
satisfied for the submodel $\mathcal{S}^{*}$ even if all four assumptions hold
for the full model $\mathcal{S}^{F}$.

To see this, consider the following two-mediator causal model:

$\displaystyle Y_{d,m_{1},m_{2}}=$
$\displaystyle\beta_{1}m_{1}+\beta_{2}m_{2}+\epsilon$ $\displaystyle M_{1d}=$
$\displaystyle\alpha_{1}d+\eta_{1}$ $\displaystyle M_{2d}=$
$\displaystyle\eta_{2}.$

We assume that the observed treatment $D$ is randomly assigned, and that the
set of baseline confounders $\mathbf{X}$ is empty. In the model above, $M_{1}$
is affected by treatment, although there is a strong lack of treatment effect
for $M_{2}$, as both potential outcomes are equivalent for every subject.
Consequently, $\mathcal{S}^{F}=\\{1,2\\}$ and $\mathcal{S}^{*}=\\{1\\}$. The
observed data is $Y=Y_{D,M_{1},M_{2}}$, $M_{1}=M_{1D}$, and $M_{2}=M_{2D}$.
Under Assumption 2.3, we must have $\epsilon\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\\{\eta_{1},\eta_{2}\\}$.

Now let’s examine if Assumption 2.1(iv) holds in $\mathcal{S}^{*}$. To do
this, we must look at the potential outcomes resulting from control only of
$d$ and $m_{1}$. This results in the potential outcomes

$Y_{d,m_{1}}=Y_{d,m_{1},M_{2}}=\beta_{1}m_{1}+\beta_{2}\eta_{2}+\epsilon.$

Due to the randomization of $D$ and empty $\mathbf{X}$, we can simplify the
requirement of the relevant assumption to

$Y_{d,m_{1}}\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}M_{1},$

which upon substitution of the previous display becomes

$\beta_{1}m_{1}+\beta_{2}\eta_{2}+\epsilon\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\alpha_{1}D+\eta_{1}.$

Since $D$ is randomized, we may effectively remove it from the RHS. The
remaining requirement is then that $\eta_{1}\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\beta_{2}\eta_{2}+\epsilon$. Since Assumption
2.3 ensures that $\eta_{1}\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\epsilon$, the last piece is
$\eta_{1}\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\beta_{2}\eta_{2}$.

Therefore, in order for the causal assumptions to hold for $\mathcal{S}^{*}$,
we must have that either $\beta_{2}=0$ or $\eta_{2}\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}\eta_{1}$. Correlations between the post-
treatment variables may violate the unmeasured confounding Assumption 2.1(iv)
if the excluded variables are predictive of the outcome.

##  Appendix F Numerical Equivalence Between Submodel and Full-model Effects

Using the Model (1), we have that

$(\mathbf{Z}-\bm{\mu}_{Z0})(\mathcal{S}_{Z}^{*})=(D-\mu_{D0})\begin{pmatrix}1\\\
\bm{\alpha}_{0}\end{pmatrix}(\mathcal{S}_{Z}^{*})+\begin{pmatrix}0\\\
\bm{\eta}\end{pmatrix}(\mathcal{S}_{Z}^{*})$

with $\mathbb{E}(\bm{\eta}|\mathbf{X},D)=0$ and
$\mathbb{E}(D-\mu_{D0}|\mathbf{X})=0$.

Recall $\mathbf{H}_{0}=\mathbb{E}(\mathbf{Z}-\bm{\mu}_{Z0})^{\otimes 2}$.
Write
$\bm{\theta}_{0\mathcal{S}^{*}}=[\mathbf{H}_{0}(\mathcal{S}_{Z}^{*})]^{-1}\mathbf{H}_{0}(\mathcal{S}_{Z}^{*},\mathcal{S}_{Z}^{F})\bm{\theta}_{0}$,
where $(\mathcal{S}_{Z}^{*},\mathcal{S}_{Z}^{F})$ represents subsetting the
rows by $\mathcal{S}_{Z}^{*}$, but including all columns in
$\mathcal{S}_{Z}^{F}$. Using the conditional expectations above and some
matrix algebra, we can show

$\mathbf{H}_{0}(\mathcal{S}_{Z}^{*},\mathcal{S}_{Z}^{F})=\mathbb{E}\left\\{(D-\mu_{D0})^{2}\right\\}\begin{pmatrix}1\\\
\bm{\alpha}_{0}(\mathcal{S}^{*})\end{pmatrix}\begin{pmatrix}1&\bm{\alpha}_{0}^{\top}\end{pmatrix}+\begin{pmatrix}0\\\
\bm{\eta}(\mathcal{S}^{*})\end{pmatrix}\begin{pmatrix}0&\bm{\eta}^{\top}\end{pmatrix}$

A similar decomposition may be shown to hold for
$\mathbf{H}_{0}(\mathcal{S}^{*})$. This allows us to decompose
$\bm{\theta}_{0\mathcal{S}^{*}}$ into separate terms depending on
$\bm{\theta}_{0}(\mathcal{S}_{Z}^{*})$ and
$\bm{\theta}_{0}(\mathcal{S}_{Z}^{*c})$:

$\bm{\theta}_{0\mathcal{S}^{*}}=\bm{\theta}_{0}(\mathcal{S}^{*})\\\
+[\mathbf{H}_{0}(\mathcal{S}_{Z}^{*})]^{-1}\bigg{\\{}\mathbb{E}\left\\{(D-\mu_{D0})^{2}\right\\}\begin{pmatrix}1\\\
\bm{\alpha}_{0}(\mathcal{S}^{*})\end{pmatrix}\bm{\alpha}_{0}(\mathcal{S}^{*c})^{\top}+\mathbb{E}\left[\begin{pmatrix}\bm{0}^{\top}\\\
\bm{\eta}(\mathcal{S}^{*})\bm{\eta}(\mathcal{S}^{*c})^{\top}\end{pmatrix}\right]\bigg{\\}}\bm{\beta}_{0}(\mathcal{S}^{*c})\\\
=\bm{\theta}_{0}(\mathcal{S}^{*})+[\mathbf{H}_{0}(\mathcal{S}_{Z}^{*})]^{-1}~{}\mathbb{E}\left[\begin{pmatrix}\bm{0}^{\top}\\\
\bm{\eta}(\mathcal{S}^{*})\bm{\eta}(\mathcal{S}^{*c})^{\top}\end{pmatrix}\right]\bm{\beta}_{0}(\mathcal{S}^{*c}).$

The final line follows because
$\bm{\alpha}_{0}(\mathcal{S}^{*c})\cdot\bm{\beta}_{0}(\mathcal{S}^{*c})=0$.
This leads to the stated conditions in Remark 2.4


#  EfficientViT: Lightweight Multi-Scale Attention for
On-Device Semantic Segmentation

Han Cai1, Junyan Li2, Muyan Hu3, Chuang Gan4, Song Han1
1MIT, 2Zhejiang University, 3Tsinghua University, 4MIT-IBM Watson AI Lab
https://github.com/mit-han-lab/efficientvit

###### Abstract

Semantic segmentation enables many appealing real-world applications, such as
computational photography, autonomous driving, etc. However, the vast
computational cost makes deploying state-of-the-art semantic segmentation
models on edge devices with limited hardware resources difficult. This work
presents EfficientViT, a new family of semantic segmentation models with a
novel lightweight multi-scale attention for on-device semantic segmentation.
Unlike prior semantic segmentation models that rely on heavy self-attention,
hardware-inefficient large-kernel convolution, or complicated topology
structure to obtain good performances, our lightweight multi-scale attention
achieves a global receptive field and multi-scale learning (two critical
features for semantic segmentation models) with only lightweight and hardware-
efficient operations. As such, EfficientViT delivers remarkable performance
gains over previous state-of-the-art semantic segmentation models across
popular benchmark datasets with significant speedup on the mobile platform.
Without performance loss on Cityscapes, our EfficientViT provides up to
15$\times$ and 9.3$\times$ mobile latency reduction over SegFormer and
SegNeXt, respectively. Maintaining the same mobile latency, EfficientViT
provides +7.4 mIoU gain on ADE20K over SegNeXt.

##  1 Introduction

Semantic segmentation is a fundamental task in computer vision, which aims to
assign a class label to each pixel in the input image. Semantic segmentation
has broad applications in real-world scenarios, including autonomous driving,
medical image processing, computational photography, etc. Therefore, deploying
state-of-the-art (SOTA) semantic segmentation models on edge devices is in
great demand to benefit a wide range of users.

However, there is a large gap between the computational cost required by SOTA
semantic segmentation models and the limited resources of edge devices. It
makes deploying these models on edge devices impractical. In particular,
semantic segmentation is a dense prediction task requiring high-resolution
images and strong context information extraction ability to work well [1, 36,
47, 52, 48, 42]. Therefore, directly porting efficient model architecture from
image classification is unsuitable for semantic segmentation.

Figure 1: Latency vs. Performance. All performance results are obtained with the single model and single-scale inference. The latency results are obtained on the Qualcomm Snapdragon 8Gen1 CPU using Tensorflow-Lite. Compared with state-of-the-art (SOTA) segmentation models, EfficientViT achieves a remarkable boost in speed while providing the same or higher performances on Cityscapes and ADE20K. In addition, EfficientViT also shows strong performances in image classification, achieving a 7.9x latency reduction over EfficientNet without accuracy loss on ImageNet. Table 1: Desirable Features for On-device Semantic Segmentation. ‘Linear computational complexity’ means the computational cost grows linearly as the input resolution increases.  Features |  SegFormer [45] |  HRFormer [49] |  SegNeXt [17] | EfficientViT
---|---|---|---|---
Global receptive field | ✓ |  |  | ✓
Multi-scale learning |  | ✓ | ✓ | ✓
Linear computational complexity |  | ✓ | ✓ | ✓
Hardware efficiency |  |  |  | ✓

This work introduces EfficientViT, a new family of models for on-device
semantic segmentation. The core of EfficientViT is a novel lightweight multi-
scale attention module that enables a global receptive field and multi-scale
learning with hardware-efficient operations. Our module is motivated by prior
SOTA semantic segmentation models. They demonstrate that the multi-scale
learning [47, 52], and global receptive field [45] play a critical role in
improving the performances for semantic segmentation. However, they do not
consider hardware efficiency when designing their models, which is essential
for on-device semantic segmentation. For example, SegFormer [45] introduces
self-attention into the backbone to have a global receptive field. But its
computational complexity is quadratic to the input resolution, making it
unable to handle high-resolution images efficiently. SegNeXt [17] proposes a
multi-branch module with large-kernel convolutions (kernel size up to 21) to
enable a large receptive field and multi-scale learning. However, large-kernel
convolution requires exceptional support on hardware to achieve good
efficiency [15], which is usually not available on edge devices.

Hence, the design principle of our module is to enable these two critical
features while avoiding hardware-inefficient operations. Specifically, to have
a global receptive field, we propose substituting the inefficient self-
attention with lightweight ReLU-based global attention [26]. By leveraging the
associative property of matrix multiplication, ReLU-based global attention can
reduce the computational complexity from quadratic to linear while preserving
functionality. In addition, it avoids hardware-inefficient operations like
softmax, making it more suitable for on-device semantic segmentation (Figure
3).

Furthermore, we propose a novel lightweight multi-scale attention module based
on the ReLU-based global attention. Specifically, we aggregate nearby tokens
with small-kernel convolutions to generate multi-scale tokens and perform
ReLU-based global attention on multi-scale tokens (Figure 2) to combine the
global receptive field with multi-scale learning. We summarize the comparison
between our work and prior SOTA semantic segmentation models in Table 1. We
can see that our model is more suitable for on-device semantic segmentation
than previous models.

We extensively evaluate EfficientViT on popular semantic segmentation
benchmark datasets, including Cityscapes [12] and ADE20K [53]. EfficientViT
provides significant performance boosts over prior SOTA semantic segmentation
models. More importantly, EfficientViT does not involve hardware-inefficient
operations, so our FLOPs reduction can easily translate to latency reduction
on mobile devices (Figure 1). On Qualcomm Snapdragon 8Gen1 CPU, EfficientViT
executes 5.8$\times$ faster than SegNeXt [17] while reaching higher mIoU on
Cityscapes and 7.9$\times$ faster than EfficientNet [39] without accuracy loss
on ImageNet. We summarize our contributions as follows:

  * •

We introduce a novel lightweight multi-scale attention for on-device semantic
segmentation. It achieves a global receptive field and multi-scale learning
while maintaining good efficiency on edge devices.

  * •

We design EfficientViT, a new family of models, based on the proposed
lightweight multi-scale attention module.

  * •

On popular semantic segmentation benchmark datasets and ImageNet, our model
demonstrates remarkable speedup on mobile over prior SOTA semantic
segmentation models.

Figure 2: Illustration of EfficientViT’s Building Block (left) and the
Proposed Lightweight Multi-Scale Attention (right). _Left_ : A building block
of EfficientViT consists of a lightweight MSA module and an MBConv. The
lightweight MSA module is responsible for capturing context information, while
the MBConv is for capturing local information. _Right_ : After getting Q/K/V
tokens via the linear projection layer, we propose to generate multi-scale
tokens by aggregating nearby tokens via lightweight small-kernel convolutions.
ReLU-based global attention is applied to multi-scale tokens, and the outputs
are concatenated and fed to the final linear projection layer for feature
fusing.

##  2 Method

This section first introduces lightweight Multi-Scale Attention (MSA). Unlike
prior works, our lightweight MSA module simultaneously achieves a global
receptive field and multi-scale learning with only hardware-efficient
operations. Then we present a new family of models named EfficientViT based on
the proposed MSA module for on-device semantic segmentation.

Figure 3: ReLU-based linear attention is 3.3-4.5$\times$ faster than softmax
attention with similar computation, thanks to removing hardware-unfriendly
operations (e.g., softmax). Latency is measured on Qualcomm Snapdragon 855 CPU
with TensorFlow-Lite. Figure 4: Macro Architecture of EfficientViT. We adopt
the standard backbone-head/encoder-decoder design. In the backbone, we insert
our lightweight MSA modules in Stages 3 and 4. Following the common practice,
we feed the features from the last three stages (P2, P3, and P4) to the head.
We use addition to fuse these features for simplicity and efficiency. As we
already have lightweight MSA modules in the backbone, we adopt a simple head
design that consists of several MBConv blocks and output layers.

###  2.1 Lightweight Multi-Scale Attention

Our lightweight MSA module balances two crucial aspects for on-device semantic
segmentation, i.e., performance and efficiency. Specifically, a global
receptive field and multi-scale learning are essential from the performance
perspective. Previous SOTA segmentation models provide strong performances by
enabling these features but fail to provide good efficiency. Our module
tackles this issue by trading slight capacity loss for significant efficiency
improvements.

An illustration of the proposed lightweight MSA module is provided in Figure 2
(right). In particular, we propose to use lightweight ReLU-based attention
[26] to enable the global receptive field instead of the heavy self-attention
[41]. While ReLU-based attention [26] and other linear attention modules [2,
11, 38, 43] has been explored in other domains, it has never been applied to
the semantic segmentation community. To the best of our knowledge, we are the
first work demonstrating ReLU-based attention’s effectiveness in semantic
segmentation. In addition, our work introduces novel designs (lightweight MSA
module) to enhance the capacity, making it much more powerful in semantic
segmentation.

#### Enabling Global Receptive Field with Lightweight ReLU-based Attention.

Given input $x\in\operatorname*{\mathbb{R}}^{N\times f}$, the generalized form
of self-attention can be written as:

$\small
O_{i}=\sum_{j=1}^{N}\frac{Sim(Q_{i},K_{j})}{\sum_{j=1}^{N}{Sim(Q_{i},K_{j})}}V_{j},$
(1)

where $Q=xW_{Q}$, $K=xW_{K}$, $V=xW_{V}$ and
$W_{Q}/W_{K}/W_{V}\in\operatorname*{\mathbb{R}}^{f\times d}$ is the learnable
linear projection matrix. $O_{i}$ represents the i-th row of matrix $O$.
$Sim(\cdot,\cdot)$ is the similarity function. When using the similarity
function $Sim(Q,K)=\exp(\frac{QK^{T}}{\sqrt{d}})$, Eq. (1) becomes the
original self-attention [41].

Apart from $\exp(\frac{QK^{T}}{\sqrt{d}})$, we can use other similarity
functions. In this work, we use ReLU-based global attention [26] to achieve
both the global receptive field and linear computational complexity. In ReLU-
based global attention, the similarity function is defined as

$\small Sim(Q,K)=\text{ReLU}(Q)\text{ReLU}(K)^{T}.$ (2)

With $Sim(Q,K)=\text{ReLU}(Q)\text{ReLU}(K)^{T}$, Eq. (1) can be rewritten as:

$\displaystyle O_{i}$
$\displaystyle=\sum_{j=1}^{N}\frac{\text{ReLU}(Q_{i})\text{ReLU}(K_{j})^{T}}{\sum_{j=1}^{N}{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\text{ReLU}(Q_{i})}\text{ReLU}(K_{j})^{T}}}V_{j}$
$\displaystyle=\frac{\sum_{j=1}^{N}(\text{ReLU}(Q_{i})\text{ReLU}(K_{j})^{T})V_{j}}{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\text{ReLU}(Q_{i})}\sum_{j=1}^{N}\text{ReLU}(K_{j})^{T}}.$

Figure 5: Illustration of the Aggregation Process for Generating Multi-Scale
Tokens. The information aggregation is done independently for each Q, K, and V
in each head. ‘d’ denotes the dimension of each token. The typical value of d
is 32.

Then, we can leverage the associative property of matrix multiplication to
reduce the computational complexity and memory footprint from quadratic to
linear without changing the functionality:

$\displaystyle O_{i}$
$\displaystyle=\frac{\sum_{j=1}^{N}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}[\text{ReLU}(Q_{i})\text{ReLU}(K_{j})^{T}]}V_{j}}{\text{ReLU}(Q_{i})\sum_{j=1}^{N}\text{ReLU}(K_{j})^{T}}$
$\displaystyle=\frac{\sum_{j=1}^{N}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\text{ReLU}(Q_{i})}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}[(\text{ReLU}(K_{j})^{T}V_{j})]}}{\text{ReLU}(Q_{i})\sum_{j=1}^{N}\text{ReLU}(K_{j})^{T}}$
$\displaystyle=\frac{{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\text{ReLU}(Q_{i})}(\sum_{j=1}^{N}\text{ReLU}(K_{j})^{T}V_{j})}{\text{ReLU}(Q_{i})(\sum_{j=1}^{N}\text{ReLU}(K_{j})^{T})}.$
(3)

As shown in Eq. (2.1), we only need to compute
$(\sum_{j=1}^{N}\text{ReLU}(K_{j})^{T}V_{j})$
$\in\operatorname*{\mathbb{R}}^{d\times d}$ and
$(\sum_{j=1}^{N}\text{ReLU}(K_{j})^{T})$
$\in\operatorname*{\mathbb{R}}^{d\times 1}$ once, then can reuse them for each
query, thereby only requires $\mathcal{O}(N)$ computational cost and
$\mathcal{O}(N)$ memory.

Another key merit of ReLU-based global attention is that it does not involve
hardware-unfriendly operations like softmax, making it more efficient on
hardware. For example, Figure 3 shows the latency comparison between softmax
attention and ReLU-based linear attention. With similar computation, ReLU-
based linear attention is significantly faster than softmax attention on
mobile.

#### Generate Multi-Scale Tokens.

ReLU-based attention alone has limited model capacity. To enhance ReLU-based
global attention with multi-scale learning ability, we propose to aggregate
the information from nearby Q/K/V tokens to get multi-scale tokens. The
aggregation process is illustrated in Figure 5. This information aggregation
process is independent for each Q, K, and V in each head. We only use small-
kernel convolutions for information aggregation to avoid hurting hardware
efficiency.

In the practical implementation, independently executing these aggregation
operations is inefficient on GPU. Therefore, we take advantage of the
infrastructure of group convolution in modern deep learning frameworks to
reduce the number of total operations. Specifically, all DWConvs are fused
into a single DWConv while all 1x1 Convs are combined into a single 1x1 group
convolution (Figure 2 right) where the number of groups is
$3\times\text{\\#heads}$ and the number of channels in each group is d.

After getting multi-scale tokens, we perform global attention upon them to
extract multi-scale global features. Finally, we concatenate the features from
different scales along the head dimension and feed them to the final linear
projection layer to fuse the features.

Table 2: Detailed Architecture Configurations of Different EfficientViT Variants. We build a series of models to fit different efficiency constraints. ‘C’ denotes the number of channels. ‘L’ denotes the number of blocks. ‘H’ is the height of the feature map, and ‘W’ is the width of the feature map.  Variants | Feature Map Shape | EfficientViT-B0 | EfficientViT-B1 | EfficientViT-B2 | EfficientViT-B3
---|---|---|---|---|---
Input Stem |  C $\times\frac{H}{2}\times\frac{W}{2}$ | C = 8, L = 1 | C = 16, L = 1 | C = 24, L = 1 | C = 32, L = 1
Stage1 |  C $\times\frac{H}{4}\times\frac{W}{4}$ | C = 16, L = 2 | C = 32, L = 2 | C = 48, L = 3 | C = 64, L = 4
Stage2 |  C $\times\frac{H}{8}\times\frac{W}{8}$ | C = 32, L = 2 | C = 64, L = 3 | C = 96, L = 4 | C = 128, L = 6
Stage3 |  C $\times\frac{H}{16}\times\frac{W}{16}$ | C = 64, L = 2 | C = 128, L = 3 | C = 192, L = 4 | C = 256, L = 6
Stage4 |  C $\times\frac{H}{32}\times\frac{W}{32}$ | C = 128, L = 2 | C = 256, L = 4 | C = 384, L = 6 | C = 512, L = 9
Head |  C $\times\frac{H}{8}\times\frac{W}{8}$ | C = 32, L = 1 | C = 64, L = 3 | C = 96, L = 3 | C = 128, L = 3

###  2.2 EfficientViT Architecture

We build a new family of models based on the proposed lightweight MSA module.
The core building block (denoted as ‘EfficientViT Module’) is illustrated in
Figure 2 (left). Specifically, an EfficientViT module comprises a lightweight
MSA module and an MBConv [37]. The lightweight MSA module is for context
information extraction, while the MBConv is for local information extraction.

The macro architecture of EfficientViT is demonstrated in Figure 4. We use the
standard backbone-head/encoder-decoder architecture design.

  * •

Backbone. The backbone of EfficientViT also follows the standard design, which
consists of the input stem and four stages with gradually decreased feature
map size and gradually increased channel number. We insert the EfficientViT
module in Stages 3 and 4. For downsampling, we use an MBConv with stride 2.

  * •

Head. P2, P3, and P4 denote the outputs of Stages 2, 3, and 4, forming a
pyramid of feature maps. For simplicity and efficiency, we use 1x1 convolution
and standard upsampling operation (e.g., bilinear/bicubic upsampling) to match
their spatial and channel size and fuse them via addition. Since our backbone
already has a strong context information extraction capacity, we adopt a
simple head design that comprises several MBConv blocks and the output layers
(i.e., prediction and upsample). In the experiments, we empirically find this
simple head design is sufficient for achieving SOTA performances thanks to our
lightweight MSA module.

In addition to semantic segmentation, our model can be applied to other vision
tasks, such as image classification, by combining the backbone with task-
specific heads.

Following the same macro architecture, we design a series of models with
different sizes to satisfy various efficiency constraints. The detailed
configurations are demonstrated in Table 2. We name these models as
EfficientViT-B0, EfficientViT-B1, EfficientViT-B2, and EfficientViT-B3,
respectively.

Figure 6: MACs vs. Performance. EfficientViT provides a better trade-off
between MACs and performance than SOTA semantic segmentation and image
classification models. On Cityscapes, EfficientViT provides 13$\times$ and
2.4$\times$ MACs reduction than SegFormer and SegNeXt, respectively, while
achieving the same or higher performances. On ImageNet, EfficientViT achieves
2.9x MACs reduction than EfficientNet without accuracy loss.

##  3 Experiments

###  3.1 Setups

Table 3: Ablation Study on Two Key Components of Our Lightweight MSA Module. The mIoU and MACs are measured on Cityscapes with 1024x2048 input resolution. We rescale the width of the models so that they have the same MACs. Multi-scale learning and the global receptive field are essential for obtaining good semantic segmentation performance.  Components | mIoU $\uparrow$ | Params $\downarrow$ | MACs $\downarrow$
---|---|---|---
Multi-scale | Global att.
|  | 68.1 | 0.7M | 4.4G
✓ |  | 72.3 | 0.7M | 4.4G
| ✓ | 72.2 | 0.7M | 4.4G
✓ | ✓ | 74.5 | 0.7M | 4.4G
Table 4: Backbone Performance of EfficientViT on ImageNet Classification. ‘r224’ means the input resolution is 224x224. While EfficientViT is mainly designed for semantic segmentation, it also works well on ImageNet classification. With 6.5G MACs, EfficientViT-B3 achieves 84.2 top1 ImageNet accuracy, surpassing EfficientNet-B6 while reducing the MACs by 2.9x and being 7.9x faster on mobile.  Models |  Top1 Acc $\uparrow$ |  Top5 Acc $\uparrow$ |  Params $\downarrow$ |  MACs $\downarrow$ |  Mobile Latency $\downarrow$ |  Speedup $\uparrow$
---|---|---|---|---|---|---
EfficientNet-B1 [39] | 79.1 | 94.4 | 7.8M | 0.70G | 87ms | 1.0x
EfficientNetV2-B0 [40] | 78.7 | - | 7.1M | 0.72G | 53ms | 1.6x
EfficientViT-B1 (r224) | 79.4 | 94.3 | 9.1M | 0.52G | 19ms | 4.6x
EfficientNet-B2 [39] | 80.1 | 94.9 | 9.2M | 1.0G | 118ms | 1.0x
EfficientNetV2-B1 [40] | 79.8 | - | 8.1M | 1.2G | 85ms | 1.4x
EfficientViT-B1 (r288) | 80.4 | 95.0 | 9.1M | 0.86G | 31ms | 3.8x
Swin-T [29] | 81.3 | - | 29M | 4.5G | - | -
ConvNeXt-T [30] | 82.1 | - | 29M | 4.5G | 356ms | 1.0x
EfficientNet-B3 [39] | 81.6 | 95.7 | 12M | 1.8G | 208ms | 1.7x
EfficientNetV2-B3 [40] | 82.1 | - | 14M | 3.0G | 201ms | 1.8x
CoAtNet-0 [13] | 81.6 | - | 25M | 4.2G | 175ms | 2.0x
EfficientViT-B2 (r256) | 82.7 | 96.1 | 24M | 2.1G | 72ms | 4.9x
Swin-S [29] | 83.0 | - | 50M | 8.7G | - | -
ConvNeXt-S [30] | 83.1 | - | 50M | 8.7G | 622ms | 1.0x
EfficientNet-B4 [39] | 82.9 | 96.4 | 19M | 4.2G | 464ms | 1.3x
CoAtNet-1 [13] | 83.3 | - | 42M | 8.4G | 332ms | 1.9x
EfficientViT-B3 (r224) | 83.5 | 96.4 | 49M | 4.0G | 140ms | 4.4x
Swin-B [29] | 83.5 | - | 88M | 15G | - | -
EfficientNet-B6 [39] | 84.0 | 96.8 | 43M | 19G | 1804ms | 1.0x
ConvNeXt-B [30] | 83.8 | - | 89M | 15G | 1025ms | 1.8x
CoAtNet-2 [13] | 84.1 | - | 75M | 16G | 600ms | 3.0x
EfficientNetV2-S [40] | 83.9 | - | 22M | 8.8G | 509ms | 3.5x
EfficientViT-B3 (r288) | 84.2 | 96.7 | 49M | 6.5G | 228ms | 7.9x

#### Datasets.

We evaluate the effectiveness of EfficientViT on two representative semantic
segmentation datasets, including Cityscapes [12] and ADE20K [53]. Cityscapes
is an autonomous driving dataset that mainly focuses on urban scenes. It
contains 5,000 fine-annotated high-resolution (1024x2048) images with 19
classes divided into three subsets of size 2,975/500/1,525 for
training/validation/testing. ADE20K is a scene-parsing dataset with 150
classes. It contains 20,210/2,000/3,352 images for training, validation, and
testing, respectively.

Apart from Cityscapes and ADE20K, we also study the effectiveness of
EfficientViT for image classification using the ImageNet dataset [14].

#### Latency Measurement.

We measure the latency of the models on Qualcomm Snapdragon 8Gen1 CPU with
Tensorflow-Lite111https://www.tensorflow.org/lite, batch size 1 and fp32.

Table 5: Comparison with SOTA Semantic Segmentation Models on Cityscapes. ‘r1024x2048’ denotes the input resolution is 1024x2048. Models with similar mIoU are grouped for efficiency comparison. Compared with SegNeXt-T, EfficientViT-B1 achieves 2.7x MACs reduction, 9.3x latency reduction, and 0.3 higher mIoU. Compared with SegFormer-B1, EfficientViT-B1 obtains 13x MACs saving, 15x measured speedup, and 1.6 higher mIoU. Models |  mIoU $\uparrow$ |  Params $\downarrow$ |  MACs $\downarrow$ |  Mobile Latency $\downarrow$ |  Speedup $\uparrow$
---|---|---|---|---|---
DeepLabV3plus-Mbv2 [7] | 75.2 | 15M | 555G | - | -
PSPNet-Mbv2 [52] | 70.2 | 14M | 423G | - | -
FCN-Mbv2 [31] | 61.5 | 9.8M | 317G | - | -
SegFormer-B0 (r768) [45] | 75.3 | 3.8M | 52G | 2.8s | 1.0x
EfficientViT-B0 (r960x1920) | 75.5 | 0.7M | 3.9G | 0.20s | 14x
HRFormer-S [49] | 80.0 | 14M | 836G | - | -
SegFormer-B1 [45] | 78.5 | 14M | 244G | 12s | 1.0x
SegNeXt-T [17] | 79.8 | 4.3M | 51G | 7.6s | 1.6x
EfficientViT-B1 (r896x1792) | 80.1 | 4.8M | 19G | 0.82s | 15x
HRFormer-B [49] | 81.9 | 56M | 2224G | - | -
SegFormer-B3 [45] | 81.7 | 47M | 963G | - | -
SegNeXt-S [17] | 81.3 | 14M | 125G | 18s | 1.0x
EfficientViT-B2 (r1024x2048) | 82.1 | 15M | 74G | 3.1s | 5.8x
SegFormer-B5 [45] | 82.4 | 85M | 1460G | - | -
SegNeXt-L [17] | 83.2 | 49M | 578G | - | -
EfficientViT-B3 (r1184x2368) | 83.2 | 40M | 240G | 10s | -

#### Implementation Details.

We implement our models using Pytorch [34] and train them on GPUs. We use the
AdamW optimizer with cosine learning rate decay for training our models. For
lightweight multi-scale attention, we use a two-branch design for the best
trade-off between performance and efficiency, where 5x5 nearby tokens are
aggregated to generate multi-scale tokens.

For semantic segmentation experiments, we use the mean Intersection over Union
(mIoU) as our evaluation metric. The backbone is initialized with weights
pretrained on ImageNet and the head is initialized randomly, following the
common practice. Common data augmentation strategies such as random scaling,
random horizontal flip, and random cropping are employed following prior
works.

###  3.2 Ablation Study

#### Effectiveness of Our Lightweight MSA Module.

We conduct ablation study experiments on Cityscapes to study the effectiveness
of two key design components of our lightweight MSA module, i.e., multi-scale
learning and global attention. To eliminate the impact of pre-training, we
train all models from random initialization. In addition, we rescale the width
of the models so that they have the same #MACs. The results are summarized in
Table 3. We can see that removing either global attention or multi-scale
learning will significantly hurt the performances. It shows that all of them
are essential for achieving a better trade-off between performance and
efficiency.

#### Backbone Performance on ImageNet.

To understand the effectiveness of EfficientViT’s backbone in image
classification, we train our models on ImageNet following the standard
training strategy (300 epochs with random initialization, no knowledge
distillation). We summarize the results and compare our models with SOTA image
classification models in Table 4.

Though EfficientViT is designed for semantic segmentation, it achieves highly
competitive performances on ImageNet. In particular, EfficientViT-B3 obtains
84.2 top1 accuracy on ImageNet, providing +0.2 accuracy gain over
EfficientNet-B6 and 7.9x speedup.

Table 6: Comparison with SOTA Semantic Segmentation Models on ADE20K. Compared with SegNeXt-S, EfficientViT-B2 provides a 5.2x speedup and 0.8 mIoU gain. Compared with SegFormer-B1, EfficientViT-B1 achieves 0.5 higher mIoU with a 6.5x speedup.  Models |  mIoU $\uparrow$ |  Params $\downarrow$ |  MACs $\downarrow$ |  Mobile Latency $\downarrow$ |  Speedup $\uparrow$
---|---|---|---|---|---
SegFormer-B1 [45] | 42.2 | 14M | 16G | 0.65s | 1.0x
SegNeXt-T [17] | 41.1 | 4.3M | 6.6G | 0.61s | 1.1x
EfficientViT-B1 (r480) | 42.7 | 4.8M | 2.7G | 0.10s | 6.5x
HRFormer-S [49] | 44.0 | 14M | 110G | - | -
SegNeXt-S [17] | 44.3 | 14M | 16G | 1.1s | 1.0x
EfficientViT-B2 (r416) | 45.1 | 15M | 6.0G | 0.21s | 5.2x
Mask2Former [9] | 47.7 | 47M | 74G | - | -
MaskFormer [10] | 46.7 | 42M | 55G | - | -
SegFormer-B2 [45] | 46.5 | 28M | 62G | 2.6s | 1.0x
EfficientViT-B3 (r384) | 48.0 | 39M | 12G | 0.45s | 5.8x
HRFormer-B [49] | 48.7 | 56M | 280G | - | -
SegNeXt-B [17] | 48.5 | 28M | 35G | 2.2s | 1.0x
EfficientViT-B3 (r512) | 49.0 | 39M | 22G | 0.80s | 2.8x

###  3.3 Main Results

#### Cityscapes.

Table 5 reports the comparison between EfficientViT and SOTA semantic
segmentation models on Cityscapes. EfficientViT achieves remarkable efficiency
improvements over prior SOTA semantic segmentation models without sacrificing
performances. Specifically, compared with SegFormer, EfficientViT obtains up
to 13x MACs saving and up to 15x latency reduction with higher mIoU. Compared
with SegNeXt, EfficientViT provides up to 2.7x MACs reduction and 9.3x speedup
on mobile while maintaining higher mIoU.

Having similar computational cost, EfficientViT yields significant performance
gains over previous SOTA models. For example, EfficientViT-B3 yields +4.7 mIoU
gain over SegFormer-B1 with similar MACs.

#### ADE20K.

Table 6 summarizes the comparison between EfficientViT and SOTA semantic
segmentation models on ADE20K. Similar to Cityscapes, we can see that
EfficientViT also achieves significant efficiency improvements on ADE20K. For
example, with +0.5 mIoU gain, EfficientViT-B1 provides 5.9x MACs reduction and
6.5x latency reduction than SegFormer-B1. With +0.8 mIoU gain, EfficientViT-B2
requires 2.7x fewer MACs and runs 5.2x faster than SegNeXt-S.

##  4 Related Work

#### Semantic Segmentation.

Semantic segmentation targets producing a class prediction for each pixel
given the input image. It can be viewed as an extension of image
classification from per-image prediction to per-pixel predictions. Since the
groundbreaking work FCN [31], which designs a fully convolutional neural
network for end-to-end pixel-to-pixel prediction, extensive studies have been
done to improve the performance for semantic segmentation [1, 36, 47, 52, 48,
42].

In addition, there are also some works targeting improving the efficiency of
semantic segmentation models [51, 35, 27, 46, 50]. Representative examples
include ICNet [51], DFANet [27], BiSeNet [46], etc. While these models provide
good efficiency, their performances are far behind SOTA semantic segmentation
models, especially on the challenging Cityscapes dataset.

Compared to these works, our models provide a better trade-off between
performance and efficiency by enabling a global receptive field and multi-
scale learning with lightweight operations.

#### Efficient Vision Transformer.

While ViT provides impressive performances in the high-computation region, it
is usually inferior to previous efficient CNNs [39, 24, 5, 18] when targeting
the low-computation region. To close the gap, MobileViT [33] proposes to
combine the strength of CNN and ViT by replacing local processing in
convolutions with global processing using transformers. MobileFormer [8]
proposes to parallelize MobileNet and Transformer with a two-way bridge in
between for feature fusing. NASViT [16] proposes to leverage neural
architecture search to search for efficient ViT architectures.

However, these models mainly focus on image classification and still rely on
self-attention with quadratic computational complexity, thus unsuitable for
on-device semantic segmentation.

#### Efficient Deep Learning.

Our work is also related to efficient deep learning, which aims at improving
the efficiency of deep neural networks so that we can deploy them on hardware
platforms with limited resources, such as mobile phones and IoT devices.
Typical technologies in efficient deep learning include network pruning [20,
22, 28], quantization [19], efficient model architecture design [25, 32], and
training techniques [23, 4]. In addition to manual designs, many recent works
use AutoML techniques [54, 3, 6] to automatically design [5], prune [21] and
quantize [44] neural networks.

##  5 Conclusion

In this work, we studied efficient architecture design for on-device semantic
segmentation. We introduced a lightweight multi-scale attention module that
simultaneously achieves a global receptive field, and multi-scale learning
with lightweight and hardware-efficient operations, thus providing significant
speedup on edge devices without performance loss than SOTA semantic
segmentation models. For future work, we will explore applying EfficientViT to
other vision tasks and further scaling up our EfficientViT models.

## References

  * [1] Vijay Badrinarayanan, Alex Kendall, and Roberto Cipolla.  Segnet: A deep convolutional encoder-decoder architecture for image segmentation.  IEEE transactions on pattern analysis and machine intelligence, 39(12):2481–2495, 2017.
  * [2] Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, and Judy Hoffman.  Hydra attention: Efficient attention with many heads.  In Computer Vision–ECCV 2022 Workshops: Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part VII, pages 35–49. Springer, 2023.
  * [3] Han Cai, Tianyao Chen, Weinan Zhang, Yong Yu, and Jun Wang.  Efficient architecture search by network transformation.  In AAAI, 2018.
  * [4] Han Cai, Chuang Gan, Ji Lin, and Song Han.  Network augmentation for tiny deep learning.  arXiv preprint arXiv:2110.08890, 2021.
  * [5] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han.  Once for all: Train one network and specialize it for efficient deployment.  In ICLR, 2020.
  * [6] Han Cai, Ligeng Zhu, and Song Han.  ProxylessNAS: Direct neural architecture search on target task and hardware.  In ICLR, 2019.
  * [7] Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam.  Encoder-decoder with atrous separable convolution for semantic image segmentation.  In Proceedings of the European conference on computer vision (ECCV), pages 801–818, 2018.
  * [8] Yinpeng Chen, Xiyang Dai, Dongdong Chen, Mengchen Liu, Xiaoyi Dong, Lu Yuan, and Zicheng Liu.  Mobile-former: Bridging mobilenet and transformer.  In Proceedings of the IEEE conference on computer vision and pattern recognition, 2022.
  * [9] Bowen Cheng, Ishan Misra, Alexander G Schwing, Alexander Kirillov, and Rohit Girdhar.  Masked-attention mask transformer for universal image segmentation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1290–1299, 2022.
  * [10] Bowen Cheng, Alex Schwing, and Alexander Kirillov.  Per-pixel classification is not all you need for semantic segmentation.  Advances in Neural Information Processing Systems, 34, 2021.
  * [11] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al.  Rethinking attention with performers.  arXiv preprint arXiv:2009.14794, 2020.
  * [12] Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele.  The cityscapes dataset for semantic urban scene understanding.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3213–3223, 2016.
  * [13] Zihang Dai, Hanxiao Liu, Quoc V Le, and Mingxing Tan.  Coatnet: Marrying convolution and attention for all data sizes.  Advances in Neural Information Processing Systems, 34:3965–3977, 2021.
  * [14] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In CVPR, 2009.
  * [15] Xiaohan Ding, Xiangyu Zhang, Yizhuang Zhou, Jungong Han, Guiguang Ding, and Jian Sun.  Scaling up your kernels to 31x31: Revisiting large kernel design in cnns.  In Proceedings of the IEEE conference on computer vision and pattern recognition, 2022.
  * [16] Chengyue Gong, Dilin Wang, Meng Li, Xinlei Chen, Zhicheng Yan, Yuandong Tian, qiang liu, and Vikas Chandra.  NASVit: Neural architecture search for efficient vision transformers with gradient conflict aware supernet training.  In International Conference on Learning Representations, 2022.
  * [17] Meng-Hao Guo, Cheng-Ze Lu, Qibin Hou, Zheng-Ning Liu, Ming-Ming Cheng, and Shi min Hu.  Segnext: Rethinking convolutional attention design for semantic segmentation.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  * [18] Kai Han, Yunhe Wang, Qi Tian, Jianyuan Guo, Chunjing Xu, and Chang Xu.  Ghostnet: More features from cheap operations.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1580–1589, 2020.
  * [19] Song Han, Huizi Mao, and William J Dally.  Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding.  In ICLR, 2016.
  * [20] Song Han, Jeff Pool, John Tran, and William Dally.  Learning both weights and connections for efficient neural network.  In NeurIPS, 2015.
  * [21] Yihui He, Ji Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han.  Amc: Automl for model compression and acceleration on mobile devices.  In ECCV, 2018.
  * [22] Yihui He, Xiangyu Zhang, and Jian Sun.  Channel pruning for accelerating very deep neural networks.  In ICCV, 2017.
  * [23] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  arXiv preprint arXiv:1503.02531, 2015.
  * [24] Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, et al.  Searching for mobilenetv3.  In ICCV, 2019.
  * [25] Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  arXiv preprint arXiv:1704.04861, 2017.
  * [26] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In International Conference on Machine Learning, pages 5156–5165. PMLR, 2020.
  * [27] Hanchao Li, Pengfei Xiong, Haoqiang Fan, and Jian Sun.  Dfanet: Deep feature aggregation for real-time semantic segmentation.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9522–9531, 2019.
  * [28] Zhuang Liu, Jianguo Li, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang.  Learning efficient convolutional networks through network slimming.  In ICCV, 2017.
  * [29] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10012–10022, 2021.
  * [30] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie.  A convnet for the 2020s.  In Proceedings of the IEEE conference on computer vision and pattern recognition, 2022.
  * [31] Jonathan Long, Evan Shelhamer, and Trevor Darrell.  Fully convolutional networks for semantic segmentation.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3431–3440, 2015.
  * [32] Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun.  Shufflenet v2: Practical guidelines for efficient cnn architecture design.  In ECCV, 2018.
  * [33] Sachin Mehta and Mohammad Rastegari.  Mobilevit: Light-weight, general-purpose, and mobile-friendly vision transformer.  In International Conference on Learning Representations, 2022.
  * [34] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  Advances in neural information processing systems, 32, 2019.
  * [35] Rudra PK Poudel, Stephan Liwicki, and Roberto Cipolla.  Fast-scnn: Fast semantic segmentation network.  arXiv preprint arXiv:1902.04502, 2019.
  * [36] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234–241. Springer, 2015.
  * [37] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In CVPR, 2018.
  * [38] Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li.  Efficient attention: Attention with linear complexities.  In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pages 3531–3539, 2021.
  * [39] Mingxing Tan and Quoc Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In ICML, 2019.
  * [40] Mingxing Tan and Quoc Le.  Efficientnetv2: Smaller models and faster training.  In International Conference on Machine Learning, pages 10096–10106. PMLR, 2021.
  * [41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In NeurIPS, 2017.
  * [42] Jingdong Wang, Ke Sun, Tianheng Cheng, Borui Jiang, Chaorui Deng, Yang Zhao, Dong Liu, Yadong Mu, Mingkui Tan, Xinggang Wang, et al.  Deep high-resolution representation learning for visual recognition.  IEEE transactions on pattern analysis and machine intelligence, 43(10):3349–3364, 2020.
  * [43] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity.  arXiv preprint arXiv:2006.04768, 2020.
  * [44] Tianzhe Wang, Kuan Wang, Han Cai, Ji Lin, Zhijian Liu, Hanrui Wang, Yujun Lin, and Song Han.  Apq: Joint search for network architecture, pruning and quantization policy.  In CVPR, 2020.
  * [45] Enze Xie, Wenhai Wang, Zhiding Yu, Anima Anandkumar, Jose M Alvarez, and Ping Luo.  Segformer: Simple and efficient design for semantic segmentation with transformers.  Advances in Neural Information Processing Systems, 34, 2021.
  * [46] Changqian Yu, Jingbo Wang, Chao Peng, Changxin Gao, Gang Yu, and Nong Sang.  Bisenet: Bilateral segmentation network for real-time semantic segmentation.  In Proceedings of the European conference on computer vision (ECCV), pages 325–341, 2018.
  * [47] Fisher Yu and Vladlen Koltun.  Multi-scale context aggregation by dilated convolutions.  arXiv preprint arXiv:1511.07122, 2015.
  * [48] Yuhui Yuan, Xilin Chen, and Jingdong Wang.  Object-contextual representations for semantic segmentation.  In European conference on computer vision, pages 173–190. Springer, 2020.
  * [49] Yuhui Yuan, Rao Fu, Lang Huang, Weihong Lin, Chao Zhang, Xilin Chen, and Jingdong Wang.  Hrformer: High-resolution vision transformer for dense predict.  Advances in Neural Information Processing Systems, 34, 2021.
  * [50] Wenqiang Zhang, Zilong Huang, Guozhong Luo, Tao Chen, Xinggang Wang, Wenyu Liu, Gang Yu, and Chunhua Shen.  Topformer: Token pyramid transformer for mobile semantic segmentation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12083–12093, 2022.
  * [51] Hengshuang Zhao, Xiaojuan Qi, Xiaoyong Shen, Jianping Shi, and Jiaya Jia.  Icnet for real-time semantic segmentation on high-resolution images.  In Proceedings of the European conference on computer vision (ECCV), pages 405–420, 2018.
  * [52] Hengshuang Zhao, Jianping Shi, Xiaojuan Qi, Xiaogang Wang, and Jiaya Jia.  Pyramid scene parsing network.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2881–2890, 2017.
  * [53] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba.  Scene parsing through ade20k dataset.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 633–641, 2017.
  * [54] Barret Zoph and Quoc V Le.  Neural architecture search with reinforcement learning.  In ICLR, 2017.