Vision Transformers

import torch
import torch.nn as nn
import numpy as np

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        
        # Linear projection of flattened patches
        self.proj = nn.Conv2d(
            in_channels, 
            embed_dim,
            kernel_size=patch_size,
            stride=patch_size
        )
        
        # Position embeddings
        self.pos_embed = nn.Parameter(
            torch.randn(1, self.n_patches + 1, embed_dim)
        )
        
        # CLS token
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
    
    def forward(self, x):
        # x: (B, C, H, W)
        B = x.shape[0]
        
        # Part 1: Reshape image to patches and project
        # Conv2d with stride=patch_size extracts non-overlapping patches
        x = self.proj(x)  # (B, embed_dim, H/P, W/P)
        x = x.flatten(2)  # (B, embed_dim, n_patches)
        x = x.transpose(1, 2)  # (B, n_patches, embed_dim)
        
        # Part 2: Add CLS token
        cls_tokens = self.cls_token.expand(B, -1, -1)  # (B, 1, embed_dim)
        x = torch.cat([cls_tokens, x], dim=1)  # (B, n_patches+1, embed_dim)
        
        # Part 3: Add position embeddings
        x = x + self.pos_embed
        
        return x

# Example usage
img_size = 224
patch_size = 16
in_channels = 3
embed_dim = 768

# Create model
patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)

# Create sample image
batch_size = 4
image = torch.randn(batch_size, in_channels, img_size, img_size)

# Forward pass
output = patch_embed(image)

print(f"Input image shape: {image.shape}")
print(f"Number of patches: {(img_size // patch_size) ** 2}")
print(f"Output shape: {output.shape}")
print(f"Expected: (batch_size, n_patches+1, embed_dim)")
print(f"Actual: ({batch_size}, {(img_size//patch_size)**2 + 1}, {embed_dim})")

Detailed Breakdown:

Part (a): Reshape Image to Patches

Original image: $224 \times 224 \times 3$

Patch size: $16 \times 16$

Number of patches: $\frac{224}{16} \times \frac{224}{16} = 14 \times 14 = 196$ patches

Each patch: $16 \times 16 \times 3 = 768$ values

Reshaped: $(196, 768)$

Part (b): Linear Projection

Using Conv2d with kernel\_size=16, stride=16:

• Input: $(B, 3, 224, 224)$
• Output: $(B, 768, 14, 14)$
• Flatten spatial dimensions: $(B, 768, 196)$
• Transpose: $(B, 196, 768)$

This is equivalent to:

1. Extract $14 \times 14 = 196$ non-overlapping patches
2. Flatten each patch: $16 \times 16 \times 3 = 768$ values
3. Apply linear projection: $\mathbb{R}^{768} \to \mathbb{R}^{768}$

Part (c): Add Position Embeddings

Position embeddings: learnable parameters of shape $(1, 197, 768)$

• 197 = 196 patches + 1 CLS token
• Each position has unique 768-dimensional embedding
• Added element-wise to patch embeddings

$$\vx_i = \text{PatchEmbed}(\text{patch}_i) + \vpos_i$$

Part (d): Output Shape Verification

Input image shape: torch.Size([4, 3, 224, 224])
Number of patches: 196
Output shape: torch.Size([4, 197, 768])
Expected: (batch_size, n_patches+1, embed_dim)
Actual: (4, 197, 768)

Output shape: $(B, 197, 768)$ where:

• $B = 4$: batch size
• $197 = 196 + 1$: patches + CLS token
• $768$: embedding dimension

Key Design Choices:

1. Conv2d for patching: More efficient than manual reshaping
2. CLS token: Special token for classification (like BERT)
3. Learnable position embeddings: Unlike sinusoidal in original Transformer
4. No overlap: Patches don't overlap (stride = patch\_size)

Alternative Implementation (Manual):

