Reconstruction and masked autoencoders

Fix integers $M,q\ge 1$, set $[M]=\{1,\dots ,M\}$, and write the image space on a fixed patch grid as $\mathcal{X}=({\mathbb{R}}^q{)}^M$. A measurable space is a set with a specified sigma-algebra; equip $\mathcal{X}$ and a finite-dimensional latent space $\mathcal{Z}$ with such structures. An image is $x=(x_1,\dots ,x_M)\in \mathcal{X}$, with $q$ real pixel values in each of its $M$ patches. Let $S=(x^{(1)},\dots ,x^{(n)})$ be a finite training set with $n\ge 1$. For $u\in {\mathbb{R}}^q$, set $\mathrm{\parallel }u{\mathrm{\parallel }}_2^2={\sum }_{r=1}^q{u}_r^2$. For $V\subseteq [M]$ let ${\mathcal{Z}}^V=\{(z_j{)}_{j\in V}:z_j\in \mathcal{Z}\}$. For unmasked reconstruction, the encoder and decoder are deterministic measurable maps $E_{\varphi }:\mathcal{X}\to \mathcal{Z}$ and $D_{\psi }:\mathcal{Z}\to \mathcal{X}$. For masked reconstruction, the encoder is a measurable deterministic map $(V,(x_j{)}_{j\in V})\mapsto (z_j{)}_{j\in V}\in {\mathcal{Z}}^V$, and the decoder maps the ordered length-$M$ sequence of encoded visible tokens and mask tokens to an element of $\mathcal{X}$. Each mask token is one learned vector $m\in \mathcal{Z}$ reused at every masked position. All maps take finite real values; when implemented by feedforward neural networks, each is a finite computation graph evaluated in the layered order of its vertices.

Let $X:\mathrm{\Omega }\to \mathcal{X}$ be a measurable random image on a probability space. Assume that the displayed squared-loss random variables are measurable. The population reconstruction risk is $R_{\mathrm{r}\mathrm{e}\mathrm{c}}(\varphi ,\psi ;X):=\mathbb{E}[\frac{1}{Mq}{\sum }_{j=1}^M\mathrm{\parallel }{D}_{\psi }(E_{\varphi }(X){)}_j-X_j{\mathrm{\parallel }}_2^2]\in [0,+\mathrm{\infty }]\mathrm{.}$ A stochastic denoising corruption is any measurable random image $\tilde{X}:\mathrm{\Omega }\to \mathcal{X}$ on the same probability space as the clean image $X$; their dependence may be arbitrary, with no restriction on support. The population denoising risk reconstructs the clean image: $R_{\mathrm{d}\mathrm{e}\mathrm{n}}(\varphi ,\psi ;X,\tilde{X}):=\mathbb{E}[\frac{1}{Mq}{\sum }_{j=1}^M\mathrm{\parallel }{D}_{\psi }(E_{\varphi }(\tilde{X}){)}_j-X_j{\mathrm{\parallel }}_2^2]\in [0,+\mathrm{\infty }]\mathrm{.}$ These expectations use the extended integral of a nonnegative measurable random variable, so they may be infinite.

The empirical reconstruction risk on uncorrupted images is $R_{\mathrm{r}\mathrm{e}\mathrm{c}}^{\mathrm{e}\mathrm{m}\mathrm{p}}(\varphi ,\psi ;S):=\frac{1}{nMq}{\sum }_{i=1}^n{\sum }_{j=1}^M\mathrm{\parallel }{D}_{\psi }(E_{\varphi }(x^{(i)}){)}_j-x_j^{(i)}{\mathrm{\parallel }}_2^2\mathrm{.}$

A finite denoising corruption law assigns to each training image $x^{(i)}$ a finite set $C_i$ of corrupted images and weights $c_i(\tilde{x})\ge 0$ with ${\sum }_{\tilde{x}\in C_i}{c}_i(\tilde{x})=1$. The empirical denoising risk is the finite-support specialization of the population risk: $R_{\mathrm{d}\mathrm{e}\mathrm{n}}^{\mathrm{e}\mathrm{m}\mathrm{p}}(\varphi ,\psi ;S,C):=\frac{1}{n}{\sum }_{i=1}^n{\sum }_{\tilde{x}\in C_i}{c}_i(\tilde{x})\frac{1}{Mq}{\sum }_{j=1}^M\mathrm{\parallel }{D}_{\psi }(E_{\varphi }(\tilde{x}){)}_j-x_j^{(i)}{\mathrm{\parallel }}_2^2\mathrm{.}$

For masked autoencoding, fix an integer $k$ with $1\le k<M$ and let ${\mathcal{V}}_k=\{V\subseteq [M]:\mathrm{\mid }V\mathrm{\mid }=k\}$, a set with cardinality $\left(\genfrac{}{}{0px}{}{M}{k}\right)$. A mask $V$ is sampled uniformly from ${\mathcal{V}}_k$; the patches in $V$ are visible and the patches in $[M]\setminus V$ are masked. The encoder receives only the visible patches, with their positions, and no mask tokens. The decoder receives the encoded visible patches and a shared learned mask token at each missing position, with position information for all $M$ slots, and outputs predictions ${\widehat{x}}_j(\varphi ,\psi ;x,V)\in {\mathbb{R}}^q$ for every patch position.

