Masked autoencoder pretraining

Let $S=(x^{(1)},\dots ,x^{(n)})$ be a finite image training set with $n\ge 1$ on a fixed grid with $M,q\ge 1$ (that is, $M$ patches of $q$ pixel values each). Choose a visible count $k$ with $1\le k<M$, a finite update count $J\ge 0$, differentiable encoder and decoder families parameterized by $\theta =(\varphi ,\psi ,m)\in \mathrm{\Theta }\subseteq {\mathbb{R}}^p$, an initial parameter vector ${\theta }_0\in \mathrm{\Theta }$, an optimizer state set $\mathcal{S}$, and an initial optimizer state $s_0\in \mathcal{S}$. Choose either the original-pixel target $T(u)=u$ or, provided every patch of every training image has positive standard deviation, the normalized target $T(u)=(u-\bar{u}\mathbf{1})\mathrm{/}s(u)$, where $\bar{u}=q^{-1}{\sum }_{r=1}^q{u}_r$ and $s(u)=(q^{-1}{\sum }_{r=1}^q(u_r-\bar{u}{)}^2{)}^{1\mathrm{/}2}$. For each $V\subseteq [M]$ with $\mathrm{\mid }V\mathrm{\mid }=k$, the encoder is a map from $V$ and the visible patch values to one encoded token at each position in $V$; it receives no mask tokens. The decoder is a map from the ordered length-$M$ sequence of encoded visible tokens and one shared learned vector $m$ at every position in $[M]\setminus V$ to predictions ${\widehat{x}}_j\in {\mathbb{R}}^q$ at every patch position. Position $j$ is supplied to the maps as its index. Require that the masked mean squared error $\frac{1}{(M-k)q}{\sum }_{j\in [M]\setminus V}\mathrm{\parallel }{\widehat{x}}_j-T(x_j){\mathrm{\parallel }}_2^2$ has a finite parameter gradient at every iterate reached by the procedure.

An optimizer update schedule is a supplied finite family of maps $U_t:(\mathrm{\Theta }\times \mathcal{S})\times {\mathbb{R}}^p\to \mathrm{\Theta }\times \mathcal{S}$ for $0\le t<J$, each defined on its entire displayed domain. When $\mathrm{\Theta }={\mathbb{R}}^p$, plain stochastic gradient descent with step sizes ${\eta }_t>0$ is the special case $U_t((\theta ,s),g)=(\theta -{\eta }_{t}g,s)$.

For each $t=0,\dots ,J-1$, perform the following steps.

1. Sample $I_t$ uniformly from $\{1,\dots ,n\}$ and, independently, sample $V_t$ uniformly from ${\mathcal{V}}_k=\{V\subseteq [M]:\mathrm{\mid }V\mathrm{\mid }=k\}$. Sampling is with replacement across different update steps.
2. Apply the encoder to the visible patch values at positions in $V_t$. Restore their original order, insert the shared learned vector $m$ at every position in $[M]\setminus V_t$, and supply each position index to the decoder. Apply the decoder to this full ordered sequence to obtain ${\widehat{x}}_{t,j}\in {\mathbb{R}}^q$ for every $j\in [M]$.
3. Compute $L_t=((M-k)q{)}^{-1}{\sum }_{j\in [M]\setminus V_t}\mathrm{\parallel }{\widehat{x}}_{t,j}-T(x_j^{(I_t)}){\mathrm{\parallel }}_2^2$, compute $g_t={\mathrm{\nabla }}_{\theta }{L}_t$, and update $({\theta }_{t+1},s_{t+1})=U_t(({\theta }_t,s_t),g_t)$.

Return the trained encoder $E_{{\varphi }_J}$ and parameter vector ${\theta }_J$. The decoder is used during pretraining and may be discarded afterward. The mask ratio is $\rho =(M-k)\mathrm{/}M$. The procedure defines finite stochastic updates only; it does not assert that the updates converge to a minimizer or that the returned representation is semantically sufficient.

The diagram records the data flow. Solid arrows pass data or predictions; the bottom comparison uses only positions in $[M]\setminus V$. The definitions and numbered instructions specify the algorithm completely, and no proof or property relies on the diagram.

diagram failed to render:
\begin{tikzpicture}[>=stealth,scale=0.78,transform shape]
  \node[draw,rounded corners,align=center,minimum width=2.1cm,minimum height=0.8cm] (input) at (0,0) {Input patch\\sequence};
  \node[draw,rounded corners,align=center,minimum width=2.1cm,minimum height=0.8cm] (mask) at (2.7,0) {Uniformly sample\\visible positions};
  \node[draw,rounded corners,align=center,minimum width=2.1cm,minimum height=0.8cm] (enc) at (5.4,0) {Encoder sees\\visible patches only};
  \node[draw,rounded corners,align=center,minimum width=2.4cm,minimum height=0.8cm] (restore) at (8.3,0) {Insert learned tokens\\at missing positions};
  \node[draw,rounded corners,align=center,minimum width=2.1cm,minimum height=0.8cm] (dec) at (11.2,0) {Decoder};
  \node[draw,rounded corners,align=center,minimum width=2.3cm,minimum height=0.8cm] (pred) at (14,0) {Predictions at\\all positions};
  \node[draw,rounded corners,align=center,minimum width=3.0cm,minimum height=0.8cm] (loss) at (8.3,-1.7) {Masked-patch MSE\\only};
  \node[draw,rounded corners,align=center,minimum width=2.8cm,minimum height=0.8cm] (target) at (3.4,-1.7) {Clean targets at\\missing positions};
  \draw[->] (input) -- (mask);
  \draw[->] (mask) -- node[above] {visible subset} (enc);
  \draw[->] (enc) -- (restore);
  \draw[->] (restore) -- (dec);
  \draw[->] (dec) -- (pred);
  \draw[->] (pred) |- (loss);
  \draw[->] (target) -- (loss);
\end{tikzpicture}

Reconstruction autoencoder and denoising risk: fixes the image-patch domain, visible/masked convention, raw and normalized target maps, and finite masked-MSE objective.

Finite feedforward computation graphs and layered MLPs: defines the finite feedforward computation graph used for encoder/decoder network notation.

He et al., §§3.1–3.3: mask sampling, visible-only encoder, mask-token decoder, and masked-pixel objective

He et al., Appendix A.1, pretraining setting: AdamW and reported hyperparameters

Goodfellow et al., §14.5: corrupt-input/clean-target denoising procedure

Vincent et al., §2.3, equation (5): sampling a corruption and updating toward reconstruction of the clean example