Multimodal Transformers

import torch
import torch.nn as nn
import torch.nn.functional as F

def clip_contrastive_loss(image_embeddings, text_embeddings, temperature=0.07):
    """
    Compute CLIP contrastive loss
    Args:
        image_embeddings: (B, D) normalized image embeddings
        text_embeddings: (B, D) normalized text embeddings
        temperature: temperature parameter tau
    Returns:
        loss: scalar contrastive loss
    """
    # Normalize embeddings
    image_embeddings = F.normalize(image_embeddings, dim=-1)
    text_embeddings = F.normalize(text_embeddings, dim=-1)
    
    # Compute similarity matrix (B, B)
    logits = torch.matmul(image_embeddings, text_embeddings.t()) / temperature
    
    # Labels: diagonal elements are positive pairs
    batch_size = image_embeddings.shape[0]
    labels = torch.arange(batch_size, device=image_embeddings.device)
    
    # Symmetric loss: image-to-text + text-to-image
    loss_i2t = F.cross_entropy(logits, labels)
    loss_t2i = F.cross_entropy(logits.t(), labels)
    
    loss = (loss_i2t + loss_t2i) / 2
    
    return loss, logits

# Part (a): Generate random embeddings
batch_size = 8
embed_dim = 512

image_embeddings = torch.randn(batch_size, embed_dim)
text_embeddings = torch.randn(batch_size, embed_dim)

print(f"Image embeddings shape: {image_embeddings.shape}")
print(f"Text embeddings shape: {text_embeddings.shape}")

# Part (b): Normalize embeddings
image_embeddings = F.normalize(image_embeddings, dim=-1)
text_embeddings = F.normalize(text_embeddings, dim=-1)

print(f"\nAfter normalization:")
print(f"Image embedding norms: {torch.norm(image_embeddings, dim=-1)}")
print(f"Text embedding norms: {torch.norm(text_embeddings, dim=-1)}")

# Part (c): Compute similarity matrix
temperature = 0.07
similarity_matrix = torch.matmul(image_embeddings, text_embeddings.t()) / temperature

print(f"\nSimilarity matrix shape: {similarity_matrix.shape}")
print(f"Similarity matrix:\n{similarity_matrix}")

# Part (d): Calculate contrastive loss
loss, logits = clip_contrastive_loss(image_embeddings, text_embeddings, temperature)

print(f"\nContrastive loss: {loss.item():.4f}")
print(f"Logits shape: {logits.shape}")

# Analyze the loss
labels = torch.arange(batch_size)
predictions_i2t = logits.argmax(dim=1)
predictions_t2i = logits.t().argmax(dim=1)

accuracy_i2t = (predictions_i2t == labels).float().mean()
accuracy_t2i = (predictions_t2i == labels).float().mean()

