Reconstruction autoencoder and denoising risk

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