def manual_patch_embedding(image, patch_size=16):
    """Manual patch extraction without Conv2d"""
    B, C, H, W = image.shape
    P = patch_size
    
    # Reshape to patches
    patches = image.unfold(2, P, P).unfold(3, P, P)  # (B, C, H/P, W/P, P, P)
    patches = patches.contiguous().view(B, C, -1, P, P)  # (B, C, n_patches, P, P)
    patches = patches.permute(0, 2, 1, 3, 4)  # (B, n_patches, C, P, P)
    patches = patches.reshape(B, -1, C * P * P)  # (B, n_patches, C*P*P)
    
    return patches

# Verify equivalence
manual_patches = manual_patch_embedding(image, patch_size)
print(f"Manual patches shape: {manual_patches.shape}")  # (4, 196, 768)

Both methods produce identical results, but Conv2d is more efficient and commonly used in practice.

Part (a): Parameter Count

ViT-Base:

• Layers: $L = 12$
• Hidden size: $d = 768$
• Attention heads: $h = 12$
• MLP size: $d_{mlp} = 3072$
• Image size: $224 \times 224$
• Patch size: $16 \times 16$
• Number of patches: $196$

Parameters:

• Patch embedding: $3 \times 16 \times 16 \times 768 = 589{,}824$
• Position embeddings: $197 \times 768 = 151{,}296$
• CLS token: $768$
• Per transformer layer:
  • Attention: $4 \times 768^2 = 2{,}359{,}296$
  • MLP: $768 \times 3072 + 3072 \times 768 = 4{,}718{,}592$
  • Layer norm: $2 \times 2 \times 768 = 3{,}072$
  • Total per layer: $7{,}080{,}960$
• 12 layers: $12 \times 7{,}080{,}960 = 84{,}971{,}520$
• Classification head: $768 \times 1000 = 768{,}000$

Total ViT-Base: $86{,}481{,}408 \approx 86$M parameters

ResNet-50:

• Conv1: $7 \times 7 \times 3 \times 64 = 9{,}408$
• Layer 1 (3 blocks): $\approx 215{,}808$
• Layer 2 (4 blocks): $\approx 1{,}219{,}648$
• Layer 3 (6 blocks): $\approx 7{,}098{,}880$
• Layer 4 (3 blocks): $\approx 14{,}964{,}800$
• FC layer: $2048 \times 1000 = 2{,}048{,}000$

Total ResNet-50: $25{,}556{,}032 \approx 25.6$M parameters

Ratio: ViT-Base has $3.4\times$ more parameters than ResNet-50

Part (b): FLOPs for $224 \times 224$ Image

ViT-Base FLOPs:

1. Patch Embedding:

• Conv2d: $3 \times 16 \times 16 \times 768 \times 14 \times 14 = 115{,}605{,}504$ FLOPs

2. Per Transformer Layer (12 layers):

Multi-Head Attention:

• $Q, K, V$ projections: $3 \times 197 \times 768 \times 768 = 347{,}054{,}592$ FLOPs
• Attention scores: $197 \times 197 \times 768 = 29{,}859{,}712$ FLOPs
• Attention weights $\times V$: $197 \times 197 \times 768 = 29{,}859{,}712$ FLOPs
• Output projection: $197 \times 768 \times 768 = 115{,}684{,}864$ FLOPs
• Total attention: $522{,}458{,}880$ FLOPs

MLP:

• First linear: $197 \times 768 \times 3072 = 464{,}739{,}456$ FLOPs
• Second linear: $197 \times 3072 \times 768 = 464{,}739{,}456$ FLOPs
• Total MLP: $929{,}478{,}912$ FLOPs

Total per layer: $1{,}451{,}937{,}792$ FLOPs

12 layers: $12 \times 1{,}451{,}937{,}792 = 17{,}423{,}253{,}504$ FLOPs

3. Classification Head:

• $768 \times 1000 = 768{,}000$ FLOPs

Total ViT-Base: $\approx 17.5$ GFLOPs

ResNet-50 FLOPs:

• Conv1 + BN: $\approx 118$ MFLOPs
• Layer 1: $\approx 1{,}219$ MFLOPs
• Layer 2: $\approx 1{,}627$ MFLOPs
• Layer 3: $\approx 3{,}254$ MFLOPs
• Layer 4: $\approx 1{,}627$ MFLOPs
• FC: $\approx 2$ MFLOPs

