Arbitrary-Rank Block Masks

Masked prediction is not intrinsically an image operation. A model may hide intervals in a signal, rectangles in an image, cuboids in a volume, or blocks in a higher-dimensional simulation field. This module describes a mask by a rank-indexed policy tensor. The extents come from the input tensor's type, while blocks : Tensor (Option Nat) [d] selects the axes that form a block grid.

none means that an axis does not participate in the block index. some k groups that axis into consecutive blocks of width k. The selected block-grid coordinates are flattened in row-major order, and one congruence class modulo period is hidden. A zero period, a zero block width, a rank mismatch, or an out-of-bounds coordinate hides nothing.

The executable mask and its coordinate theorems use the same finite predicate. Consequently the runtime training sample cannot silently use a different patch convention from the one stated in Lean.

def TorchLean.ssl.BlockMask.Internal.blockIndex (shape : Shape) (policies : List (Option ℕ)) (coordinates : List ℕ) (flatIndex : ℕ) (used : Bool) : Option ℕ

Flat index of the block a coordinate falls in, or none when the coordinate is out of range or no axis is blocked at all.

The walk goes over shape and policy together. An axis with none contributes nothing; an axis with some blockSize contributes its block coordinate as one digit of a mixed-radix number whose radix is the number of blocks along that axis. used is what distinguishes "block 0" from "nothing was blocked", which the caller has to tell apart.

def TorchLean.ssl.BlockMask.Internal.blockHidden (shape : Shape) (blocks : List (Option ℕ)) (period offset : ℕ) (coordinate : List ℕ) : Bool

Is the block containing this coordinate hidden by the mask with the given period and offset?

Hiding every period-th block rather than sampling at random is a deliberate choice: a mask is then reproducible from two numbers, which is what makes the self-supervised examples in the guide comparable across machines and across runs.

def TorchLean.ssl.BlockMask.Internal.apply (shape : Shape) (blocks : List (Option ℕ)) (period offset : ℕ) (coordinatePrefix : List ℕ) (remainingShape : Shape) : Tensor Float remainingShape → Tensor Float remainingShape

Zero out the hidden blocks of a tensor, recursing over the leading axes and collecting the coordinate prefix on the way down.

The result is rebuilt with stackLeading rather than written in place, so a masked tensor is an ordinary value and the mask cannot depend on evaluation order.

theorem TorchLean.ssl.BlockMask.Internal.apply_scalar_at (shape : Shape) (blocks : List (Option ℕ)) (period offset : ℕ) (coordinatePrefix : List ℕ) (remainingShape : Shape) (x : Tensor Float remainingShape) (coordinates : List ℕ) : Spec.getSpec (apply shape blocks period offset coordinatePrefix remainingShape x) coordinates = Option.map (fun (value : Float) => if blockHidden shape blocks period offset (coordinatePrefix ++ coordinates) = true then 0.0 else value) (Spec.getSpec x coordinates)

Reading one coordinate out of a masked tensor is the same as reading it out of the original and zeroing it when its block is hidden.

This is the lemma that makes apply usable in proofs. apply is defined by recursion on the remaining shape, unstacking one axis at a time, so its unfolded form talks about stackLeading and prefixes of coordinates rather than about masking. Stated pointwise it says what a reader expects a mask to mean, and the Option.map on the right is just getSpec reporting an out-of-range coordinate the same way on both sides.

def TorchLean.ssl.BlockMask.index {shape : Shape} (blocks : Tensor (Option ℕ) [shape.rank]) (coordinate : Tensor ℕ [shape.rank]) : Option ℕ

Row-major block index, or none for an invalid/degenerate block description.

def TorchLean.ssl.BlockMask.hidden {shape : Shape} (blocks : Tensor (Option ℕ) [shape.rank]) (period offset : ℕ) (coordinate : Tensor ℕ [shape.rank]) : Bool

Whether a coordinate belongs to the selected congruence class of blocks.

def TorchLean.ssl.BlockMask.apply {shape : Shape} (x : Tensor Float shape) (blocks : Tensor (Option ℕ) [shape.rank]) (period offset : ℕ) : Tensor Float shape

Set every scalar in a selected block to zero, preserving the tensor's arbitrary-rank shape.

For example, policies [none, some 4, some 4] repeat a 4-by-4 block mask across the first axis; [some 8] masks intervals in a signal; and [some 2, some 2, some 2] masks volume blocks.

theorem TorchLean.ssl.BlockMask.apply_scalar_at {shape : Shape} (blocks : Tensor (Option ℕ) [shape.rank]) (period offset : ℕ) (x : Tensor Float shape) (coordinate : Tensor ℕ [shape.rank]) : Spec.getSpec (apply x blocks period offset) (coordinate.to (List ℕ)) = Option.map (fun (value : Float) => if hidden blocks period offset coordinate = true then 0.0 else value) (Spec.getSpec x (coordinate.to (List ℕ)))

Exact coordinate semantics of apply, including out-of-bounds coordinates.

