Masking primitives for self-supervised objectives

This file gives a small finite-index vocabulary for masked prediction objectives. It is deliberately independent of any particular image or transformer implementation: a patch/token collection is just Fin n → α, and a mask is a Boolean predicate on Fin n.

The goal is to make MAE/JEPA-style objectives precise enough to prove local invariants before we connect them to larger executable models.

@[reducible, inline]

abbrev NN.MLTheory.SelfSupervised.Mask (n : ℕ) : Type

A finite mask over n patches/tokens. true means the index is selected.

def NN.MLTheory.SelfSupervised.selected {n : ℕ} (m : Mask n) (i : Fin n) : Prop

Predicate form of a mask.

def NN.MLTheory.SelfSupervised.allMask (n : ℕ) : Mask n

The all-visible/all-target mask.

def NN.MLTheory.SelfSupervised.emptyMask (n : ℕ) : Mask n

The empty mask.

def NN.MLTheory.SelfSupervised.complement {n : ℕ} (m : Mask n) : Mask n

Complement a mask.

@[simp]

theorem NN.MLTheory.SelfSupervised.allMask_selected {n : ℕ} (i : Fin n) : selected (allMask n) i

@[simp]

theorem NN.MLTheory.SelfSupervised.emptyMask_not_selected {n : ℕ} (i : Fin n) : ¬selected (emptyMask n) i

@[simp]

theorem NN.MLTheory.SelfSupervised.complement_selected_iff {n : ℕ} (m : Mask n) (i : Fin n) : selected (complement m) i ↔ ¬selected m i

def NN.MLTheory.SelfSupervised.maskedLoss {n : ℕ} (idxs : List (Fin n)) (perPatchLoss : Fin n → ℕ) : ℕ

Generic masked loss over an explicit list of selected indices.

For theory files we keep the scalar loss as Nat; concrete runtime losses can instantiate this with squared-error bins, quantized patch losses, or any other executable per-patch score.

@[simp]

theorem NN.MLTheory.SelfSupervised.maskedLoss_nil {n : ℕ} (perPatchLoss : Fin n → ℕ) : maskedLoss [] perPatchLoss = 0

@[simp]

theorem NN.MLTheory.SelfSupervised.maskedLoss_cons {n : ℕ} (i : Fin n) (idxs : List (Fin n)) (perPatchLoss : Fin n → ℕ) : maskedLoss (i :: idxs) perPatchLoss = perPatchLoss i + maskedLoss idxs perPatchLoss

theorem NN.MLTheory.SelfSupervised.maskedLoss_append {n : ℕ} (xs ys : List (Fin n)) (perPatchLoss : Fin n → ℕ) : maskedLoss (xs ++ ys) perPatchLoss = maskedLoss xs perPatchLoss + maskedLoss ys perPatchLoss

theorem NN.MLTheory.SelfSupervised.maskedLoss_reverse {n : ℕ} (idxs : List (Fin n)) (perPatchLoss : Fin n → ℕ) : maskedLoss idxs.reverse perPatchLoss = maskedLoss idxs perPatchLoss

theorem NN.MLTheory.SelfSupervised.maskedLoss_eq_zero_of_all_zero {n : ℕ} (idxs : List (Fin n)) (perPatchLoss : Fin n → ℕ) (h : ∀ (i : Fin n), i ∈ idxs → perPatchLoss i = 0) : maskedLoss idxs perPatchLoss = 0