Vision Transformers Masked Image Modeling

# Vision Transformers & Masked Image Modeling

## Introduction & Motivation

For most of the 2010s, convolutional neural networks were the unchallenged default architecture for computer vision, their design (local receptive fields, weight sharing across spatial positions, hierarchical downsampling) encoding strong assumptions about images: that nearby pixels are more related than distant ones, and that useful visual features are translation-invariant. The Vision Transformer (ViT), introduced by Dosovitskiy and colleagues in 2020, challenged this default by showing that a nearly unmodified Transformer architecture, the same architecture family that had already come to dominate natural language processing, could match or exceed convolutional networks on image classification, provided it was pretrained on a sufficiently large dataset.

This result was notable precisely because the Transformer architecture, unlike a convolutional network, encodes almost no built-in assumptions about spatial structure: an image is simply cut into a grid of fixed-size patches, each patch is linearly projected into a token embedding, and the resulting sequence of patch tokens is processed by a standard Transformer encoder exactly as a sequence of word tokens would be processed in language modeling. The Transformer must learn essentially from scratch, purely from data, spatial relationships that a convolutional network would have built in as an architectural prior, an approach the original ViT paper explicitly framed as trading inductive bias for scale: with enough training data, learning these relationships from data outperforms hand-encoding them as architectural assumptions, whereas with limited data, the convolutional network's built-in assumptions provide a helpful head start that a data-hungry Transformer lacks.

The practical significance of this shift extends well beyond a single architecture swap. Adopting the Transformer as computer vision's standard building block enabled much closer architectural and methodological convergence between vision and language research, facilitating the multimodal vision-language models that fuse visual and textual understanding within closely related or even shared architectures, and enabling vision-specific adaptations of self-supervised pretraining techniques originally developed for language, most notably masked image modeling approaches directly inspired by the masked language modeling objective used to pretrain models such as BERT.

Masked Autoencoders (MAE) and related masked image modeling techniques subsequently demonstrated that Vision Transformers could be pretrained highly effectively using a self-supervised reconstruction objective, masking out a large majority of an image's patches and training the network to reconstruct the missing content from the small subset of visible patches, achieving strong downstream performance while requiring no manually annotated labels at all, further cementing the Vision Transformer, combined with self-supervised masked pretraining, as a central paradigm in modern computer vision research and practice.

## Core Concepts & Theory

The Vision Transformer processes an image by first dividing it into a grid of non-overlapping fixed-size patches (commonly 16 by 16 pixels), flattening each patch into a vector, and linearly projecting each flattened patch vector into a token embedding of the model's working dimensionality, exactly analogous to how a word or subword token is embedded into a vector in a language Transformer. Because the standard Transformer encoder's self-attention operation is permutation-invariant (it has no inherent notion of the order or spatial arrangement of its input tokens), a learned positional embedding is added to each patch token to inject information about that patch's original spatial location within the image, without which the model would be unable to distinguish an image from any arbitrary rearrangement of its own patches.