Total ResNet-50: $\approx 4.1$ GFLOPs

Ratio: ViT-Base requires $4.3\times$ more FLOPs than ResNet-50

Part (c): Memory Footprint

ViT-Base Memory (Inference):

1. Activations per layer:

• Input: $197 \times 768 = 151{,}296$ values
• Attention scores: $12 \times 197 \times 197 = 465{,}228$ values
• MLP intermediate: $197 \times 3072 = 605{,}184$ values

Peak activation memory per layer: $\approx 1.2$M values $\times$ 4 bytes = 4.8 MB

Total for 12 layers: $\approx 58$ MB

Parameters: $86$M $\times$ 4 bytes = 344 MB

Total ViT-Base inference: $\approx 402$ MB

ResNet-50 Memory (Inference):

Peak activation memory:

• Layer 1 output: $56 \times 56 \times 256 = 802{,}816$ values
• Layer 2 output: $28 \times 28 \times 512 = 401{,}408$ values
• Layer 3 output: $14 \times 14 \times 1024 = 200{,}704$ values
• Layer 4 output: $7 \times 7 \times 2048 = 100{,}352$ values

Peak: $\approx 3.2$ MB

Parameters: $25.6$M $\times$ 4 bytes = 102 MB

Total ResNet-50 inference: $\approx 105$ MB

Ratio: ViT-Base uses $3.8\times$ more memory than ResNet-50

Part (d): Which is More Efficient?

Efficiency Analysis:

Conclusion:

ResNet-50 is more computationally efficient in terms of:

• Fewer parameters ($3.4\times$ less)
• Lower FLOPs ($4.3\times$ less)
• Smaller memory footprint ($3.8\times$ less)

However, ViT-Base has advantages:

1. Better scaling: Performance improves more with larger datasets
2. Transfer learning: Pre-trained ViT generalizes better
3. Parallelization: Self-attention is more parallelizable than convolutions
4. Long-range dependencies: Global receptive field from layer 1
5. Interpretability: Attention maps show what model focuses on

Trade-off:

• Small datasets: ResNet-50 is better (more efficient, better inductive bias)
• Large datasets (ImageNet-21k, JFT-300M): ViT-Base is better (superior accuracy)
• Edge deployment: ResNet-50 preferred (lower resource requirements)
• Cloud deployment: ViT-Base viable (resources available, better accuracy)

Practical Recommendation:

For ImageNet-1k from scratch: ResNet-50

For transfer learning with pre-training: ViT-Base

For production with limited compute: ResNet-50

For research and maximum accuracy: ViT-Base

import torch
import torch.nn as nn
import numpy as np

