Self-Supervised Model Constructors

Most SSL machinery belongs in NN.API.ssl: masks, tensor-to-training-sample transforms, and objective-facing helpers should work with any compatible model.

This file keeps only architecture-level conveniences. The compact MAE constructor below is useful for examples and smoke tests, but the SSL idea itself is not tied to this model.

Compatibility aliases

These names preserve the model API while the model-independent path lives under NN.API.ssl.

def NN.API.nn.models.vectorMaeHiddenMask (cfg : VectorGenerativeConfig) (period offset : ℕ) : MLTheory.SelfSupervised.Mask cfg.dataDim

def NN.API.nn.models.vectorMaeMask (cfg : VectorGenerativeConfig) (period offset : ℕ) (x : Spec.Tensor Float (vectorDataShape cfg)) : Spec.Tensor Float (vectorDataShape cfg)

def NN.API.nn.models.vectorMaeSample (cfg : VectorGenerativeConfig) (period offset : ℕ) (x : Spec.Tensor Float (vectorDataShape cfg)) : sample.Supervised Float (vectorDataShape cfg) (vectorDataShape cfg)

def NN.API.nn.models.tensorPrefixMaeSample {source : Shape} (cfg : VectorGenerativeConfig) (hData : cfg.dataDim ≤ Shape.size source) (period offset : ℕ) (x : Spec.Tensor Float (Spec.Shape.dim cfg.batch source)) : sample.Supervised Float (vectorDataShape cfg) (vectorDataShape cfg)

Alias for NN.API.ssl.tensorPrefixMaeSample; prefer the ssl namespace in new code.

ViT-MAE

structure NN.API.nn.models.VitMaeConfig : Type

Configuration for a compact ViT-MAE image reconstructor.

The input/output contract is MAE-style:

• input: a masked image tensor, N×C×H×W;
• output: a flattened reconstruction vector, N×reconDim.

reconDim can be the full image size (C*H*W) or a prefix for faster experiments.

• batch : ℕ
• inC : ℕ
• inH : ℕ
• inW : ℕ
• patchH : ℕ
• patchW : ℕ
• stride : ℕ
• padding : ℕ
• dModel : ℕ
• reconDim : ℕ
• numHeads : ℕ
• headDim : ℕ
• ffnHidden : ℕ

@[implicit_reducible]

instance NN.API.nn.models.instReprVitMaeConfig : Repr VitMaeConfig

def NN.API.nn.models.instReprVitMaeConfig.repr : VitMaeConfig → ℕ → Std.Format

def NN.API.nn.models.VitMaeConfig.toVitConfig (cfg : VitMaeConfig) : VitConfig

@[reducible, inline]

abbrev NN.API.nn.models.vitMaeInShape (cfg : VitMaeConfig) : Shape

@[reducible, inline]

abbrev NN.API.nn.models.vitMaeOutShape (cfg : VitMaeConfig) : Shape

def NN.API.nn.models.VitMaeConfig.seqLen (cfg : VitMaeConfig) : ℕ

def NN.API.nn.models.VitMaeConfig.flatDim (cfg : VitMaeConfig) : ℕ

def NN.API.nn.models.vitMaskedAutoencoder (cfg : VitMaeConfig) (h_inC : cfg.inC ≠ 0 := by decide) (h_patchH : cfg.patchH ≠ 0 := by decide) (h_patchW : cfg.patchW ≠ 0 := by decide) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.dModel ≠ 0 := by decide) : M (Sequential (vitMaeInShape cfg) (vitMaeOutShape cfg))

Compact ViT-MAE image reconstructor.

This is a real image/patch transformer path:

1. patch embedding by strided convolution,
2. tokenization to N×numPatches×dModel,
3. one transformer encoder block,
4. a linear pixel decoder from encoded patch tokens to a reconstruction vector.

The masking objective is provided by NN.API.ssl.imagePatchMaeSample, so any image model with this input/output shape can use the same SSL training sample.

def NN.API.nn.models.vectorMaskedAutoencoder (cfg : VectorGenerativeConfig) : M (Sequential (vectorDataShape cfg) (vectorDataShape cfg))

Compact vector masked autoencoder.

Architecturally this reuses the vector autoencoder body; the self-supervised part is in NN.API.ssl.vectorMaeSample or NN.API.ssl.tensorPrefixMaeSample, which mask the input while keeping the original tensor content as the target.