Masked Autoencoder CIFAR Example

This is the compact ViT-MAE-style training path in TorchLean.

The data path is explicit:

1. load real CIFAR-10 .npy arrays through Data;
2. take a typed image batch with shape [batch, channels, height, width];
3. hide deterministic spatial blocks with ssl.BlockMAE.sample;
4. run a ViT encoder over patch tokens;
5. train a decoder head to reconstruct the original image vector.

The architecture uses one transformer encoder block and a linear pixel decoder rather than a large asymmetric MAE decoder. It exercises image patch masking, patch embedding, transformer tokens, and reconstruction of the original image.

def NN.Examples.Models.Generative.Mae.exeName : String

Command name used in error messages and CLI output.

def NN.Examples.Models.Generative.Mae.defaultLogPath : System.FilePath

Default JSON loss-curve path for this command.

def NN.Examples.Models.Generative.Mae.batchSize : ℕ

CIFAR minibatch size used by the typed MAE command.

def NN.Examples.Models.Generative.Mae.inputChannels : ℕ

Number of CIFAR image channels.

def NN.Examples.Models.Generative.Mae.cropHeight : ℕ

Cropped CIFAR image height for the compact runnable example.

def NN.Examples.Models.Generative.Mae.cropWidth : ℕ

Cropped CIFAR image width for the compact runnable example.

def NN.Examples.Models.Generative.Mae.patchHeight : ℕ

Patch height for the image-to-token projection.

def NN.Examples.Models.Generative.Mae.patchWidth : ℕ

Patch width for the image-to-token projection.

def NN.Examples.Models.Generative.Mae.stride : ℕ

Patch stride; equal to patch size here, so patches do not overlap.

def NN.Examples.Models.Generative.Mae.padding : ℕ

Zero padding around the image before patch extraction.

def NN.Examples.Models.Generative.Mae.modelWidth : ℕ

Width of each patch token after projection into the encoder stream.

def NN.Examples.Models.Generative.Mae.attentionHeads : ℕ

Number of self-attention heads in the compact ViT encoder.

def NN.Examples.Models.Generative.Mae.attentionHeadWidth : ℕ

Per-head attention width; attentionHeads * attentionHeadWidth = modelWidth.

def NN.Examples.Models.Generative.Mae.feedForwardWidth : ℕ

Hidden width of the feed-forward block inside the encoder.

def NN.Examples.Models.Generative.Mae.reconstructionWidth : ℕ

Number of flattened crop values predicted by the decoder head.

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Mae.modelConfig : TorchLean.nn.models.ViT.MaskedPatchReconstructor.Config 2

Small ViT-MAE configuration.

The command crops CIFAR images to 4×4, divides them into four 2×2 patches, and reconstructs the entire flattened crop. This keeps the command quick while making masking and attention genuinely operate across multiple patch positions.

def NN.Examples.Models.Generative.Mae.maskPeriod : ℕ

Hide one patch-index class every four patch positions.

The image remains an image tensor; the mask zeros whole patch regions before patch embedding.

def NN.Examples.Models.Generative.Mae.maskOffset : ℕ

Phase of the deterministic patch mask. Changing this selects a different patch-index class.

def NN.Examples.Models.Generative.Mae.maskBlocks : TorchLean.Tensor (Option ℕ) [3]

Per-axis block policy used by both the input mask and the reconstruction loss.

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Mae.batch : Spec.Shape

Batch shape used by this training example.

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Mae.input : Spec.Shape

Input shape: a real batched CIFAR image tensor.

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Mae.output : Spec.Shape

Output shape: flattened image reconstruction.

def NN.Examples.Models.Generative.Mae.model : TorchLean.nn.Builder (TorchLean.nn.Sequential input output)

Construct the trainable model.

The architecture lives in the public self-supervised model API; this example only chooses a config, loads data, and trains it.

def NN.Examples.Models.Generative.Mae.maskedAutoencoderSample (b : TorchLean.Sample.Supervised Float [batchSize, modelConfig.encoder.inputChannels, cropHeight, cropWidth] [batchSize, RealData.cifarClasses]) : Except String (TorchLean.Sample.Supervised Float input output)

Turn a typed CIFAR image batch into the compact MAE training sample.

The input stays an image tensor with some patches zeroed out. The target is the original image flattened to a vector because the current decoder head predicts a batched matrix.

def NN.Examples.Models.Generative.Mae.lossWeights {α : Type} [TorchLean.Storage α] [Context α] : TorchLean.Tensor α output

Normalized reconstruction weights repeated across the batch.

Only coordinates hidden in the model input have nonzero weight, and those weights sum to one across each batch because this compact example uses a singleton batch.

def NN.Examples.Models.Generative.Mae.hiddenReconstructionLoss {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (prediction target : TorchLean.Runtime.ValueRef m α output) : m (TorchLean.Runtime.ValueRef m α [])

Mean squared reconstruction error over hidden coordinates only.

def NN.Examples.Models.Generative.Mae.hiddenReconstructionLossProgram {α : Type} [TorchLean.Storage α] [Context α] : Runtime.Autograd.Model.Program α [output, output] []

Runtime-polymorphic form of hiddenReconstructionLoss for Trainer.

def NN.Examples.Models.Generative.Mae.data (flags : RealData.CifarModelTrainFlags) : TorchLean.Trainer.Dataset input output

Public singleton dataset for masked-image reconstruction on one real CIFAR batch.

Like the compact vector generative examples, the sample itself is loaded as Float from the real data boundary, then cast into the runtime-selected arithmetic representation by the public dataset constructor.

def NN.Examples.Models.Generative.Mae.train (runtime : TorchLean.Runtime.Config) (flags : RealData.CifarModelTrainFlags) : IO (TorchLean.Trainer.Result input output)

Train the compact MAE model with the public Trainer surface.

def NN.Examples.Models.Generative.Mae.main (args : List String) : IO UInt32

CLI entrypoint.

Useful flags:

• --device cuda runs the public trainer on the CUDA runtime.
• --steps <n> controls optimization steps.
• --x <path> --y <path> selects custom CIFAR-style .npy arrays.
• --log <path> writes the standard TorchLean training log JSON.