Masked Autoencoder Objective Semantics

This module formalizes the finite patch/token core of a masked autoencoder (MAE):

• an input is a finite collection of patches Fin n → Patch;
• an encoder/decoder computation is abstracted to a reconstruction function;
• the objective is a sum of per-patch losses over the masked indices.

The formalization focuses on the semantic core. It captures the semantics that examples and future model helpers should preserve, while leaving ViT blocks, convolutional patch embeddings, and image IO in the executable API layer.

Paper anchor: “Masked Autoencoders Are Scalable Vision Learners” (He, Chen, Xie, Li, Dollár, Girshick, 2021), arXiv:2111.06377. The key objective-level fact we encode is that the reconstruction loss is taken over the masked patch set. Therefore the objective should not depend on an arbitrary ordering of masked patch indices; maeLoss_reverse is the small finite theorem capturing that property.

@[reducible, inline]

abbrev NN.MLTheory.SelfSupervised.PatchBatch (n : Nat) (Patch : Type) : Type

A finite patch collection.

def NN.MLTheory.SelfSupervised.reconstruct {n : Nat} {Patch Pred : Type} (decode : Fin n → Pred → Patch) (pred : Fin n → Pred) : PatchBatch n Patch

Reconstruct every patch using a reconstruction function.

def NN.MLTheory.SelfSupervised.ExactReconstruction {n : Nat} {Patch : Type} (x y : PatchBatch n Patch) : Prop

Exact reconstruction predicate for all patches.

def NN.MLTheory.SelfSupervised.maeLoss {n : Nat} {Patch Pred : Type} (maskedIdxs : Array (Fin n)) (target : PatchBatch n Patch) (pred : Fin n → Pred) (patchLoss : Patch → Pred → Nat) : Nat

MAE-style masked reconstruction loss over an explicit masked-index array.

The objective sums over the array, counting duplicate indices repeatedly. A set interpretation requires distinct indices; the theorems below cover reversal and concatenation of arrays.

@[simp]

theorem NN.MLTheory.SelfSupervised.maeLoss_nil {n : Nat} {Patch Pred : Type} (target : PatchBatch n Patch) (pred : Fin n → Pred) (patchLoss : Patch → Pred → Nat) : maeLoss #[] target pred patchLoss = 0

Reconstructing no patches costs nothing.

theorem NN.MLTheory.SelfSupervised.maeLoss_append {n : Nat} {Patch Pred : Type} (xs ys : Array (Fin n)) (target : PatchBatch n Patch) (pred : Fin n → Pred) (patchLoss : Patch → Pred → Nat) : maeLoss (xs ++ ys) target pred patchLoss = maeLoss xs target pred patchLoss + maeLoss ys target pred patchLoss

The masked-autoencoder loss is additive in the masked patch list, for the same reason.

theorem NN.MLTheory.SelfSupervised.maeLoss_reverse {n : Nat} {Patch Pred : Type} (idxs : Array (Fin n)) (target : PatchBatch n Patch) (pred : Fin n → Pred) (patchLoss : Patch → Pred → Nat) : maeLoss idxs.reverse target pred patchLoss = maeLoss idxs target pred patchLoss

The MAE loss is invariant under reversing the masked-index array. Multiplicities are preserved; this does not identify arrays that differ by duplicate indices.

theorem NN.MLTheory.SelfSupervised.maeLoss_eq_zero_of_patch_losses_zero {n : Nat} {Patch Pred : Type} (idxs : Array (Fin n)) (target : PatchBatch n Patch) (pred : Fin n → Pred) (patchLoss : Patch → Pred → Nat) (h : ∀ (i : Fin n), i ∈ idxs → patchLoss (target i) (pred i) = 0) : maeLoss idxs target pred patchLoss = 0

If every selected patch has zero reconstruction loss, the masked MAE loss is zero.

theorem NN.MLTheory.SelfSupervised.exactReconstruction_identity {n : Nat} {Patch : Type} (x : PatchBatch n Patch) : ExactReconstruction x (reconstruct (fun (x : Fin n) (p : Patch) => p) x)

Reconstructing with the identity decoder/prediction is exact.