A special classification token, prepended to the sequence of patch tokens (borrowing directly from BERT's analogous classification token convention), accumulates information from across the full patch sequence through the Transformer's self-attention layers, and the final-layer representation of this token is used as the aggregate image representation fed into a classification head, though alternative pooling strategies (such as simple average pooling over all patch tokens' final representations) are also used in various ViT variants and have been shown to perform comparably or even somewhat better in some settings.

Self-attention within the Vision Transformer allows every patch token to directly attend to, and be influenced by, every other patch token in a single layer, regardless of spatial distance between the corresponding image regions, in contrast to a convolutional network, where information from distant image regions can only be combined after passing through multiple layers of local convolutional and pooling operations that progressively expand the effective receptive field. This gives the Vision Transformer a global receptive field from its very first layer, a property understood to be part of why it can perform so well given sufficient training data, since it need not wait for depth to accumulate long-range spatial context the way a convolutional network structurally must.

The original ViT paper's central empirical finding was that this architecture underperforms comparable convolutional networks when trained on small to medium datasets (such as ImageNet-1k trained from scratch), but matches or exceeds them when pretrained on much larger datasets (such as the 300-million-image JFT-300M dataset used in the original work) before fine-tuning on the downstream task, directly reflecting the architecture's comparative lack of built-in spatial inductive bias: with abundant data, the Transformer can learn effective spatial reasoning patterns directly from examples, but with limited data, it lacks the head start that a convolutional network's built-in translation-equivariance and locality assumptions provide.

## Mathematical Formulation

Given an input image, the patch embedding process first reshapes the image into a sequence of N flattened 2D patches, each patch of size P by P pixels (with N determined by the image resolution divided by the patch size in each spatial dimension), and linearly projects each flattened patch into a D-dimensional embedding using a learned projection matrix E, with a learnable class token embedding prepended and learnable position embeddings added elementwise:

$$ z_0 = [x_{class}; \, x_p^1 E; \, x_p^2 E; \, \ldots; \, x_p^N E] + E_{pos} $$

where x_p^i denotes the i-th flattened image patch, E is the shared learned linear patch-embedding projection, x_class is the learnable classification token embedding, and E_pos is the learned positional embedding added to the whole sequence, producing the initial token sequence z_0 that is then passed through a standard Transformer encoder consisting of alternating multi-head self-attention and feedforward blocks, each wrapped with residual connections and layer normalization, exactly matching the standard Transformer encoder architecture used in language modeling.

The self-attention operation applied within each Transformer layer, mapping queries Q, keys K, and values V (all linear projections of the layer's input token sequence) to an updated token representation, follows the standard scaled dot-product attention formulation:

$$ ext{Attention}(Q, K, V) = ext{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} ight) V $$

where d_k is the dimensionality of the key vectors, used to scale the dot products before the softmax to keep gradients well-behaved, exactly as in the standard Transformer architecture used across natural language processing, underscoring how directly the Vision Transformer reuses the language Transformer's core computational machinery with only the input tokenization (patches rather than words or subwords) substantively changed.

The Masked Autoencoder's training objective masks a large fraction (commonly 75 percent) of the image's patches, encoding only the small remaining subset of visible patches with a (relatively large) ViT encoder, then reconstructing the pixel values of the masked patches using a (relatively small and lightweight) decoder that takes the encoded visible-patch representations along with learned mask tokens marking the masked patch positions, trained via a simple mean squared error reconstruction loss computed only over the masked patches:

$$ L_{MAE} = \frac{1}{|M|} \sum_{i \in M} \| \hat{x}_i - x_i \|^2 $$

where M is the set of masked patch indices, x_i is the original pixel content of masked patch i, and x-hat_i is the decoder's reconstruction of that patch, with the loss deliberately computed only over the masked patches (not the visible ones), since the reconstruction task is meant to test and train the model's ability to infer missing content from context, not merely to copy already-visible input.

## Advanced Theory & Extensions

The very high masking ratio used by MAE, substantially higher than the roughly 15 percent masking ratio typically used in BERT-style masked language modeling, is motivated by a key difference between images and text: images contain substantial spatial redundancy (a missing image patch can often be predicted fairly well just by interpolating from its immediate visible neighbors, at low masking ratios, without requiring the model to learn much genuinely useful high-level visual understanding), so a much more aggressive masking ratio is needed to prevent the pretraining task from being solvable through simple local interpolation, forcing the model to instead develop genuinely useful, more global visual representations to succeed at the reconstruction task.

MAE's asymmetric encoder-decoder design, processing only the small subset of unmasked patches with the large, expressive encoder (rather than processing the full patch sequence, including masked-token placeholders, through the entire encoder), was a key practical innovation enabling substantially faster pretraining, since the computationally expensive encoder only ever processes roughly a quarter of the full patch sequence at high masking ratios, while the lightweight decoder, applied only briefly to reconstruct pixel values, can afford to process the full sequence (visible token representations plus mask tokens) without materially increasing overall training cost.

BEiT and related discretized-target masked image modeling approaches take a different approach to defining the reconstruction target than MAE's direct raw-pixel prediction: rather than reconstructing raw pixel values, these methods first train (or reuse) a separate discrete visual tokenizer that maps each image patch to a discrete "visual word" from a learned, finite vocabulary (conceptually analogous to how text is tokenized into a finite vocabulary of subword tokens), and the masked image modeling pretraining task becomes predicting the discrete visual token of each masked patch, framing the objective as a classification task over a finite, learned visual vocabulary, rather than as regression toward raw, continuous pixel values, a design choice argued to encourage the model to focus on higher-level, more semantically meaningful visual content rather than being overly influenced by low-level pixel statistics like exact color and lighting.

Hierarchical and windowed variants such as the Swin Transformer reintroduce some convolutional-style spatial inductive bias into the Vision Transformer framework, restricting self-attention computation to local, non-overlapping windows of patches (rather than full global self-attention across the entire patch sequence) and periodically shifting the window boundaries between successive layers to allow information to eventually flow across window boundaries, alongside a hierarchical, progressively coarsening multi-stage architecture that produces multi-scale feature maps directly analogous to a convolutional network's feature pyramid, a design that substantially reduces self-attention's quadratic computational cost (since attention is now computed only within small local windows rather than across the entire, much larger, patch sequence) while still retaining much of the flexibility and long-range modeling capacity of the standard Transformer's self-attention mechanism, particularly valuable for dense prediction tasks like object detection and segmentation that benefit from exactly this kind of multi-scale feature representation.

## Computational Considerations

Standard Vision Transformer self-attention incurs a computational and memory cost that scales quadratically with the number of image patches, which itself scales with image resolution (higher resolution images, or smaller patch sizes, both directly increase patch count), making naive full self-attention increasingly expensive for high-resolution image processing and pushing much practical ViT research and engineering toward efficiency-improving modifications, whether windowed and hierarchical attention patterns as used in Swin Transformer, more computationally efficient exact or approximate attention implementations, or simply operating at coarser patch granularities that trade spatial resolution for reduced computational cost.

Masked pretraining approaches like MAE substantially reduce pretraining compute relative to processing full, unmasked image sequences through the full encoder, since, as discussed above, only the small unmasked patch subset passes through the expensive full encoder at high masking ratios, a design that has made MAE-style pretraining an attractive, comparatively compute-efficient path to obtaining strong pretrained Vision Transformer representations relative to earlier, less compute-efficient self-supervised pretraining objectives requiring the full encoder to process every image patch during pretraining.

Fine-tuning and inference-time compute for a pretrained Vision Transformer scale with the same patch-count-dependent quadratic attention cost as pretraining, meaning deploying Vision Transformers for high-resolution image inference in latency-sensitive applications requires the same efficiency considerations (window-restricted attention, hierarchical multi-scale processing, or reduced input resolution) already discussed as relevant during pretraining and fine-tuning, an important practical consideration when selecting between a plain ViT and a more specialized, efficiency-oriented variant for a given deployment target.

## Practical Implementation Strategies

Selecting the patch size used to tokenize input images represents a direct trade-off between the resulting sequence length (and thus computational cost, scaling quadratically as discussed above) and the spatial granularity at which the model can distinguish fine visual detail, with smaller patches producing longer sequences (more compute, finer detail) and larger patches producing shorter sequences (less compute, coarser detail), a choice that should be informed by the specific downstream task's sensitivity to fine spatial detail (dense prediction tasks like segmentation generally benefit more from finer patch granularity than coarse whole-image classification tasks do).

Given the well-established finding that Vision Transformers particularly benefit from, and in the small-data regime particularly require, either very large-scale pretraining or a self-supervised pretraining objective like masked image modeling before fine-tuning on a smaller downstream dataset, practitioners working with limited labeled data for a specific downstream task should strongly prefer starting from an existing, publicly available pretrained ViT or MAE checkpoint rather than training a Vision Transformer from random initialization directly on their own smaller dataset, where a convolutional network alternative might reasonably be trained from scratch with less risk of underperformance due to insufficient inductive bias.

When fine-tuning a pretrained Vision Transformer on a downstream task, practitioners commonly use techniques such as layer-wise learning rate decay (applying progressively smaller learning rates to earlier, more general-purpose layers of the pretrained network relative to later, more task-specific layers) and, particularly at smaller downstream dataset sizes, stronger data augmentation and regularization than would typically be applied when fine-tuning a comparable convolutional network, reflecting the Vision Transformer's continued relative data-hungriness even after large-scale pretraining, compared to convolutional architectures.

## Benchmark Datasets & Evaluation

ImageNet-1k remains the standard benchmark for evaluating Vision Transformer image classification accuracy, both for models pretrained from scratch directly on ImageNet-1k (where plain ViT variants have historically underperformed comparable convolutional networks, as discussed above, absent additional pretraining data or self-supervised pretraining) and for models pretrained on larger external datasets (JFT-300M, ImageNet-21k, or through self-supervised masked image modeling) and subsequently fine-tuned on ImageNet-1k, where pretrained Vision Transformers have demonstrated performance matching or exceeding strong convolutional baselines.

Beyond image classification, standard dense prediction benchmarks including COCO for object detection and instance segmentation, and ADE20k for semantic segmentation, are widely used to evaluate Vision Transformer backbones (particularly hierarchical variants like Swin Transformer, given their natural fit for producing the multi-scale feature representations these dense prediction tasks typically require) against comparable convolutional network backbones, testing generalization of the Transformer-based visual representation beyond simple whole-image classification to more spatially detailed prediction tasks.

For self-supervised masked image modeling approaches specifically, linear probing (freezing the pretrained encoder and training only a simple linear classification head on top of its frozen representations) and full fine-tuning are both commonly reported evaluation protocols, providing complementary signals: linear probing more directly measures the quality and linear separability of the pretrained representations themselves, while full fine-tuning measures the practically more relevant question of ultimate downstream task performance after adapting all of the model's parameters to the specific downstream task.

## Key Challenges & Limitations

The comparative data-hungriness of plain Vision Transformers relative to convolutional networks, stemming directly from their weaker built-in spatial inductive bias, remains a real practical consideration: while large-scale pretraining and self-supervised masked pretraining have substantially mitigated this issue for many common application settings, practitioners working in genuinely data-scarce domains (highly specialized scientific or medical imaging domains, for instance, where even large-scale generic pretraining may transfer imperfectly) may still find convolutional architectures, or hybrid convolutional-Transformer architectures that combine both design philosophies, offer better practical performance than a plain Vision Transformer trained or fine-tuned with only limited task-specific data.

Interpreting exactly what Vision Transformer self-attention patterns represent, and how they relate to human-interpretable visual reasoning, remains an active and only partially resolved research question, and while attention-map visualization is commonly used to provide some qualitative insight into which image regions a trained ViT is attending to when making a given prediction, these visualizations, much like analogous saliency-map techniques for convolutional networks, do not always provide a fully reliable or complete account of the actual computational basis for the model's predictions.

The quadratic computational scaling of standard self-attention with patch count, discussed above as a computational consideration, also represents a genuine architectural limitation for very high-resolution image processing applications (such as gigapixel medical or satellite imagery), where even efficiency-oriented variants like windowed or hierarchical attention may still present meaningful computational challenges relative to a purely convolutional approach, whose per-pixel computational cost scales more favorably (linearly, for a fixed-size convolutional kernel) with image resolution.

## Hyperparameter Tuning

Model scale (depth, width, and the resulting total parameter count, following the standard ViT naming convention of Base, Large, and Huge variants of increasing size) interacts strongly with available pretraining data scale: the original ViT paper's central finding that larger ViT variants only outperform smaller ones, and only outperform comparable convolutional networks, once pretraining data scale is sufficiently large, means the appropriate model scale to select for a given application depends heavily on the scale of available pretraining (or fine-tuning) data, not simply on maximizing model capacity in isolation.

For masked image modeling pretraining specifically, the masking ratio is a particularly consequential and somewhat unusual (relative to typical deep learning hyperparameters) design choice, with MAE's ablation studies showing that reconstruction quality and, more importantly, downstream fine-tuning performance both depend substantially on this ratio, with the optimal masking ratio for downstream task transfer performance (75 percent, in MAE's reported results) notably higher than the ratio that would minimize raw pixel reconstruction error alone, illustrating that the pretraining objective's difficulty, not merely its solvability, is an important lever for producing genuinely useful learned representations.

The choice of positional embedding scheme (learned absolute positional embeddings, as in the original ViT, versus relative positional embeddings or other alternatives explored in various ViT variants) and how positional embeddings are adapted when fine-tuning at a different input resolution than that used during pretraining (commonly requiring some form of interpolation of the pretrained positional embeddings to match the new, different patch-count sequence length) both meaningfully affect fine-tuning performance, particularly when the fine-tuning resolution differs substantially from the original pretraining resolution.

## Real-World Applications & Case Studies

Medical imaging analysis has adopted Vision Transformers and masked-image-modeling-pretrained variants for tasks such as radiological image classification and segmentation, often initializing from large-scale natural-image pretrained checkpoints (ImageNet-scale or larger) before fine-tuning on comparatively much smaller, domain-specific labeled medical datasets, leveraging the transfer-learning benefits of large-scale pretraining to partially compensate for the relatively limited quantity of labeled medical imaging data typically available for any single specific clinical task.

Vision Transformer backbones, particularly hierarchical variants like Swin Transformer, have become standard components within modern object detection and semantic segmentation systems (discussed in dedicated treatments of those topics), replacing or complementing convolutional backbones within detection frameworks and demonstrating competitive or superior accuracy on standard dense-prediction benchmarks, reflecting the broader trend of Transformer-based architectures displacing or augmenting convolutional networks across an increasingly wide range of core computer vision tasks beyond simple image classification.

Multimodal vision-language models (discussed in dedicated treatments of that topic) commonly use a Vision Transformer, frequently one pretrained via masked image modeling or large-scale image-text contrastive pretraining such as CLIP, as their visual encoder component, converting an input image into a sequence of patch-token representations that can then be processed jointly with text tokens by a subsequent multimodal fusion architecture, directly leveraging the architectural convergence between vision and language processing that the Vision Transformer's adoption of the standard Transformer architecture enabled.

## Integration with Other Methods

CLIP-style contrastive vision-language pretraining, discussed in dedicated treatments of multimodal models, commonly uses a Vision Transformer as its image encoder, trained jointly with a text encoder using a contrastive objective that pulls matching image-text pairs together in a shared embedding space, an approach that, like masked image modeling, provides an effective self-supervised (in this case, weakly supervised by naturally paired image-text web data rather than manually annotated labels) pretraining signal for Vision Transformers, complementary to and sometimes combined with masked image modeling objectives within a single, multi-task pretraining pipeline.

Knowledge distillation techniques, discussed in dedicated treatments of model compression, have been applied specifically to Vision Transformers, notably in the DeiT (Data-efficient image Transformer) approach, which introduces an additional distillation token trained to match the predictions of a separate, already-trained convolutional network teacher, providing an additional, convolutional-network-derived training signal that helps the Vision Transformer train effectively with substantially less data and compute than the original ViT paper's very large-scale pretraining approach required.

Diffusion models for image generation, discussed in dedicated treatments of that topic, have increasingly adopted Transformer-based architectures (Diffusion Transformers, or DiT) in place of the originally more common convolutional U-Net backbone, directly applying Vision-Transformer-style patch-based tokenization within the image generation setting, illustrating the Vision Transformer's influence extending from discriminative visual understanding tasks into generative image modeling as well.

## Future Research Directions

Improving Vision Transformer data efficiency, closing the gap with convolutional networks in genuinely small-data training regimes without relying on very large-scale external pretraining, remains an active research direction, including further architectural hybrids that combine convolutional and Transformer design elements, and continued refinement of self-supervised pretraining objectives (extending beyond masked image modeling and contrastive pretraining) that extract more downstream-useful signal per unit of available pretraining data.

Extending efficient attention mechanisms and architectural innovations developed for very long text sequences in language modeling (discussed in dedicated treatments of long-context sequence architectures) to the Vision Transformer setting, where very high resolution imagery imposes analogous long-sequence computational challenges, represents an active area of cross-pollination between vision and language architecture research, given the close underlying architectural similarity between the two domains that the Vision Transformer's design established.

Better understanding what Vision Transformers actually learn during large-scale pretraining, whether through masked image modeling, contrastive image-text pretraining, or supervised classification pretraining, and how these different pretraining objectives shape the resulting learned visual representations differently, connects Vision Transformer research to the broader interpretability and representation-learning research agenda, with continued potential to inform both more effective pretraining objective design and more principled downstream model and objective selection for specific applications.

## Summary & Key Takeaways

The Vision Transformer applies the standard Transformer architecture to images by dividing them into fixed-size patches, linearly embedding each patch as a token, and processing the resulting patch sequence through standard self-attention layers, trading the strong built-in spatial inductive bias of convolutional networks for greater flexibility and improved performance given sufficiently large-scale pretraining data.

Masked Autoencoders and related masked image modeling techniques provide an effective, self-supervised pretraining objective for Vision Transformers directly inspired by masked language modeling, using a very high masking ratio (substantially higher than typical text masking ratios) and an efficient asymmetric encoder-decoder design to produce strong pretrained visual representations without requiring manually annotated labels.

Architectural variants such as the hierarchical, windowed Swin Transformer reintroduce some convolutional-style spatial inductive bias and multi-scale structure into the Vision Transformer framework, substantially improving computational efficiency and performance on dense prediction tasks like object detection and segmentation relative to a plain, globally-attending Vision Transformer.

The Vision Transformer's adoption of the standard Transformer architecture has enabled substantial architectural and methodological convergence between computer vision and natural language processing research, directly facilitating multimodal vision-language models, contrastive image-text pretraining approaches like CLIP, and Transformer-based generative image models, cementing it as a central, broadly influential architecture in modern computer vision.

Keywords: Vision Transformer, ViT, patch embedding, self-attention, positional embedding, masked autoencoder, MAE, masked image modeling, BEiT, discrete visual tokenizer, Swin Transformer, windowed attention, hierarchical Transformer, inductive bias, classification token, linear probing, DeiT distillation token, CLIP contrastive pretraining, Diffusion Transformer, layer-wise learning rate decay

---

## Appendix: Practical Labs

### Lab 1: Image-to-Patch Tokenization and Linear Patch Embedding

import numpy as np

np.random.seed(0)


def image_to_patches(image, patch_size):
    """Splits an (H, W, C) image into non-overlapping (patch_size, patch_size, C) patches,
    flattening each into a 1D vector, exactly as done at the input of a Vision Transformer."""
    h, w, c = image.shape
    assert h % patch_size == 0 and w % patch_size == 0, "Image dims must be divisible by patch_size"

    patches = []
    for i in range(0, h, patch_size):
        for j in range(0, w, patch_size):
            patch = image[i:i + patch_size, j:j + patch_size, :]
            patches.append(patch.flatten())
    return np.stack(patches)  # shape (num_patches, patch_size*patch_size*C)


class PatchEmbedding:
    def __init__(self, patch_dim, embed_dim, seed=0):
        rng = np.random.RandomState(seed)
        self.E = rng.randn(patch_dim, embed_dim) * (1.0 / np.sqrt(patch_dim))
        self.class_token = rng.randn(embed_dim) * 0.02
        self.pos_embedding = None  # initialized once num_patches is known

    def forward(self, patches):
        num_patches, patch_dim = patches.shape
        embed_dim = self.E.shape[1]

        if self.pos_embedding is None or self.pos_embedding.shape[0] != num_patches + 1:
            rng = np.random.RandomState(42)
            self.pos_embedding = rng.randn(num_patches + 1, embed_dim) * 0.02

        patch_tokens = patches @ self.E  # (num_patches, embed_dim)
        tokens = np.concatenate([self.class_token[None, :], patch_tokens], axis=0)
        tokens = tokens + self.pos_embedding
        return tokens


def test_patch_embedding_shapes():
    image = np.random.rand(32, 32, 3)
    patch_size = 8

    patches = image_to_patches(image, patch_size)
    expected_num_patches = (32 // patch_size) * (32 // patch_size)
    expected_patch_dim = patch_size * patch_size * 3

    print(f"Number of patches: {patches.shape[0]} (expected {expected_num_patches})")
    print(f"Flattened patch dimensionality: {patches.shape[1]} (expected {expected_patch_dim})")

    assert patches.shape == (expected_num_patches, expected_patch_dim), "Patch extraction shape mismatch"

    embed_dim = 64
    embedder = PatchEmbedding(patch_dim=expected_patch_dim, embed_dim=embed_dim, seed=1)
    tokens = embedder.forward(patches)

    print(f"Token sequence shape (including class token): {tokens.shape}")
    assert tokens.shape == (expected_num_patches + 1, embed_dim), (
        "Token sequence should include the prepended class token"
    )
    assert np.allclose(tokens[0] - embedder.pos_embedding[0], embedder.class_token), (
        "First token should correspond to the class token (plus its positional embedding)"
    )

    print("Patch tokenization and embedding test passed.")


if __name__ == "__main__":
    test_patch_embedding_shapes()

### Lab 2: Scaled Dot-Product Self-Attention Over Patch Tokens

import numpy as np

np.random.seed(1)


def softmax(x, axis=-1):
    shifted = x - np.max(x, axis=axis, keepdims=True)
    exp = np.exp(shifted)
    return exp / np.sum(exp, axis=axis, keepdims=True)


def scaled_dot_product_attention(Q, K, V):
    """Q, K, V: (seq_len, d_k). Returns attention output (seq_len, d_v) and
    the attention weight matrix (seq_len, seq_len)."""
    d_k = Q.shape[-1]
    scores = (Q @ K.T) / np.sqrt(d_k)
    weights = softmax(scores, axis=-1)
    output = weights @ V
    return output, weights


class SingleHeadSelfAttention:
    def __init__(self, embed_dim, seed=0):
        rng = np.random.RandomState(seed)
        self.W_q = rng.randn(embed_dim, embed_dim) * (1.0 / np.sqrt(embed_dim))
        self.W_k = rng.randn(embed_dim, embed_dim) * (1.0 / np.sqrt(embed_dim))
        self.W_v = rng.randn(embed_dim, embed_dim) * (1.0 / np.sqrt(embed_dim))

    def forward(self, tokens):
        Q = tokens @ self.W_q
        K = tokens @ self.W_k
        V = tokens @ self.W_v
        return scaled_dot_product_attention(Q, K, V)


def test_self_attention_properties():
    seq_len, embed_dim = 10, 16
    tokens = np.random.randn(seq_len, embed_dim)

    attn_layer = SingleHeadSelfAttention(embed_dim, seed=2)
    output, weights = attn_layer.forward(tokens)

    print(f"Attention output shape: {output.shape}")
    print(f"Attention weight matrix shape: {weights.shape}")
    print(f"Row sums of attention weights (should all be ~1.0): {weights.sum(axis=-1)[:3]}")

    assert output.shape == (seq_len, embed_dim), "Attention output shape mismatch"
    assert weights.shape == (seq_len, seq_len), "Attention weight matrix should be seq_len x seq_len"
    assert np.allclose(weights.sum(axis=-1), 1.0, atol=1e-6), (
        "Each row of the attention weight matrix (a softmax output) should sum to 1"
    )

    # Test global receptive field: perturbing a DISTANT token should measurably
    # change the output at every other token's position, unlike a local convolution
    tokens_perturbed = tokens.copy()
    tokens_perturbed[-1] += 5.0  # large perturbation to the LAST token only

    output_perturbed, _ = attn_layer.forward(tokens_perturbed)
    change_at_first_token = np.abs(output_perturbed[0] - output[0]).mean()

    print(f"Change in FIRST token's output after perturbing the LAST token: {change_at_first_token:.4f}")
    assert change_at_first_token > 1e-4, (
        "Self-attention should propagate influence from a distant token to every "
        "other token in a single layer (global receptive field)"
    )

    print("Self-attention properties test passed.")


if __name__ == "__main__":
    test_self_attention_properties()

### Lab 3: MAE-Style Random Masking and Asymmetric Encoder-Decoder Reconstruction

import numpy as np

np.random.seed(2)


def random_masking(num_patches, mask_ratio, seed=None):
    """Randomly selects which patch indices are masked (hidden) vs. visible (kept),
    matching MAE's random masking strategy."""
    rng = np.random.RandomState(seed)
    num_masked = int(num_patches * mask_ratio)
    perm = rng.permutation(num_patches)
    masked_indices = perm[:num_masked]
    visible_indices = perm[num_masked:]
    return visible_indices, masked_indices


class TinyLinearLayer:
    def __init__(self, in_dim, out_dim, seed):
        rng = np.random.RandomState(seed)
        self.W = rng.randn(in_dim, out_dim) * (1.0 / np.sqrt(in_dim))

    def forward(self, x):
        return x @ self.W


class ToyMAE:
    """A drastically simplified MAE: linear 'encoder' applied only to visible
    patches, linear 'decoder' applied to encoded visible tokens + learned mask
    tokens (placed at masked positions) to reconstruct all patches."""

    def __init__(self, patch_dim, encoder_dim, seed=0):
        self.encoder = TinyLinearLayer(patch_dim, encoder_dim, seed=seed)
        self.decoder = TinyLinearLayer(encoder_dim, patch_dim, seed=seed + 1)
        rng = np.random.RandomState(seed + 2)
        self.mask_token = rng.randn(encoder_dim) * 0.02

    def forward(self, patches, mask_ratio, seed=None):
        num_patches = patches.shape[0]
        visible_idx, masked_idx = random_masking(num_patches, mask_ratio, seed=seed)

        # Encoder only processes VISIBLE patches (the key MAE efficiency trick)
        encoded_visible = self.encoder.forward(patches[visible_idx])

        # Reassemble full sequence: encoded visible tokens + mask tokens at masked positions
        encoder_dim = encoded_visible.shape[1]
        full_encoded = np.zeros((num_patches, encoder_dim))
        full_encoded[visible_idx] = encoded_visible
        full_encoded[masked_idx] = self.mask_token

        reconstruction = self.decoder.forward(full_encoded)
        return reconstruction, visible_idx, masked_idx


def test_mae_masking_and_reconstruction_target():
    num_patches, patch_dim, encoder_dim = 16, 12, 8
    patches = np.random.randn(num_patches, patch_dim)

    mae = ToyMAE(patch_dim, encoder_dim, seed=3)
    mask_ratio = 0.75

    reconstruction, visible_idx, masked_idx = mae.forward(patches, mask_ratio, seed=5)

    print(f"Total patches: {num_patches}")
    print(f"Visible patches: {len(visible_idx)} ({len(visible_idx)/num_patches:.0%})")
    print(f"Masked patches: {len(masked_idx)} ({len(masked_idx)/num_patches:.0%})")

    assert len(masked_idx) == int(num_patches * mask_ratio), "Masked patch count should match mask_ratio"
    assert len(visible_idx) + len(masked_idx) == num_patches, "Visible + masked should cover all patches"
    assert set(visible_idx.tolist()).isdisjoint(set(masked_idx.tolist())), (
        "Visible and masked index sets must be disjoint"
    )

    # Reconstruction loss should be computed ONLY over masked patches, per the MAE objective
    reconstruction_loss_masked = np.mean((reconstruction[masked_idx] - patches[masked_idx]) ** 2)
    print(f"MSE reconstruction loss on MASKED patches (untrained model): {reconstruction_loss_masked:.3f}")

    assert reconstruction.shape == patches.shape, "Reconstruction should cover all patch positions"
    assert np.isfinite(reconstruction_loss_masked), "Reconstruction loss should be a finite number"

    print("MAE-style masking and reconstruction test passed.")


if __name__ == "__main__":
    test_mae_masking_and_reconstruction_target()

### Lab 4: Windowed vs. Global Self-Attention Compute Cost Comparison

import numpy as np

np.random.seed(3)


def global_attention_flops(num_patches, embed_dim):
    """Approximate FLOPs for one global self-attention layer: dominated by the
    Q@K^T and attention@V matrix multiplications, each O(num_patches^2 * embed_dim)."""
    qk_flops = num_patches * num_patches * embed_dim
    av_flops = num_patches * num_patches * embed_dim
    return qk_flops + av_flops


def windowed_attention_flops(num_patches, embed_dim, window_size):
    """Approximate FLOPs for windowed self-attention (Swin-Transformer-style):
    attention is computed independently within each non-overlapping window of
    window_size^2 patches, so total cost scales linearly in num_patches rather
    than quadratically, at the cost of restricting each layer's receptive field."""
    assert num_patches % (window_size * window_size) == 0, "Patches must divide evenly into windows"
    num_windows = num_patches // (window_size * window_size)
    patches_per_window = window_size * window_size

    per_window_flops = global_attention_flops(patches_per_window, embed_dim)
    return per_window_flops * num_windows


def test_windowed_attention_scales_better():
    embed_dim = 96
    window_size = 7  # 7x7 = 49 patches per window, a common Swin Transformer choice

    # Test across a range of increasingly large images (more total patches)
    patch_counts = [196, 784, 3136, 12544]  # e.g. 14x14, 28x28, 56x56, 112x112 patch grids, each a multiple of the 7x7 window

    print(f"{'num_patches':>12} | {'global_flops':>15} | {'windowed_flops':>15} | speedup")
    speedups = []
    for num_patches in patch_counts:
        global_cost = global_attention_flops(num_patches, embed_dim)
        windowed_cost = windowed_attention_flops(num_patches, embed_dim, window_size)
        speedup = global_cost / windowed_cost
        speedups.append(speedup)
        print(f"{num_patches:>12} | {global_cost:>15,} | {windowed_cost:>15,} | {speedup:6.1f}x")

    # Global attention cost grows quadratically, windowed cost grows linearly,
    # so the speedup factor of windowed over global should INCREASE as num_patches grows
    assert all(s > 1.0 for s in speedups), "Windowed attention should always be cheaper than global"
    assert speedups[-1] > speedups[0], (
        "The efficiency advantage of windowed attention over global attention should "
        "grow as the number of patches (image resolution) increases, reflecting "
        "quadratic vs. linear scaling"
    )

    print(f"
Speedup at smallest resolution: {speedups[0]:.1f}x")
    print(f"Speedup at largest resolution:  {speedups[-1]:.1f}x")
    print("Windowed vs. global attention scaling test passed.")


if __name__ == "__main__":
    test_windowed_attention_scales_better()