theorem TorchLean.ssl.BlockMask.hidden_scalar_eq_zero {shape : Shape} (blocks : Tensor (Option ℕ) [shape.rank]) (period offset : ℕ) (x : Tensor Float shape) (coordinate : Tensor ℕ [shape.rank]) (value : Float) (hValue : Spec.getSpec x (coordinate.to (List ℕ)) = some value) (hHidden : hidden blocks period offset coordinate = true) : Spec.getSpec (apply x blocks period offset) (coordinate.to (List ℕ)) = some 0.0

A selected in-bounds coordinate is exactly zero after masking.

theorem TorchLean.ssl.BlockMask.visible_scalar_eq_input {shape : Shape} (blocks : Tensor (Option ℕ) [shape.rank]) (period offset : ℕ) (x : Tensor Float shape) (coordinate : Tensor ℕ [shape.rank]) (value : Float) (hValue : Spec.getSpec x (coordinate.to (List ℕ)) = some value) (hVisible : hidden blocks period offset coordinate = false) : Spec.getSpec (apply x blocks period offset) (coordinate.to (List ℕ)) = some value

A visible in-bounds coordinate is copied unchanged by the mask.

def TorchLean.ssl.BlockMAE.Internal.sample (batchShape : Shape) {dataShape : Shape} (reconstructionWidth : ℕ) (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) (hReconstruction : reconstructionWidth ≤ dataShape.size) (x : Tensor Float (batchShape.concat dataShape)) : Sample.Supervised Float (batchShape.concat dataShape) (batchShape.appendDim reconstructionWidth)

Create a masked-reconstruction sample after validating the requested target width.

The model input retains its original shape. The target is a row-major prefix of the unmasked source because TorchLean's compact decoder heads produce matrices.

def TorchLean.ssl.BlockMAE.Internal.hiddenReconstructionIndices {dataShape : Shape} (reconstructionWidth : ℕ) (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) (hReconstruction : reconstructionWidth ≤ dataShape.size) : Array (Fin reconstructionWidth)

Flattened reconstruction coordinates hidden by the block mask.

def TorchLean.ssl.BlockMAE.hiddenMask {dataShape : Shape} (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) : Tensor Bool [dataShape.size]

Tensor-valued indicator of hidden reconstruction coordinates in row-major order.

def TorchLean.ssl.BlockMAE.reconstructionWeights {α : Type} [Storage α] [Context α] {dataShape : Shape} (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) : Tensor α [dataShape.size]

Uniform reconstruction weights on hidden entries, or all zeros when no entry is hidden.

def TorchLean.ssl.BlockMAE.hiddenIndices {dataShape : Shape} (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) : Array (Fin dataShape.size)

Return every hidden coordinate in the flattened data tensor.

The result covers the whole data shape, so no reconstruction-width validation is needed.

def TorchLean.ssl.BlockMAE.sample (batchShape : Shape) {dataShape : Shape} (reconstructionWidth : ℕ) (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) (x : Tensor Float (batchShape.concat dataShape)) : Except String (Sample.Supervised Float (batchShape.concat dataShape) (batchShape.appendDim reconstructionWidth))

Create a masked-reconstruction sample with arbitrary batch and data shapes.

The model input retains its original shape. The target is a row-major prefix of the unmasked source. An invalid reconstruction width is reported at the ordinary executable boundary rather than requiring callers to provide a theorem.

def TorchLean.ssl.BlockMAE.hiddenReconstructionIndices {dataShape : Shape} (reconstructionWidth : ℕ) (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) : Except String (Array (Fin reconstructionWidth))

Return the flattened reconstruction coordinates hidden by the block mask.

An invalid reconstruction width is reported explicitly.

def TorchLean.ssl.BlockMAE.Proof.rowPredictiveContract {dataShape : Shape} (batch reconstructionWidth : ℕ) (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) (hReconstruction : reconstructionWidth ≤ dataShape.size) (x : Tensor Float (dataShape.prependDim batch)) (prediction : Tensor Float [batch, reconstructionWidth]) (row : Fin batch) (loss : Float → Float → ℕ) : NN.MLTheory.SelfSupervised.PredictiveViewContract reconstructionWidth Unit Float Float Float

One batch row of block-MAE training as a finite predictive-view contract.

theorem TorchLean.ssl.BlockMAE.Proof.row_predictive_objective_eq_mae_loss {dataShape : Shape} (batch reconstructionWidth : ℕ) (blocks : Tensor (Option ℕ) [dataShape.rank]) (period offset : ℕ) (hReconstruction : reconstructionWidth ≤ dataShape.size) (x : Tensor Float (dataShape.prependDim batch)) (prediction : Tensor Float [batch, reconstructionWidth]) (row : Fin batch) (loss : Float → Float → ℕ) : NN.MLTheory.SelfSupervised.predictiveViewObjective (rowPredictiveContract batch reconstructionWidth blocks period offset hReconstruction x prediction row loss) = NN.MLTheory.SelfSupervised.maeLoss (Internal.hiddenReconstructionIndices reconstructionWidth blocks period offset hReconstruction) (fun (j : Fin reconstructionWidth) => (Spec.get (Spec.get (Internal.sample [batch] reconstructionWidth blocks period offset hReconstruction x).target row) j).item) (fun (j : Fin reconstructionWidth) => prediction[row][j]) loss

The runnable block-MAE row objective is exactly the finite MAE objective.