print(f"\nImage-to-Text accuracy: {accuracy_i2t.item():.2
print(f"Text-to-Image accuracy: {accuracy_t2i.item():.2

Mathematical Derivation:

Part (a) \& (b): Embeddings

Image embeddings: $\vI = [\vi_1, \vi_2, \ldots, \vi_8] \in \mathbb{R}^{8 \times 512}$

Text embeddings: $\vT = [\vt_1, \vt_2, \ldots, \vt_8] \in \mathbb{R}^{8 \times 512}$

Normalize to unit sphere: $\hat{\vi}_i = \frac{\vi_i}{\|\vi_i\|_2}, \quad \hat{\vt}_i = \frac{\vt_i}{\|\vt_i\|_2}$

Part (c): Similarity Matrix

Cosine similarity matrix: $\vS_{ij} = \frac{\hat{\vi}_i \cdot \hat{\vt}_j}{\tau}$

where $\tau = 0.07$ is the temperature parameter.

Full matrix: $\vS = \frac{1}{\tau} \hat{\vI} \hat{\vT}^T \in \mathbb{R}^{8 \times 8}$

Example: \[ \vS = \begin{bmatrix} s_{11} & s_{12} & \cdots & s_{18} \\ s_{21} & s_{22} & \cdots & s_{28} \\ \vdots & \vdots & \ddots & \vdots \\ s_{81} & s_{82} & \cdots & s_{88} \end{bmatrix} \]

Diagonal elements $s_{ii}$ are positive pairs (matched image-text).

Off-diagonal elements $s_{ij}$ ($i \neq j$) are negative pairs.

Part (d): Contrastive Loss

Image-to-Text Loss:

For each image $i$, predict its matching text from 8 candidates:

$\mathcal{L}_{i2t} = -\frac{1}{8} \sum_{i=1}^{8} \log \frac{\exp(s_{ii})}{\sum_{j=1}^{8} \exp(s_{ij})}$

This is cross-entropy with labels $y_i = i$ (diagonal).

Text-to-Image Loss:

For each text $j$, predict its matching image from 8 candidates:

$\mathcal{L}_{t2i} = -\frac{1}{8} \sum_{j=1}^{8} \log \frac{\exp(s_{jj})}{\sum_{i=1}^{8} \exp(s_{ij})}$

Total CLIP Loss:

$\mathcal{L}_{\text{CLIP}} = \frac{1}{2}(\mathcal{L}_{i2t} + \mathcal{L}_{t2i})$

Symmetric loss ensures both modalities learn aligned representations.

Why Temperature $\tau = 0.07$?

• Sharpens distribution: Small $\tau$ makes softmax more peaked
• Emphasizes hard negatives: Distinguishes similar but incorrect pairs
• Empirically optimal: Found through hyperparameter search
• Typical range: $\tau \in [0.01, 0.1]$

Effect of temperature:

• $\tau \to 0$: Approaches hard assignment (argmax)
• $\tau \to \infty$: Uniform distribution (no learning)
• $\tau = 0.07$: Good balance for contrastive learning

Numerical Example:

Suppose for image 1:

• $s_{11} = 0.9$ (correct text)
• $s_{12} = 0.3, s_{13} = 0.2, \ldots, s_{18} = 0.1$ (incorrect texts)

Softmax probabilities: $p_1 = \frac{\exp(0.9/0.07)}{\exp(0.9/0.07) + \sum_{j=2}^{8} \exp(s_{1j}/0.07)}$

Loss for image 1: $\ell_1 = -\log p_1$

If $p_1 \approx 1$, then $\ell_1 \approx 0$ (good alignment).

If $p_1 \approx 0.125$ (uniform), then $\ell_1 \approx 2.08$ (poor alignment).

Training Dynamics:

1. Initial: Random embeddings, $\mathcal{L} \approx \log(8) = 2.08$
2. Training: Embeddings align, diagonal elements increase
3. Converged: $s_{ii} \gg s_{ij}$ for $i \neq j$, $\mathcal{L} \to 0$

Key Insights:

• Batch size acts as number of negative samples
• Larger batches improve contrastive learning (more negatives)
• CLIP uses batch sizes up to 32,768 in practice
• Symmetric loss prevents modality collapse
• Temperature is a critical hyperparameter

import torch
import clip
from PIL import Image
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm

# Part (a): Load pre-trained CLIP model
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

print(f"CLIP model loaded on {device}")
print(f"Model: ViT-B/32")

# Part (b): Create text prompts for 10 CIFAR-10 classes
cifar10_classes = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck"
]

# Template-based prompts (improves accuracy)
templates = [
    "a photo of a {}.",
    "a blurry photo of a {}.",
    "a photo of many {}.",
    "a photo of the small {}.",
    "a photo of the large {}.",
]

# Encode text prompts
def encode_text_prompts(model, classes, templates):
    """Encode text prompts with multiple templates"""
    text_features = []
    
    for classname in classes:
        # Create prompts from templates
        texts = [template.format(classname) for template in templates]
        texts = clip.tokenize(texts).to(device)
        
        # Encode texts
        with torch.no_grad():
            class_embeddings = model.encode_text(texts)
            class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True)
            
            # Average over templates
            class_embedding = class_embeddings.mean(dim=0)
            class_embedding = class_embedding / class_embedding.norm()
            
            text_features.append(class_embedding)
    
    text_features = torch.stack(text_features, dim=0)
    return text_features

text_features = encode_text_prompts(model, cifar10_classes, templates)
print(f"\nText features shape: {text_features.shape}")  # (10, 512)

# Part (c): Load CIFAR-10 test set
test_dataset = torchvision.datasets.CIFAR10(
    root='./data', 
    train=False, 
    download=True,
    transform=preprocess
)