class MAEMasking(nn.Module):
    def __init__(self, n_patches=196, embed_dim=768, mask_ratio=0.75):
        super().__init__()
        self.n_patches = n_patches
        self.embed_dim = embed_dim
        self.mask_ratio = mask_ratio
        self.n_visible = int(n_patches * (1 - mask_ratio))
        
        # Mask token for decoder
        self.mask_token = nn.Parameter(torch.randn(1, 1, embed_dim))
    
    def random_masking(self, x):
        """
        Randomly mask patches
        Args:
            x: (B, N, D) where N = n_patches + 1 (including CLS)
        Returns:
            x_visible: (B, n_visible+1, D) visible patches + CLS
            mask: (B, N) binary mask (0 = keep, 1 = remove)
            ids_restore: (B, N) indices to restore original order
        """
        B, N, D = x.shape
        N_patches = N - 1  # Exclude CLS token
        
        # Generate random noise for shuffling
        noise = torch.rand(B, N_patches, device=x.device)
        
        # Sort noise to get shuffle indices
        ids_shuffle = torch.argsort(noise, dim=1)
        ids_restore = torch.argsort(ids_shuffle, dim=1)
        
        # Keep first n_visible patches
        ids_keep = ids_shuffle[:, :self.n_visible]
        
        # Extract CLS token
        cls_token = x[:, :1, :]  # (B, 1, D)
        
        # Extract patch tokens (exclude CLS)
        x_patches = x[:, 1:, :]  # (B, N_patches, D)
        
        # Gather visible patches
        x_visible = torch.gather(
            x_patches, 
            dim=1, 
            index=ids_keep.unsqueeze(-1).expand(-1, -1, D)
        )  # (B, n_visible, D)
        
        # Concatenate CLS token
        x_visible = torch.cat([cls_token, x_visible], dim=1)  # (B, n_visible+1, D)
        
        # Generate binary mask: 0 = keep, 1 = remove
        mask = torch.ones(B, N_patches, device=x.device)
        mask[:, :self.n_visible] = 0
        mask = torch.gather(mask, dim=1, index=ids_restore)
        
        return x_visible, mask, ids_restore
    
    def add_mask_tokens(self, x_visible, ids_restore):
        """
        Add mask tokens for decoder
        Args:
            x_visible: (B, n_visible+1, D)
            ids_restore: (B, N_patches)
        Returns:
            x_full: (B, N, D) with mask tokens
        """
        B, _, D = x_visible.shape
        
        # Extract CLS token
        cls_token = x_visible[:, :1, :]
        
        # Extract visible patches
        x_patches = x_visible[:, 1:, :]  # (B, n_visible, D)
        
        # Create mask tokens
        n_mask = self.n_patches - self.n_visible
        mask_tokens = self.mask_token.expand(B, n_mask, -1)
        
        # Concatenate visible and mask tokens
        x_combined = torch.cat([x_patches, mask_tokens], dim=1)  # (B, N_patches, D)
        
        # Restore original order
        x_restored = torch.gather(
            x_combined,
            dim=1,
            index=ids_restore.unsqueeze(-1).expand(-1, -1, D)
        )  # (B, N_patches, D)
        
        # Add CLS token back
        x_full = torch.cat([cls_token, x_restored], dim=1)  # (B, N, D)
        
        return x_full
    
    def forward(self, x):
        """
        Complete MAE masking pipeline
        """
        # Part 1: Random masking
        x_visible, mask, ids_restore = self.random_masking(x)
        
        # Part 2: Add mask tokens for decoder
        x_full = self.add_mask_tokens(x_visible, ids_restore)
        
        return x_visible, x_full, mask, ids_restore



# Example usage
n_patches = 196
embed_dim = 768
mask_ratio = 0.75
batch_size = 4

# Create MAE masking module
mae_mask = MAEMasking(n_patches, embed_dim, mask_ratio)

# Simulate patch embeddings (including CLS token)
x = torch.randn(batch_size, n_patches + 1, embed_dim)

# Apply masking
x_visible, x_full, mask, ids_restore = mae_mask(x)

print(f"Input shape: {x.shape}")
print(f"Visible patches shape: {x_visible.shape}")
print(f"Full with mask tokens shape: {x_full.shape}")
print(f"Mask shape: {mask.shape}")
print(f"Number of masked patches: {mask.sum(dim=1)[0].item()}")
print(f"Number of visible patches: {(1 - mask).sum(dim=1)[0].item()}")

Input shape: torch.Size([4, 197, 768])
Visible patches shape: torch.Size([4, 50, 768])
Full with mask tokens shape: torch.Size([4, 197, 768])
Mask shape: torch.Size([4, 196])
Number of masked patches: 147.0
Number of visible patches: 49.0

Part (a): Randomly Mask 75\% of 196 Patches

Masking Strategy:

1. Generate random noise: $\text{noise} \sim \mathcal{U}(0, 1)^{196}$
2. Sort noise to get shuffle indices
3. Keep first $49$ patches (25\% of 196)
4. Mask remaining $147$ patches (75\% of 196)

Mathematical Formulation:

Let $\vx = [\vx_{\text{cls}}, \vx_1, \vx_2, \ldots, \vx_{196}]$ be patch embeddings.

Random permutation: $\pi: \{1, \ldots, 196\} \to \{1, \ldots, 196\}$

