Self-Supervised Model Constructors

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

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

ViT-MAE

structure TorchLean.nn.models.VitMaeConfig (d : ℕ) : Type

Configuration for a compact masked patch-transformer reconstructor.

The input/output contract is MAE-style:

• input: a masked tensor, (batch, channels, spatial...);
• output: a flattened reconstruction vector, N×reconDim.

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

• encoder : VitConfig d

  Patch-transformer encoder configuration.

• reconDim : ℕ

  Number of reconstructed output coordinates.

def TorchLean.nn.models.vitMaeInShape {d : ℕ} (cfg : VitMaeConfig d) : Shape

Masked input shape.

def TorchLean.nn.models.vitMaeOutShape {d : ℕ} (cfg : VitMaeConfig d) : Shape

Reconstruction-vector output shape.

def TorchLean.nn.models.VitMaeConfig.seqLen {d : ℕ} (cfg : VitMaeConfig d) : ℕ

Number of patch tokens produced by the ViT-MAE patch embedding.

def TorchLean.nn.models.VitMaeConfig.flatDim {d : ℕ} (cfg : VitMaeConfig d) : ℕ

Flattened encoded-token representation size before the MAE decoder head.

def TorchLean.nn.models.vitMaskedAutoencoder {d : ℕ} (cfg : VitMaeConfig d) (h_inC : cfg.encoder.inChannels ≠ 0 := by decide) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.encoder.patch.outChannels ≠ 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 TorchLean.ssl.blockMaeSample. Its axis policy is independent of the model architecture and spatial rank, so this constructor uses the same checked operation as signal, volume, and higher-dimensional masked-prediction models.

def TorchLean.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 TorchLean.ssl.vectorMaeSample or TorchLean.ssl.tensorPrefixMaeSample, which mask the input while keeping the original tensor content as the target.