test_loader = DataLoader(test_dataset, batch_size=100, shuffle=False)

# Zero-shot classification
def zero_shot_classify(model, loader, text_features):
    """Perform zero-shot classification"""
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in tqdm(loader):
            images = images.to(device)
            labels = labels.to(device)
            
            # Encode images
            image_features = model.encode_image(images)
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            
            # Compute similarity with text features
            similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
            
            # Predict
            predictions = similarity.argmax(dim=-1)
            
            correct += (predictions == labels).sum().item()
            total += labels.size(0)
    
    accuracy = 100.0 * correct / total
    return accuracy

# Part (d): Compute accuracy
zero_shot_accuracy = zero_shot_classify(model, test_loader, text_features)
print(f"\nZero-shot accuracy: {zero_shot_accuracy:.2f}



# Part (e): Compare to supervised baseline
# Train a simple supervised classifier
class SimpleCNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = torch.nn.Conv2d(32, 64, 3, padding=1)
        self.pool = torch.nn.MaxPool2d(2, 2)
        self.fc1 = torch.nn.Linear(64 * 8 * 8, 128)
        self.fc2 = torch.nn.Linear(128, 10)
    
    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Train supervised model (simplified)
supervised_model = SimpleCNN().to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(supervised_model.parameters(), lr=0.001)

# Training loop (10 epochs for quick comparison)
train_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True,
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

for epoch in range(10):
    supervised_model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = supervised_model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# Evaluate supervised model
supervised_model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = supervised_model(images)
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

supervised_accuracy = 100.0 * correct / total

print(f"\nComparison:")
print(f"CLIP Zero-shot: {zero_shot_accuracy:.2f}
print(f"Supervised CNN (10 epochs): {supervised_accuracy:.2f}

Expected Results:

Analysis:

Part (a): Pre-trained CLIP Model

CLIP models available:

• RN50: ResNet-50 image encoder
• ViT-B/32: ViT-Base with patch size 32
• ViT-B/16: ViT-Base with patch size 16 (better)
• ViT-L/14: ViT-Large with patch size 14 (best)

Pre-training:

• Dataset: 400M image-text pairs from internet
• Training: Contrastive learning for 32 epochs
• Batch size: 32,768 (large-scale)
• Compute: 256 V100 GPUs for 12 days

Part (b): Text Prompts

"airplane", "automobile", "bird", ...

"a photo of a airplane."
"a blurry photo of a airplane."
"a photo of many airplanes."

Why templates help:

• Match training distribution (natural sentences)
• Provide context for ambiguous classes
• Ensemble over multiple descriptions
• Improve robustness to variations

Prompt engineering tips:

• Use natural language sentences
• Include domain-specific context
• Try multiple templates and average
• Avoid overly specific descriptions

Part (c): Encoding Process

Image encoding:

1. Preprocess: Resize to $224 \times 224$, normalize
2. ViT encoder: Extract features
3. Projection: Map to shared embedding space (512-dim)
4. Normalize: $\hat{\vi} = \vi / \|\vi\|_2$

Text encoding:

1. Tokenize: Convert text to token IDs
2. Text encoder: Transformer processes tokens
3. Projection: Map to shared embedding space (512-dim)
4. Normalize: $\hat{\vt} = \vt / \|\vt\|_2$

Part (d): Zero-Shot Classification

Algorithm:

For each test image $\vx$:

1. Encode image: $\vi = \text{ImageEncoder}(\vx)$
2. Compute similarity with all class embeddings: $s_k = \vi \cdot \vt_k$ for $k = 1, \ldots, 10$
3. Apply softmax: $p_k = \frac{\exp(s_k / \tau)}{\sum_{j=1}^{10} \exp(s_j / \tau)}$
4. Predict: $\hat{y} = \argmax_k p_k$

Temperature $\tau = 0.01$ (learned during training).

Mathematical Formulation:

$P(y = k | \vx) = \frac{\exp(\text{sim}(\vi, \vt_k) / \tau)}{\sum_{j=1}^{10} \exp(\text{sim}(\vi, \vt_j) / \tau)}$

where $\text{sim}(\vi, \vt) = \vi \cdot \vt$ (cosine similarity after normalization).

Part (e): Comparison with Supervised Baseline

Why CLIP Zero-Shot Outperforms Supervised CNN:

1. Pre-training scale: 400M image-text pairs vs 50k CIFAR-10 images
2. Transfer learning: Leverages knowledge from diverse data
3. Better architecture: ViT-B/32 vs simple CNN
4. Semantic understanding: Learns concepts, not just patterns
5. Robustness: Generalizes better to distribution shifts

When Supervised Wins:

• Sufficient training data: ResNet-50 with 200 epochs reaches 95-96\%
• Domain-specific: Fine-tuned models beat zero-shot on specialized tasks
• Computational constraints: Smaller models are faster

CLIP Advantages:

1. No training required: Instant deployment
2. Flexible: Change classes without retraining
3. Interpretable: Natural language descriptions
4. Robust: Handles distribution shifts better
5. Multimodal: Can do image-text retrieval, captioning, etc.

Practical Recommendations:

Improving CLIP Zero-Shot:

1. Better prompts: Domain-specific templates
2. Larger model: ViT-L/14 instead of ViT-B/32
3. Ensemble: Average predictions from multiple prompts
4. Few-shot: Add a few examples with linear probe
5. Fine-tuning: Adapt to target domain

Key Takeaways:

• CLIP achieves strong zero-shot performance through large-scale pre-training
• Natural language prompts enable flexible classification
• Zero-shot CLIP often matches or exceeds supervised baselines
• Prompt engineering is crucial for optimal performance
• CLIP's multimodal nature enables many downstream tasks

Part (a): Encoder Parameters (24 layers, $d=1024$)

Whisper Encoder Configuration:

• Layers: $L = 24$
• Hidden size: $d = 1024$
• Attention heads: $h = 16$
• MLP ratio: $4.0$ (MLP size = $4096$)
• Audio features: 80-dimensional log-mel spectrogram
• Sequence length: $T = 3000$ (30 seconds at 100 Hz)

Parameter Breakdown:

1. Input Convolution Layers:

• Conv1: $80 \times 3 \times 1024 = 245{,}760$
• Conv2: $1024 \times 3 \times 1024 = 3{,}145{,}728$
• Total: $3{,}391{,}488$ parameters

2. Position Embeddings:

• Sinusoidal (not learned): 0 parameters

3. Per Transformer Layer:

Multi-Head Attention:

• $Q, K, V$ projections: $3 \times 1024^2 = 3{,}145{,}728$
• Output projection: $1024^2 = 1{,}048{,}576$
• Total attention: $4{,}194{,}304$

MLP:

• First linear: $1024 \times 4096 = 4{,}194{,}304$
• Second linear: $4096 \times 1024 = 4{,}194{,}304$
• Total MLP: $8{,}388{,}608$

Layer Normalization:

• 2 LayerNorms: $2 \times 2 \times 1024 = 4{,}096$

Total per layer: $12{,}587{,}008$ parameters

4. All 24 Encoder Layers: $24 \times 12{,}587{,}008 = 302{,}088{,}192$ parameters

Total Encoder: $\approx 305.5$M parameters

Part (b): Decoder Parameters (24 layers)

Whisper Decoder Configuration:

• Layers: $L = 24$
• Hidden size: $d = 1024$
• Attention heads: $h = 16$
• Vocabulary size: $V = 51{,}865$
• Max sequence length: $448$ tokens

Parameter Breakdown:

1. Token Embedding:

• $51{,}865 \times 1024 = 53{,}109{,}760$ parameters

2. Position Embeddings:

• $448 \times 1024 = 458{,}752$ parameters

3. Per Decoder Layer:

Masked Self-Attention:

• Same as encoder: $4{,}194{,}304$ parameters

Cross-Attention:

• $Q$ projection: $1024^2 = 1{,}048{,}576$
• $K, V$ projections (from encoder): $2 \times 1024^2 = 2{,}097{,}152$
• Output projection: $1024^2 = 1{,}048{,}576$
• Total cross-attention: $4{,}194{,}304$

MLP:

• Same as encoder: $8{,}388{,}608$ parameters

Layer Normalization:

• 3 LayerNorms: $3 \times 2 \times 1024 = 6{,}144$

Total per decoder layer: $16{,}783{,}360$ parameters

4. All 24 Decoder Layers: $24 \times 16{,}783{,}360 = 402{,}800{,}640$ parameters

5. Output Projection:

• Shared with token embedding: 0 additional parameters

Total Decoder: $\approx 456.4$M parameters

Total Whisper Model: $305.5 + 456.4 = 761.9$M parameters

(Actual Whisper-large: $\approx 1.55$B parameters due to additional components)

Part (c): Memory for 30-Second Audio

Input Processing:

1. Audio Preprocessing:

• Sample rate: 16 kHz
• 30 seconds: $30 \times 16{,}000 = 480{,}000$ samples
• Raw audio: $480{,}000 \times 4$ bytes = 1.92 MB

2. Log-Mel Spectrogram:

• Window size: 25 ms (400 samples)
• Hop length: 10 ms (160 samples)
• Number of frames: $\frac{480{,}000}{160} = 3{,}000$
• Mel bins: 80
• Features: $3{,}000 \times 80 = 240{,}000$ values
• Memory: $240{,}000 \times 4$ bytes = 0.96 MB

Encoder Memory (Inference):

1. Activations per layer:

• Input: $3{,}000 \times 1024 = 3{,}072{,}000$ values
• Attention scores: $16 \times 3{,}000 \times 3{,}000 = 144{,}000{,}000$ values
• MLP intermediate: $3{,}000 \times 4096 = 12{,}288{,}000$ values

Peak per layer: $\approx 159$M values $\times$ 4 bytes = 636 MB

2. Total encoder activations: $24 \times 636$ MB = 15.3 GB (if storing all layers)

With activation checkpointing: $\approx 1.3$ GB

Decoder Memory (Inference):

For generating 448 tokens:

• Decoder activations: $448 \times 1024 = 458{,}752$ values per layer
• Cross-attention: $448 \times 3{,}000 = 1{,}344{,}000$ values per layer
• KV cache: $2 \times 24 \times 448 \times 1024 = 22{,}020{,}096$ values

Decoder memory: $\approx 500$ MB

Total Memory (Inference):

• Model parameters: $1.55$B $\times$ 4 bytes = 6.2 GB
• Encoder activations: $\approx 1.3$ GB (with checkpointing)
• Decoder activations: $\approx 0.5$ GB
• KV cache: $\approx 0.1$ GB
• Total: $\approx 8.1$ GB

For FP16: $\approx 4.1$ GB

For INT8 quantization: $\approx 2.1$ GB

Part (d): Compare to Text-Only GPT-2

GPT-2 (1.5B parameters):

• Layers: 48
• Hidden size: 1600
• Attention heads: 25
• Vocabulary: 50,257
• Context length: 1024 tokens

Comparison Table:

Key Differences:

1. Architecture:
   • Whisper: Encoder-decoder (like T5)
   • GPT-2: Decoder-only
2. Input Processing:
   • Whisper: Audio $\to$ Log-mel $\to$ Encoder
   • GPT-2: Text $\to$ Tokens $\to$ Decoder
3. Computational Cost:
   • Whisper encoder: $O(T^2 d)$ where $T = 3000$
   • GPT-2: $O(n^2 d)$ where $n = 1024$
   • Whisper is $\approx 9\times$ more expensive for encoder
4. Memory Footprint:
   • Whisper: Larger due to long audio sequences
   • GPT-2: Smaller, text-only
5. Use Cases:
   • Whisper: Speech recognition, translation, transcription
   • GPT-2: Text generation, completion, summarization

Why Whisper Needs Encoder-Decoder:

• Cross-modal: Audio input, text output
• Compression: Encoder compresses 3000 audio frames
• Attention: Decoder attends to compressed audio
• Efficiency: Encoder processes audio once, decoder generates text autoregressively

Performance Comparison:

Practical Considerations:

1. Deployment:
   • Whisper: Requires audio preprocessing
   • GPT-2: Simple tokenization
2. Latency:
   • Whisper: Higher (audio encoding + decoding)
   • GPT-2: Lower (text-only)
3. Hardware:
   • Whisper: Needs GPU for real-time (8+ GB VRAM)
   • GPT-2: Can run on CPU for small batches

Key Insights:

• Whisper and GPT-2 have similar parameter counts but different architectures
• Encoder-decoder is essential for cross-modal tasks
• Audio sequences are much longer than text, requiring more memory
• Both models benefit from large-scale pre-training
• Whisper's multimodal nature enables speech-to-text applications

Part (a): Proposed Architecture

import torch
import torch.nn as nn

class MultimodalVideoTransformer(nn.Module):
    def __init__(self, 
                 visual_dim=768,      # ViT features
                 audio_dim=512,       # Audio features
                 text_dim=768,        # BERT features
                 hidden_dim=1024,     # Fusion dimension
                 num_layers=12,       # Fusion transformer layers
                 num_heads=16,
                 num_classes=400):    # Action recognition classes
        super().__init__()
        
        # Modality-specific encoders
        self.visual_encoder = VisualEncoder(visual_dim, hidden_dim)
        self.audio_encoder = AudioEncoder(audio_dim, hidden_dim)
        self.text_encoder = TextEncoder(text_dim, hidden_dim)
        
        # Modality-specific tokens
        self.visual_token = nn.Parameter(torch.randn(1, 1, hidden_dim))
        self.audio_token = nn.Parameter(torch.randn(1, 1, hidden_dim))
        self.text_token = nn.Parameter(torch.randn(1, 1, hidden_dim))
        
        # Fusion transformer
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            batch_first=True
        )
        self.fusion_transformer = nn.TransformerEncoder(
            encoder_layer, 
            num_layers=num_layers
        )
        
        # Classification head
        self.classifier = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, num_classes)
        )
    
    def forward(self, visual_features, audio_features, text_features):
        """
        Args:
            visual_features: (B, T_v, D_v) - video frames
            audio_features: (B, T_a, D_a) - audio segments
            text_features: (B, T_t, D_t) - caption tokens
        Returns:
            logits: (B, num_classes)
        """
        B = visual_features.shape[0]
        
        # Encode each modality
        visual_emb = self.visual_encoder(visual_features)  # (B, T_v, H)
        audio_emb = self.audio_encoder(audio_features)     # (B, T_a, H)
        text_emb = self.text_encoder(text_features)        # (B, T_t, H)
        
        # Add modality tokens
        visual_token = self.visual_token.expand(B, -1, -1)
        audio_token = self.audio_token.expand(B, -1, -1)
        text_token = self.text_token.expand(B, -1, -1)
        
        visual_emb = torch.cat([visual_token, visual_emb], dim=1)
        audio_emb = torch.cat([audio_token, audio_emb], dim=1)
        text_emb = torch.cat([text_token, text_emb], dim=1)
        
        # Concatenate all modalities
        multimodal_emb = torch.cat([visual_emb, audio_emb, text_emb], dim=1)
        # Shape: (B, 1+T_v + 1+T_a + 1+T_t, H)
        
        # Fusion transformer
        fused = self.fusion_transformer(multimodal_emb)
        
        # Aggregate: average modality tokens
        visual_rep = fused[:, 0, :]
        audio_rep = fused[:, 1+visual_features.shape[1], :]
        text_rep = fused[:, 1+visual_features.shape[1]+1+audio_features.shape[1], :]
        
        # Combine representations
        combined = (visual_rep + audio_rep + text_rep) / 3
        
        # Classification
        logits = self.classifier(combined)
        
        return logits