The original-pixel target map is $T_{\mathrm{r}\mathrm{a}\mathrm{w}}(u)=u$. The optional normalized-pixel target is defined for $u=(u_1,\dots ,u_q)$ with positive standard deviation by $\bar{u}:=\frac{1}{q}{\sum }_{r=1}^q{u}_r,s(u):={(\frac{1}{q}{\sum }_{r=1}^q(u_r-\bar{u}{)}^2)}^{1\mathrm{/}2}>0,T_{\mathrm{n}\mathrm{o}\mathrm{r}\mathrm{m}}(u):=\frac{u-\bar{u}\mathbf{1}}{s(u)},$ where $\mathbf{1}=(1,\dots ,1)\in {\mathbb{R}}^q$. The target choice $T$ is one of these maps. Define ${\mathcal{D}}_{\mathrm{n}\mathrm{o}\mathrm{r}\mathrm{m}}:=\{x\in \mathcal{X}:s(x_j)>0\text{ for every }j\in [M]\}\mathrm{.}$ For $T=T_{\mathrm{n}\mathrm{o}\mathrm{r}\mathrm{m}}$, the masked-objective formulas below are defined only on ${\mathcal{D}}_{\mathrm{n}\mathrm{o}\mathrm{r}\mathrm{m}}$, so a finite empirical objective requires $x^{(i)}\in {\mathcal{D}}_{\mathrm{n}\mathrm{o}\mathrm{r}\mathrm{m}}$ for every $i$. For a population objective, require the event $A:=\{X\in {\mathcal{D}}_{\mathrm{n}\mathrm{o}\mathrm{r}\mathrm{m}}\}$ to be measurable with $\mathbb{P}(A)=1$; interpret each loss integrand by its displayed formula on $A$ and by $0$ on $A^c$. Since $A^c$ is null, any other measurable extension there gives the same expectation. The raw target has no such restriction. The per-example masked-patch mean squared error, its population objective, and its finite empirical objective are ${\mathrm{\ell }}_{\mathrm{M}\mathrm{A}\mathrm{E}}(\varphi ,\psi ;x,V,T):=\frac{1}{(M-k)q}{\sum }_{j\in [M]\setminus V}\mathrm{\parallel }{\widehat{x}}_j(\varphi ,\psi ;x,V)-T(x_j){\mathrm{\parallel }}_2^2,$ $R_{\mathrm{M}\mathrm{A}\mathrm{E}}(\varphi ,\psi ;X,k,T):=\frac{1}{\left(\genfrac{}{}{0px}{}{M}{k}\right)}{\sum }_{V\in {\mathcal{V}}_k}\mathbb{E}[{\mathrm{\ell }}_{\mathrm{M}\mathrm{A}\mathrm{E}}(\varphi ,\psi ;X,V,T)],$ $R_{\mathrm{M}\mathrm{A}\mathrm{E}}^{\mathrm{e}\mathrm{m}\mathrm{p}}(\varphi ,\psi ;S,k,T):=\frac{1}{n\left(\genfrac{}{}{0px}{}{M}{k}\right)}{\sum }_{i=1}^n{\sum }_{V\in {\mathcal{V}}_k}{\mathrm{\ell }}_{\mathrm{M}\mathrm{A}\mathrm{E}}(\varphi ,\psi ;x^{(i)},V,T)\mathrm{.}$ The population objective averages the expectation over the image law for each fixed mask; equivalently, it is the joint expectation for an independent uniform mask. The finite empirical objective averages over every training image and mask.

Because the encoder input and decoder output depend on $V$ and the loss is evaluated only on $[M]\setminus V$, this is an expectation of a mask-dependent loss. It is not defined as an estimator of the separate full-image risk $R_{\mathrm{r}\mathrm{e}\mathrm{c}}$, and no equality or unbiasedness between these objectives is asserted. Neither loss alone guarantees that a learned code retains distinctions needed for a downstream label (semantic sufficiency), uniquely determines an underlying data-generating parameter (identifiability), or improves performance on a specified later task (downstream feature quality).

Finite feedforward computation graphs and layered MLPs: a feedforward network is a finite acyclic computation graph whose node values are evaluated in a specified order.

Expectation of a real random variable on a finite probability space: expectation on a finite probability space is the weighted finite sum.

Probability measures and probability spaces: a probability measure has total mass one and defines a probability space.

Random elements and real random variables: a random image is a measurable map from a probability space into the image space.

A measurable function between measurable spaces: a map is measurable when preimages of measurable sets are measurable.

Measurable spaces and measurable sets: a measurable space consists of a set and a sigma-algebra of its measurable subsets.

Sigma-algebras: a sigma-algebra is closed under complements and countable unions.

Expectation of a nonnegative or integrable random variable: the expectation of a nonnegative measurable loss is its extended-valued integral.

He et al., §§3.1–3.3: uniform random patch masking, visible-only encoder, mask-token decoder, masked-pixel MSE, and normalized-pixel target

Goodfellow et al., §14.5, pp. 507–508: denoising corruption and clean-target reconstruction objective

Vincent et al., §§2.2–2.3, equations (1), (3), and (5): reconstruction risk and stochastic corruption with clean target

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.

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