Visible set: $\mathcal{V} = \{\pi(1), \ldots, \pi(49)\}$

Masked set: $\mathcal{M} = \{\pi(50), \ldots, \pi(196)\}$

Binary mask: $m_i = \begin{cases} 0 & \text{if } i \in \mathcal{V} \\ 1 & \text{if } i \in \mathcal{M} \end{cases}$

Part (b): Keep 49 Visible Patches

Encoder Input:

$\vx_{\text{visible}} = [\vx_{\text{cls}}, \vx_{\pi(1)}, \vx_{\pi(2)}, \ldots, \vx_{\pi(49)}]$

Shape: $(B, 50, 768)$ where $50 = 49 + 1$ (CLS token)

Computational Savings:

Encoder processes only 25\% of patches:

• Attention complexity: $O(50^2 \cdot 768)$ vs $O(197^2 \cdot 768)$
• Speedup: $\frac{197^2}{50^2} = 15.5\times$ faster
• Memory reduction: $15.5\times$ less

This is the key efficiency gain of MAE!

Part (c): Add Mask Tokens for Decoder

Decoder Input Construction:

1. Take encoder output: $\vz_{\text{visible}} = \text{Encoder}(\vx_{\text{visible}})$
2. Create mask tokens: $\vm_{\text{mask}} \in \mathbb{R}^{147 \times 768}$ (learnable)
3. Concatenate: $[\vz_{\pi(1)}, \ldots, \vz_{\pi(49)}, \vm_1, \ldots, \vm_{147}]$
4. Restore original order using $\pi^{-1}$
5. Add position embeddings

Decoder Input:

$\vx_{\text{decoder}} = [\vx_{\text{cls}}, \vz_1, \vz_2, \ldots, \vz_{196}]$

where $\vz_i = \begin{cases} \text{Encoder output} & \text{if } i \in \mathcal{V} \\ \vm_{\text{mask}} & \text{if } i \in \mathcal{M} \end{cases}$

Shape: $(B, 197, 768)$ - full sequence restored

Part (d): Compute Reconstruction Loss

Decoder Output:

$\hat{\vx}_i = \text{Decoder}(\vx_{\text{decoder}})_i$ for $i = 1, \ldots, 196$

Reconstruction Loss (MSE on masked patches only):

$\mathcal{L}_{\text{MAE}} = \frac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \|\hat{\vx}_i - \vx_i\|_2^2$

Only compute loss on masked patches (147 patches):

def compute_mae_loss(original_patches, reconstructed_patches, mask):
    """
    Compute MAE reconstruction loss
    Args:
        original_patches: (B, N, D) original patch embeddings
        reconstructed_patches: (B, N, D) decoder output
        mask: (B, N) binary mask (1 = masked, 0 = visible)
    Returns:
        loss: scalar
    """
    # Compute MSE
    mse = (reconstructed_patches - original_patches) ** 2
    mse = mse.mean(dim=-1)  # (B, N) - mean over embedding dim
    
    # Apply mask - only compute loss on masked patches
    loss = (mse * mask).sum() / mask.sum()
    
    return loss

# Example
original = torch.randn(4, 196, 768)
reconstructed = torch.randn(4, 196, 768)
mask = torch.zeros(4, 196)
mask[:, 49:] = 1  # Mask last 147 patches

loss = compute_mae_loss(original, reconstructed, mask)
print(f"MAE Loss: {loss.item():.4f}")

Why Only Masked Patches?

• Visible patches are already seen by encoder
• Predicting visible patches is trivial (identity mapping)
• Masked patches require understanding context
• Forces model to learn semantic representations

Complete MAE Training Loop:

# 1. Patch embedding
patches = patch_embed(images)  # (B, 197, 768)

# 2. Random masking
visible_patches, mask, ids_restore = random_masking(patches)  # (B, 50, 768)

# 3. Encoder (only on visible patches)
encoded = encoder(visible_patches)  # (B, 50, 768)

# 4. Add mask tokens and restore order
decoder_input = add_mask_tokens(encoded, ids_restore)  # (B, 197, 768)