class VisualEncoder(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.proj = nn.Linear(input_dim, output_dim)
        self.norm = nn.LayerNorm(output_dim)
    
    def forward(self, x):
        return self.norm(self.proj(x))

class AudioEncoder(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.proj = nn.Linear(input_dim, output_dim)
        self.norm = nn.LayerNorm(output_dim)
    
    def forward(self, x):
        return self.norm(self.proj(x))

class TextEncoder(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.proj = nn.Linear(input_dim, output_dim)
        self.norm = nn.LayerNorm(output_dim)
    
    def forward(self, x):
        return self.norm(self.proj(x))

# Example usage
model = MultimodalVideoTransformer()

# Simulate inputs
batch_size = 4
visual = torch.randn(batch_size, 16, 768)   # 16 frames
audio = torch.randn(batch_size, 32, 512)    # 32 audio segments
text = torch.randn(batch_size, 20, 768)     # 20 caption tokens

logits = model(visual, audio, text)
print(f"Output shape: {logits.shape}")  # (4, 400)

Part (b): Fusion Mechanism

Architecture Overview:

Input:
  Visual: (B, 16, 768)  - 16 video frames from ViT
  Audio:  (B, 32, 512)  - 32 audio segments from audio encoder
  Text:   (B, 20, 768)  - 20 caption tokens from BERT

Step 1: Modality-Specific Projection
  Visual -> (B, 16, 1024)
  Audio  -> (B, 32, 1024)
  Text   -> (B, 20, 1024)

Step 2: Add Modality Tokens
  Visual: [V_token, v1, v2, ..., v16]  -> (B, 17, 1024)
  Audio:  [A_token, a1, a2, ..., a32]  -> (B, 33, 1024)
  Text:   [T_token, t1, t2, ..., t20]  -> (B, 21, 1024)

Step 3: Concatenate
  Multimodal: [V_token, v1, ..., v16, A_token, a1, ..., a32, T_token, t1, ..., t20]
  Shape: (B, 71, 1024)

Step 4: Fusion Transformer (12 layers)
  Cross-modal attention enables interaction
  Output: (B, 71, 1024)

Step 5: Aggregate
  Extract modality tokens: V_token, A_token, T_token
  Average: (V_token + A_token + T_token) / 3
  Shape: (B, 1024)

Step 6: Classification
  MLP: (B, 1024) -> (B, 400)

Fusion Strategies Comparison:

1. Early Fusion (Concatenation):
   • Concatenate features before transformer
   • Simple but limited cross-modal interaction
   • Used in this design
2. Late Fusion (Ensemble):
   • Process modalities separately
   • Combine predictions at the end
   • No cross-modal learning
3. Cross-Modal Attention:
   • Visual attends to audio and text
   • Audio attends to visual and text
   • More complex but better interaction
4. Bottleneck Fusion:
   • Compress each modality to bottleneck tokens
   • Fuse bottlenecks
   • More efficient for long sequences

Why This Design:

• Modality tokens: Aggregate information from each modality
• Shared transformer: Enables cross-modal attention
• Flexible: Can handle missing modalities
• Scalable: Easy to add more modalities

Part (c): Training Objective

Primary Objective: Action Recognition

$\mathcal{L}_{\text{action}} = -\frac{1}{B} \sum_{i=1}^{B} \log P(y_i | \vv_i, \va_i, \vt_i)$

where:

• $\vv_i$: visual features for sample $i$
• $\va_i$: audio features for sample $i$
• $\vt_i$: text features for sample $i$
• $y_i$: ground truth action class

Auxiliary Objectives (Multi-Task Learning):

1. Contrastive Loss (Cross-Modal Alignment):

Align visual-audio, visual-text, audio-text pairs:

$\mathcal{L}_{\text{contrast}} = \mathcal{L}_{\text{VA}} + \mathcal{L}_{\text{VT}} + \mathcal{L}_{\text{AT}}$

where each term is CLIP-style contrastive loss:

$\mathcal{L}_{\text{VA}} = -\frac{1}{B} \sum_{i=1}^{B} \log \frac{\exp(\text{sim}(\vv_i, \va_i) / \tau)}{\sum_{j=1}^{B} \exp(\text{sim}(\vv_i, \va_j) / \tau)}$

2. Masked Modality Modeling:

Randomly mask one modality and predict it from others:

$\mathcal{L}_{\text{mask}} = \mathcal{L}_{\text{mask-V}} + \mathcal{L}_{\text{mask-A}} + \mathcal{L}_{\text{mask-T}}$

Example (mask visual): $\mathcal{L}_{\text{mask-V}} = \|\hat{\vv} - \vv\|_2^2$

where $\hat{\vv} = f(\va, \vt)$ is predicted visual features.

3. Temporal Ordering:

Predict correct temporal order of video segments:

$\mathcal{L}_{\text{temporal}} = -\log P(\text{order} | \vv, \va, \vt)$

Total Training Objective:

$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{action}} + \lambda_1 \mathcal{L}_{\text{contrast}} + \lambda_2 \mathcal{L}_{\text{mask}} + \lambda_3 \mathcal{L}_{\text{temporal}}$

Typical weights: $\lambda_1 = 0.1$, $\lambda_2 = 0.05$, $\lambda_3 = 0.05$

Training Recipe:

# Optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)

# Learning rate schedule
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

# Training loop
for epoch in range(100):
    for batch in dataloader:
        visual, audio, text, labels = batch
        
        # Forward pass
        logits = model(visual, audio, text)
        
        # Action recognition loss
        loss_action = F.cross_entropy(logits, labels)
        
        # Contrastive loss (optional)
        visual_rep = model.get_visual_rep(visual)
        audio_rep = model.get_audio_rep(audio)
        loss_contrast = contrastive_loss(visual_rep, audio_rep)
        
        # Total loss
        loss = loss_action + 0.1 * loss_contrast
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    scheduler.step()

Data Augmentation:

• Visual: Random crop, color jitter, temporal sampling
• Audio: Time stretching, pitch shifting, noise injection
• Text: Synonym replacement, back-translation
• Multimodal: Random modality dropout (robustness)

Part (d): Parameter Count Estimation

Component Breakdown:

1. Modality-Specific Encoders:

Visual Encoder:

• Projection: $768 \times 1024 = 786{,}432$
• LayerNorm: $2 \times 1024 = 2{,}048$
• Total: $788{,}480$

Audio Encoder:

• Projection: $512 \times 1024 = 524{,}288$
• LayerNorm: $2 \times 1024 = 2{,}048$
• Total: $526{,}336$

Text Encoder:

• Projection: $768 \times 1024 = 786{,}432$
• LayerNorm: $2 \times 1024 = 2{,}048$
• Total: $788{,}480$

Encoder total: $2{,}103{,}296$ parameters

2. Modality Tokens:

• 3 tokens $\times$ 1024 = $3{,}072$ parameters

3. Fusion Transformer (12 layers):

Per layer:

• Self-attention: $4 \times 1024^2 = 4{,}194{,}304$
• MLP: $2 \times 1024 \times 4096 = 8{,}388{,}608$
• LayerNorm: $2 \times 2 \times 1024 = 4{,}096$
• Total per layer: $12{,}587{,}008$

12 layers: $12 \times 12{,}587{,}008 = 151{,}044{,}096$ parameters

4. Classification Head:

• LayerNorm: $2 \times 1024 = 2{,}048$
• Linear 1: $1024 \times 1024 = 1{,}048{,}576$
• Linear 2: $1024 \times 400 = 409{,}600$
• Total: $1{,}460{,}224$

Total Model Parameters:

$2{,}103{,}296 + 3{,}072 + 151{,}044{,}096 + 1{,}460{,}224 = 154{,}610{,}688$

Total: $\approx 155$M parameters

Memory Footprint (FP32):

• Parameters: $155$M $\times$ 4 bytes = 620 MB
• Activations (batch size 4):
  • Input: $4 \times 71 \times 1024 = 290{,}816$ values
  • Per layer: $\approx 2$M values
  • Total: $\approx 24$M values $\times$ 4 bytes = 96 MB
• Gradients: 620 MB (same as parameters)
• Optimizer states (AdamW): $2 \times 620$ MB = 1.24 GB

Total training memory: $\approx 2.6$ GB

Comparison with Baselines:

Design Trade-offs:

1. Parameter efficiency:
   • Shared fusion transformer reduces parameters
   • Modality-specific encoders are lightweight
   • Could use pre-trained encoders (ViT, BERT, etc.)
2. Computational cost:
   • Sequence length: 71 tokens (manageable)
   • Attention complexity: $O(71^2 \times 1024) \approx 5$M operations
   • Inference time: $\approx 50$ ms on GPU
3. Scalability:
   • Easy to add more modalities (depth, optical flow, etc.)
   • Can increase fusion layers for better interaction
   • Bottleneck fusion for longer sequences

Practical Recommendations:

1. Use pre-trained encoders: ViT for visual, Wav2Vec for audio, BERT for text
2. Freeze encoders initially: Train fusion transformer first
3. Fine-tune end-to-end: Unfreeze all parameters later
4. Modality dropout: Randomly drop modalities during training for robustness
5. Temporal modeling: Add temporal attention for video sequences

Key Insights:

• Multimodal fusion requires careful architecture design
• Modality tokens enable flexible aggregation
• Shared transformer enables cross-modal learning
• Multi-task learning improves representation quality
• Parameter count is reasonable for modern GPUs
• Pre-trained encoders significantly improve performance