16 Lab 16: Multimodal & Diffusion Models
16.1 Learning Objectives
- Understand core concepts from Chapter 16
- Implement key algorithms and techniques
- Apply methods to practical problems
- Analyze results and draw insights
- Connect to broader NeuroAI themes
16.2 Prerequisites
- Reading: Chapter 16: Multimodal Models
- Libraries: NumPy, Matplotlib, relevant frameworks
- Concepts: Vision-language integration
16.3 Setup
Run this cell to import the required libraries and configure the environment:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
np.random.seed(42)
plt.rcParams['figure.figsize'] = (12, 8)16.4 Part 1: Vision Transformer (ViT) Architecture
The Vision Transformer (ViT) is a foundational architecture that enables treating images as sequences, allowing Transformer models (originally designed for text) to process visual data. This section visualizes how ViT transforms an image into a sequence of patches for processing.
16.4.1 Exercise 1: Visualize the ViT Pipeline
Run this cell to generate a visual diagram of the Vision Transformer architecture:
def visualize_vit_architecture():
"""Visualize the Vision Transformer architecture.
This function creates a schematic diagram showing how ViT processes images:
1. Input image (224×224) is split into patches (16×16 each)
2. Patches are flattened and projected to embeddings
3. Position embeddings and [CLS] token are added
4. Sequence is processed through Transformer encoder (12 layers)
5. [CLS] token representation is used for classification
Returns:
matplotlib.pyplot: Figure object displaying the architecture
"""
fig, ax = plt.subplots(figsize=(14, 7))
ax.set_xlim(0, 14)
ax.set_ylim(0, 7)
ax.axis('off')
# Colors
img_color = '#e3f2fd'
patch_color = '#bbdefb'
embed_color = '#c8e6c9'
transformer_color = '#fff3e0'
output_color = '#f3e5f5'
# Original Image
img_rect = mpatches.FancyBboxPatch((0.5, 2.5), 2, 2, boxstyle="round,pad=0.05",
facecolor=img_color, edgecolor='#1565c0', linewidth=2)
ax.add_patch(img_rect)
# Draw grid on image
for i in range(4):
ax.plot([0.5 + i*0.5, 0.5 + i*0.5], [2.5, 4.5], 'k-', alpha=0.3, linewidth=0.5)
ax.plot([0.5, 2.5], [2.5 + i*0.5, 2.5 + i*0.5], 'k-', alpha=0.3, linewidth=0.5)
ax.text(1.5, 4.8, 'Image\n224x224', ha='center', fontsize=9)
# Arrow
ax.annotate('', xy=(3.2, 3.5), xytext=(2.7, 3.5),
arrowprops=dict(arrowstyle='->', color='#666', lw=2))
ax.text(2.95, 3.9, 'Split into\npatches', ha='center', fontsize=8)
# Patches
for i, y in enumerate([4.0, 3.3, 2.6]):
for j, x in enumerate([3.4, 3.9, 4.4]):
if i < 2 or j < 2:
rect = mpatches.FancyBboxPatch((x, y), 0.4, 0.4, boxstyle="round,pad=0.02",
facecolor=patch_color, edgecolor='#1976d2', linewidth=1)
ax.add_patch(rect)
ax.text(4.4, 2.3, '...', fontsize=12)
ax.text(4.15, 4.8, 'Patches\n16x16 each', ha='center', fontsize=9)
# Arrow to embeddings
ax.annotate('', xy=(5.5, 3.5), xytext=(5.0, 3.5),
arrowprops=dict(arrowstyle='->', color='#666', lw=2))
ax.text(5.25, 3.9, 'Flatten\n& Project', ha='center', fontsize=8)
# Embeddings (vertical sequence)
embed_y = [4.3, 3.8, 3.3, 2.8, 2.3]
labels = ['[CLS]', 'P1', 'P2', 'P3', '...']
colors = ['#ffcc80'] + ['#c8e6c9']*4
for i, (y, label, color) in enumerate(zip(embed_y, labels, colors)):
rect = mpatches.FancyBboxPatch((5.7, y), 0.8, 0.4, boxstyle="round,pad=0.02",
facecolor=color, edgecolor='#388e3c', linewidth=1)
ax.add_patch(rect)
ax.text(6.1, y+0.2, label, ha='center', va='center', fontsize=8)
ax.text(6.1, 4.9, 'Embeddings\n+ Positions', ha='center', fontsize=9)
# Arrow to Transformer
ax.annotate('', xy=(7.3, 3.5), xytext=(6.8, 3.5),
arrowprops=dict(arrowstyle='->', color='#666', lw=2))
# Transformer block
trans_rect = mpatches.FancyBboxPatch((7.5, 2.0), 2.5, 3.0, boxstyle="round,pad=0.1",
facecolor=transformer_color, edgecolor='#f57c00', linewidth=2)
ax.add_patch(trans_rect)
ax.text(8.75, 4.7, 'Transformer\nEncoder', ha='center', fontsize=10, fontweight='bold')
# Inside transformer
for i, (y, txt) in enumerate([(3.8, 'Multi-Head\nAttention'), (3.0, 'FFN'), (2.4, 'x12 layers')]):
ax.text(8.75, y, txt, ha='center', va='center', fontsize=8)
if i < 2:
ax.plot([7.8, 9.7], [y-0.25, y-0.25], 'k--', alpha=0.3)
# Arrow to output
ax.annotate('', xy=(10.7, 3.5), xytext=(10.2, 3.5),
arrowprops=dict(arrowstyle='->', color='#666', lw=2))
# Output
out_rect = mpatches.FancyBboxPatch((10.9, 2.8), 1.5, 1.4, boxstyle="round,pad=0.1",
facecolor=output_color, edgecolor='#7b1fa2', linewidth=2)
ax.add_patch(out_rect)
ax.text(11.65, 3.5, '[CLS]\nEmbedding', ha='center', va='center', fontsize=9)
ax.text(11.65, 4.5, 'Image\nRepresentation', ha='center', fontsize=9)
# Arrow to classification
ax.annotate('', xy=(13.0, 3.5), xytext=(12.5, 3.5),
arrowprops=dict(arrowstyle='->', color='#666', lw=2))
ax.text(13.2, 3.5, 'Class\nLabel', ha='center', va='center', fontsize=9)
ax.set_title('Vision Transformer (ViT) Architecture', fontsize=14, fontweight='bold', pad=20)
plt.tight_layout()
return plt
# Run this cell to display the ViT architecture visualization
visualize_vit_architecture()Discussion Questions: 1. Why does ViT split the image into patches rather than processing individual pixels? 2. What is the purpose of the [CLS] token in the architecture? 3. How does the lack of inductive biases in ViT compare to CNNs?
16.4.2 Exercise 2: Explore Patch Size Effects
Modify the visualization to understand how different patch sizes affect the number of tokens:
def calculate_vit_tokens(image_size=224, patch_size=16):
"""Calculate the number of tokens produced by ViT patch embedding.
Args:
image_size: Input image dimension (assumes square image)
patch_size: Size of each patch (assumes square patches)
Returns:
tuple: (num_patches, sequence_length including [CLS] token)
"""
num_patches = (image_size // patch_size) ** 2
sequence_length = num_patches + 1 # +1 for [CLS] token
return num_patches, sequence_length
# Compare different patch sizes
patch_sizes = [8, 16, 32]
print("Patch Size | Num Patches | Sequence Length")
print("-" * 45)
for ps in patch_sizes:
num_patches, seq_len = calculate_vit_tokens(patch_size=ps)
print(f" {ps}x{ps} | {num_patches} | {seq_len}")Questions: - What is the trade-off between smaller and larger patch sizes? - How does sequence length affect computational complexity?
16.5 Part 2: Hands-On Implementation
Practical exercises exploring vision-language integration.
16.5.1 Exercise 3: Cross-Attention Mechanism
Implement a simplified cross-attention mechanism that allows text to query image features:
def softmax(x, axis=-1):
"""Numerically stable softmax."""
exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
def cross_attention(text_query, image_keys, image_values):
"""Compute cross-attention between text and image modalities.
This demonstrates how text queries can selectively attend to
different image patches, similar to how Stable Diffusion conditions
image generation on text prompts.
Args:
text_query: Query vector from text encoder (shape: [d])
image_keys: Key vectors from image encoder (shape: [num_patches, d])
image_values: Value vectors from image encoder (shape: [num_patches, d])
Returns:
tuple: (attended_output, attention_weights)
"""
# Compute attention scores: Q @ K^T
d_k = len(text_query)
scores = np.dot(image_keys, text_query) / np.sqrt(d_k)
# Apply softmax to get attention weights
attention_weights = softmax(scores)
# Weighted sum of values: weights @ V
attended_output = np.dot(attention_weights, image_values)
return attended_output, attention_weights
# Example: Text query attending to image patches
np.random.seed(42)
d = 64 # embedding dimension
num_patches = 16 # 4x4 grid of patches
# Simulated embeddings (in practice, these come from trained encoders)
text_query = np.random.randn(d) # "What animal is this?"
image_keys = np.random.randn(num_patches, d)
image_values = np.random.randn(num_patches, d)
output, weights = cross_attention(text_query, image_keys, image_values)
print(f"Text query shape: {text_query.shape}")
print(f"Image keys shape: {image_keys.shape}")
print(f"Attention weights shape: {weights.shape}")
print(f"Attended output shape: {output.shape}")
print(f"\nAttention weights (sum to 1): {weights.sum():.4f}")16.5.2 Exercise 4: CLIP-Style Contrastive Learning
Implement a simplified version of CLIP’s contrastive loss:
def clip_contrastive_loss(image_embeddings, text_embeddings, temperature=0.07):
"""Compute CLIP-style contrastive loss for image-text alignment.
This loss maximizes similarity between matching image-text pairs
while minimizing similarity between non-matching pairs.
Args:
image_embeddings: Normalized image features (batch_size, embed_dim)
text_embeddings: Normalized text features (batch_size, embed_dim)
temperature: Temperature parameter for scaling logits
Returns:
tuple: (total_loss, image_to_text_loss, text_to_image_loss)
"""
batch_size = image_embeddings.shape[0]
# Normalize embeddings (important for cosine similarity)
image_embeddings = image_embeddings / np.linalg.norm(image_embeddings, axis=1, keepdims=True)
text_embeddings = text_embeddings / np.linalg.norm(text_embeddings, axis=1, keepdims=True)
# Compute similarity matrix (scaled by temperature)
logits = np.dot(image_embeddings, text_embeddings.T) / temperature
# Target: diagonal (matching pairs)
targets = np.arange(batch_size)
# Softmax cross-entropy for image-to-text direction
i2t_probs = softmax(logits, axis=1)
i2t_loss = -np.mean(np.log(i2t_probs[np.arange(batch_size), targets] + 1e-8))
# Softmax cross-entropy for text-to-image direction
t2i_probs = softmax(logits.T, axis=1)
t2i_loss = -np.mean(np.log(t2i_probs[np.arange(batch_size), targets] + 1e-8))
# Total loss is average of both directions
total_loss = (i2t_loss + t2i_loss) / 2
return total_loss, i2t_loss, t2i_loss
# Simulate a batch of image-text pairs
batch_size = 4
embed_dim = 512
np.random.seed(42)
image_emb = np.random.randn(batch_size, embed_dim)
text_emb = np.random.randn(batch_size, embed_dim)
# Add correlation between matching pairs (simulating learned alignment)
for i in range(batch_size):
text_emb[i] += 0.5 * image_emb[i]
loss, i2t, t2i = clip_contrastive_loss(image_emb, text_emb)
print(f"Total contrastive loss: {loss:.4f}")
print(f"Image-to-text loss: {i2t:.4f}")
print(f"Text-to-image loss: {t2i:.4f}")16.6 Part 3: Diffusion Models
Understanding the denoising process in generative AI.
16.6.1 Exercise 5: Forward Diffusion Process
Implement the forward diffusion process that gradually adds noise to images:
def forward_diffusion_step(x_t, noise, beta_t):
"""Perform one step of the forward diffusion process.
The forward process gradually adds Gaussian noise:
x_{t+1} = sqrt(1 - beta_t) * x_t + sqrt(beta_t) * noise
Args:
x_t: Current state (image or latent)
noise: Standard Gaussian noise
beta_t: Noise schedule coefficient at step t
Returns:
x_{t+1}: Next (noisier) state
"""
alpha_t = 1 - beta_t
return np.sqrt(alpha_t) * x_t + np.sqrt(beta_t) * noise
def forward_diffusion(x_0, num_steps, beta_schedule):
"""Run the complete forward diffusion process.
Args:
x_0: Original clean data
num_steps: Number of diffusion steps (T)
beta_schedule: Array of beta values for each step
Returns:
list: All intermediate states [x_0, x_1, ..., x_T]
"""
trajectory = [x_0.copy()]
x_t = x_0.copy()
for t in range(num_steps):
noise = np.random.randn(*x_t.shape)
x_t = forward_diffusion_step(x_t, noise, beta_schedule[t])
trajectory.append(x_t.copy())
return trajectory
# Visualize diffusion on 1D data
np.random.seed(42)
x_0 = np.array([2.0, -1.5, 1.0, -0.5]) # Original "clean" data points
num_steps = 100
beta_schedule = np.linspace(0.001, 0.02, num_steps) # Linear schedule
trajectory = forward_diffusion(x_0, num_steps, beta_schedule)
# Plot the trajectory
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Left: Individual trajectories
for i in range(len(x_0)):
axes[0].plot([t[i] for t in trajectory], alpha=0.7, label=f'Dimension {i+1}')
axes[0].set_xlabel('Diffusion Step (t)')
axes[0].set_ylabel('Value')
axes[0].set_title('Forward Diffusion: Data to Noise')
axes[0].legend()
axes[0].axhline(y=0, color='k', linestyle='--', alpha=0.3)
# Right: Snapshots at different timesteps
snapshot_steps = [0, 25, 50, 75, 100]
for idx, step in enumerate(snapshot_steps):
axes[1].scatter([idx] * len(trajectory[step]), trajectory[step],
alpha=0.7, label=f't={step}')
axes[1].set_xlabel('Snapshot Index')
axes[1].set_ylabel('Value')
axes[1].set_title('Data Distribution at Different Timesteps')
axes[1].legend()
plt.tight_layout()
plt.show()16.7 Part 4: Analysis
Analyze and visualize results.
16.7.1 Exercise 6: Understanding Classifier-Free Guidance
Implement classifier-free guidance to control generation strength:
def classifier_free_guidance(noise_uncond, noise_cond, guidance_scale):
"""Apply classifier-free guidance to noise predictions.
This technique pushes the generation toward the conditioned
prediction while maintaining diversity from the unconditional baseline.
Formula: noise_guided = noise_uncond + w * (noise_cond - noise_uncond)
Args:
noise_uncond: Noise prediction without text conditioning
noise_cond: Noise prediction with text conditioning
guidance_scale: Guidance strength (w), typically 5-15
Returns:
Guided noise prediction
"""
return noise_uncond + guidance_scale * (noise_cond - noise_uncond)
# Visualize the effect of guidance scale
np.random.seed(42)
noise_uncond = np.random.randn(2) * 0.5 # Generic prediction
noise_cond = np.random.randn(2) * 0.5 + np.array([1.0, 0.5]) # Text-conditioned
guidance_scales = [1.0, 5.0, 10.0, 20.0]
colors = plt.cm.viridis(np.linspace(0.2, 0.9, len(guidance_scales)))
fig, ax = plt.subplots(figsize=(10, 8))
# Plot unconditional and conditional predictions
ax.scatter(*noise_uncond, s=200, c='blue', marker='o', label='Unconditional', zorder=5)
ax.scatter(*noise_cond, s=200, c='red', marker='s', label='Conditional', zorder=5)
# Plot guided predictions for different scales
for i, w in enumerate(guidance_scales):
guided = classifier_free_guidance(noise_uncond, noise_cond, w)
ax.scatter(*guided, s=150, c=[colors[i]], marker='^', label=f'w={w}')
# Draw arrow from uncond to guided
ax.annotate('', xy=guided, xytext=noise_uncond,
arrowprops=dict(arrowstyle='->', color=colors[i], alpha=0.5, lw=1.5))
ax.axhline(y=0, color='k', linestyle='--', alpha=0.3)
ax.axvline(x=0, color='k', linestyle='--', alpha=0.3)
ax.set_xlabel('Noise Dimension 1')
ax.set_ylabel('Noise Dimension 2')
ax.set_title('Effect of Classifier-Free Guidance Scale')
ax.legend(loc='upper left')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
print("Guidance Scale Interpretation:")
print(" w=1.0: Uses conditional prediction directly (weak prompt adherence)")
print(" w=5-10: Sweet spot (good balance of quality and prompt following)")
print(" w>15: Strong prompt adherence but may produce artifacts")16.8 Exercises
16.8.1 Exercise 1
Implement the basic ViT patch embedding algorithm.
16.8.2 Exercise 2
Explore parameter variations in the diffusion process (different beta schedules).
16.8.3 Exercise 3
Apply cross-attention to real image features using a pre-trained vision model.
16.8.4 Exercise 4
Compare different fusion strategies (concatenation vs. cross-attention vs. contrastive).
16.8.5 Exercise 5
Discuss NeuroAI connections: How does multimodal integration in these models compare to multisensory integration in the brain (superior temporal sulcus, parietal cortex)?
16.9 Challenge Problems
16.9.1 Challenge 1
Implement a minimal CLIP model from scratch using PyTorch and train it on a small image-text dataset (e.g., CIFAR-100 with generated captions).
16.9.2 Challenge 2
Reproduce research findings: Implement the DDPM (Denoising Diffusion Probabilistic Models) paper’s training and sampling algorithm.
16.9.3 Challenge 3
Novel application: Design a multimodal model that integrates EEG signals with text descriptions for brain-computer interface applications.
16.10 Discussion Questions
16.10.1 Question 1
How does the cross-attention mechanism in multimodal models relate to multisensory integration in the brain’s association cortex?
16.10.2 Question 2
What are the AI implications of diffusion models for creative applications, and what ethical concerns arise from their ability to generate realistic content?
16.10.3 Question 3
Future research directions: How might incorporating additional modalities (touch, proprioception, smell) lead to more robust and general AI systems?
16.11 Summary
Explored vision-language integration through theory and practice, including:
- Vision Transformer (ViT) architecture and patch-based image processing
- Cross-attention mechanisms for multimodal fusion
- CLIP-style contrastive learning for joint embedding spaces
- Diffusion models and the forward/reverse denoising process
- Classifier-free guidance for controlled generation
Key Takeaways: - Core concepts mastered: ViT, cross-attention, contrastive learning, diffusion - Practical implementation skills developed through hands-on exercises - NeuroAI connections understood: multisensory integration parallels in AI and brain