# 5. Decoder
reconstructed = decoder(decoder_input)  # (B, 197, 768)

# 6. Compute loss (only on masked patches)
loss = compute_mae_loss(patches[:, 1:], reconstructed[:, 1:], mask)

# 7. Backpropagation
loss.backward()

Key Insights:

1. High masking ratio (75\%): Forces model to learn global structure
2. Random masking: Prevents trivial solutions (interpolation)
3. Asymmetric encoder-decoder: Encoder is large, decoder is small
4. Pixel-level reconstruction: Simpler than contrastive learning
5. Efficiency: $15.5\times$ faster than processing all patches

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

class ViTTiny(nn.Module):
    def __init__(self, img_size=32, patch_size=4, in_channels=3, 
                 num_classes=10, embed_dim=192, depth=6, num_heads=3,
                 mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2  # 64 patches
        
        # Patch embedding
        self.patch_embed = nn.Conv2d(
            in_channels, embed_dim, 
            kernel_size=patch_size, stride=patch_size
        )
        
        # CLS token and position embeddings
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.randn(1, self.n_patches + 1, embed_dim))
        self.dropout = nn.Dropout(dropout)
        
        # Transformer encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=int(embed_dim * mlp_ratio),
            dropout=dropout,
            activation='gelu',
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=depth)
        
        # Classification head
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)
    
    def forward(self, x):
        B = x.shape[0]
        
        # Patch embedding
        x = self.patch_embed(x)  # (B, embed_dim, H/P, W/P)
        x = x.flatten(2).transpose(1, 2)  # (B, n_patches, embed_dim)
        
        # Add CLS token
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)
        
        # Add position embeddings
        x = x + self.pos_embed
        x = self.dropout(x)
        
        # Transformer
        x = self.transformer(x)
        
        # Classification
        x = self.norm(x[:, 0])  # CLS token
        x = self.head(x)
        
        return x



# Part (a): Patch size 4 for 32x32 images
print(f"Image size: 32x32")
print(f"Patch size: 4x4")
print(f"Number of patches: {(32 // 4) ** 2} = 64")
print(f"Each patch: 4x4x3 = 48 values")

# Part (b): Model configuration
model = ViTTiny(
    img_size=32,
    patch_size=4,
    in_channels=3,
    num_classes=10,
    embed_dim=192,
    depth=6,
    num_heads=3,
    mlp_ratio=4.0
)

# Count parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"\nViT-Tiny parameters: {total_params:,}")

# Part (c): RandAugment data augmentation
from torchvision.transforms import RandAugment

train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# Load CIFAR-10
train_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=train_transform
)
test_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=test_transform
)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4)

# Training setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.05)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)

# Training loop
def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
    
    return total_loss / len(loader), 100. * correct / total

def evaluate(model, loader, device):
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    return 100. * correct / total

# Train for 200 epochs
num_epochs = 200
best_acc = 0

