Masking primitives for self-supervised objectives

This file gives a small finite-index vocabulary for masked prediction objectives. It stays 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 definitions make MAE/JEPA-style objectives precise enough for local invariants before they are connected to larger executable models.

@[reducible, inline]

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

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

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

Proposition stating that index i is selected by the Boolean mask m.

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

The all-visible/all-target mask.

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

Mask selecting no positions.

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

Pointwise Boolean complement of a mask.

@[simp]

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

Every index is selected by the full mask.

@[simp]

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

No index is selected by the empty mask.

@[simp]

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

Complementing a mask negates selection pointwise. This supports a context/target partition when the caller chooses complementary masks; the objective definitions do not enforce that choice.

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

Generic masked loss over an explicit array of selected indices.

The scalar loss is Nat, so this models natural-valued scores such as quantized patch losses. Each array occurrence contributes once; duplicate indices are not removed.

@[simp]

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

An empty index array contributes no loss.

@[simp]

theorem NN.MLTheory.SelfSupervised.maskedLoss_push {n : Nat} (idxs : Array (Fin n)) (i : Fin n) (perPatchLoss : Fin n → Nat) : maskedLoss (idxs.push i) perPatchLoss = maskedLoss idxs perPatchLoss + perPatchLoss i

Pushing one index adds that patch's loss.

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

Masked loss is additive in the index array.

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

Masked loss does not depend on the order of the indices, so a shuffled mask scores the same.

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

Zero per-patch loss at every selected index gives zero total loss. This theorem states only that direction; it does not assume the per-patch score characterizes perfect reconstruction.