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 masked patch reconstructor below encodes every patch and reconstructs through a dense head; the SSL objective can also train other architectures.

Masked patch reconstruction

structure TorchLean.nn.models.ViT.MaskedPatchReconstructor.Config (d : ℕ) : Type

Configuration for a compact masked patch-transformer reconstructor.

The input/output contract is MAE-style: a masked channel/spatial tensor is mapped to one flattened reconstruction vector per batch position.

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

• encoder : EncoderConfig d

  Patch-transformer encoder configuration.

• reconstructionWidth : ℕ

  Number of reconstructed output coordinates.

def TorchLean.nn.models.ViT.MaskedPatchReconstructor.Config.validate {d : ℕ} (config : Config d) : Except String Unit

Validate both the encoder and decoder width before allocating either component.

@[reducible, inline]

abbrev TorchLean.nn.models.ViT.MaskedPatchReconstructor.Config.outputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Reconstruction output shape for the same batch shape as the input.

def TorchLean.nn.models.ViT.maskedPatchReconstructor {d : ℕ} (config : MaskedPatchReconstructor.Config d) (batchShape : Shape := []) : Builder (Sequential (config.encoder.inputShape batchShape) (config.outputShape batchShape))

Compact masked-patch image reconstructor.

Strided convolution embeds the patches, and the Transformer encoder processes every patch token. The encoded tokens are flattened together and passed to one dense reconstruction projection. Masked positions remain in the encoder sequence, so increasing the mask ratio does not reduce its token count. A separate Transformer decoder and token removal/restoration are not part of this compact architecture.

The masking objective is provided by TorchLean.ssl.BlockMAE.sample. 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.