for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
    test_acc = evaluate(model, test_loader, device)
    scheduler.step()
    
    if test_acc > best_acc:
        best_acc = test_acc
    
    if (epoch + 1) 
        print(f"Epoch {epoch+1}/{num_epochs}")
        print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}
        print(f"Test Acc: {test_acc:.2f}

print(f"\nFinal Best Test Accuracy: {best_acc:.2f}

Part (d): Compare to Small ResNet

# Small ResNet for CIFAR-10
class BasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x):
        out = torch.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = torch.relu(out)
        return out

class SmallResNet(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, 1, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        
        self.layer1 = self._make_layer(64, 64, 2, stride=1)
        self.layer2 = self._make_layer(64, 128, 2, stride=2)
        self.layer3 = self._make_layer(128, 256, 2, stride=2)
        
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(256, num_classes)
    
    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
        layers = [BasicBlock(in_channels, out_channels, stride)]
        for _ in range(1, num_blocks):
            layers.append(BasicBlock(out_channels, out_channels, 1))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = torch.relu(self.bn1(self.conv1(x)))
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Create and compare models
resnet = SmallResNet(num_classes=10)
vit_tiny = ViTTiny()

resnet_params = sum(p.numel() for p in resnet.parameters())
vit_params = sum(p.numel() for p in vit_tiny.parameters())

print("Model Comparison:")
print(f"ViT-Tiny parameters: {vit_params:,}")
print(f"Small ResNet parameters: {resnet_params:,}")
print(f"Ratio: {vit_params / resnet_params:.2f}x")

Expected Results:

Analysis:

Part (a): Patch Size 4 for 32×32 Images

• Image: $32 \times 32 \times 3$
• Patch size: $4 \times 4$
• Number of patches: $\frac{32}{4} \times \frac{32}{4} = 8 \times 8 = 64$
• Each patch: $4 \times 4 \times 3 = 48$ values
• Sequence length: $64 + 1 = 65$ (including CLS token)

This is appropriate for CIFAR-10 because:

• Smaller images need smaller patches
• 64 patches provide sufficient spatial resolution
• Comparable to 196 patches for ImageNet (224×224, patch 16)

Part (b): Model Configuration

ViT-Tiny Architecture:

• Layers: $L = 6$
• Hidden size: $d = 192$
• Attention heads: $h = 3$
• MLP ratio: $4.0$ (MLP size = $192 \times 4 = 768$)
• Dropout: $0.1$

Parameter Count:

• Patch embedding: $3 \times 4 \times 4 \times 192 = 9{,}216$
• Position embeddings: $65 \times 192 = 12{,}480$
• Per layer: $4 \times 192^2 + 2 \times 192 \times 768 \approx 443{,}000$
• 6 layers: $6 \times 443{,}000 = 2{,}658{,}000$
• Classification head: $192 \times 10 = 1{,}920$
• Total: $\approx 5.7$M parameters

Part (c): RandAugment

RandAugment Strategy:

• Randomly select $N$ augmentation operations
• Apply with magnitude $M$
• Operations: rotation, shear, color jitter, contrast, etc.
• Typical: $N=2$, $M=9$ for CIFAR-10

Why RandAugment for ViT?

1. Data augmentation is crucial: ViT lacks inductive bias
2. Prevents overfitting: CIFAR-10 is small (50k images)
3. Improves generalization: +2-3\% accuracy improvement
4. Simpler than AutoAugment: No search required

Training Recipe:

• Optimizer: AdamW with weight decay 0.05
• Learning rate: $10^{-3}$ with cosine annealing
• Batch size: 128
• Epochs: 200
• Warmup: 10 epochs (optional)

Part (d): Comparison with Small ResNet

Quantitative Comparison:

Why ResNet Performs Better on CIFAR-10:

1. Inductive bias: Convolutions encode spatial locality
2. Translation equivariance: Built into convolutions
3. Parameter efficiency: Fewer parameters, better generalization
4. Small dataset: CIFAR-10 (50k) is too small for ViT
5. Low resolution: $32 \times 32$ images have limited spatial information

When ViT Would Win:

• Pre-training: Train on ImageNet-21k, fine-tune on CIFAR-10
• Larger dataset: More training data (e.g., 500k images)
• Higher resolution: Upscale CIFAR-10 to $224 \times 224$
• Transfer learning: Use pre-trained ViT weights

Practical Recommendations:

1. From scratch on CIFAR-10: Use ResNet (better accuracy, faster)
2. With pre-training: Use ViT (transfer learning advantage)
3. Research purposes: Try both, compare carefully
4. Production: ResNet for efficiency, ViT for maximum accuracy with pre-training

Key Takeaways:

• ViT requires large-scale pre-training to excel
• Convolutional inductive bias helps on small datasets
• Data augmentation is critical for ViT
• Patch size should scale with image resolution
• ResNet is more sample-efficient on CIFAR-10

Experiment Variations to Try:

1. Increase ViT depth to 12 layers
2. Try different patch sizes (2, 4, 8)
3. Add more aggressive augmentation
4. Use mixup or cutmix
5. Pre-train on CIFAR-100, fine-tune on CIFAR-10
6. Compare with hybrid models (ResNet + Transformer)