23 Lab 23: NeuroAI in Action: Four Case Studies
23.1 Learning Objectives
- Understand core concepts from Chapter 23
- Implement key algorithms and techniques
- Apply methods to practical problems
- Analyze results and draw insights
- Connect to broader NeuroAI themes
23.2 Prerequisites
- Reading: Chapter 23: Case Studies
- Libraries: NumPy, Matplotlib, TensorFlow/Keras
- Concepts: Predictive coding, experience replay, transformers, multimodal fusion
23.3 Setup
# Run this cell to import required libraries
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras import layers, Model
import random
np.random.seed(42)
tf.random.set_seed(42)
random.seed(42)
plt.rcParams['figure.figsize'] = (12, 8)23.4 Part 1: Predictive Coding - PredNet Architecture
Implementing the brain’s predictive processing framework using hierarchical prediction networks.
23.4.1 Exercise 1.1: PredNet Block Implementation
Run this cell to implement a PredNet block that mimics the brain’s predictive coding mechanism.
# PredNet: A predictive coding neural network
# Concept: Higher layers predict lower layer activity; only errors propagate upward
class PredNetBlock(layers.Layer):
"""
Implements a single layer of the PredNet architecture.
This layer receives input from below and predictions from above,
computes prediction errors, and generates representations.
Parameters:
-----------
num_channels : int
Number of feature channels in this layer
"""
def __init__(self, num_channels, **kwargs):
super(PredNetBlock, self).__init__(**kwargs)
self.num_channels = num_channels
# Prediction layer: generates top-down predictions
self.conv_pred = layers.Conv2D(
num_channels, (3, 3),
padding='same',
activation='relu',
name='prediction'
)
# Error representation layer: encodes prediction errors
self.conv_error = layers.Conv2D(
num_channels, (3, 3),
padding='same',
activation='relu',
name='error_representation'
)
# Pooling for spatial downsampling
self.pool = layers.MaxPooling2D((2, 2))
def call(self, inputs):
"""
Forward pass computing prediction error.
Inputs:
-------
inputs : tuple
(current_input, higher_representation)
- current_input: input from layer below
- higher_representation: representation from layer above
Returns:
--------
error : tensor
Prediction error (difference between input and prediction)
representation : tensor
Error-based representation for next layer
pooled : tensor
Downsampled representation for higher layers
"""
current_input, higher_representation = inputs
# Generate prediction from higher layer
if higher_representation is not None:
prediction = self.conv_pred(higher_representation)
else:
# At top layer, no prediction (zero baseline)
prediction = tf.zeros_like(current_input)
# Compute prediction error (ReLU ensures positive errors)
error = tf.nn.relu(current_input - prediction)
# Generate representation from error
representation = self.conv_error(error)
return error, representation, self.pool(representation)
# Test the PredNet block
print("PredNetBlock implementation complete!")
print("This block implements hierarchical predictive coding where:")
print(" - Higher layers predict lower layer activity")
print(" - Only prediction errors propagate upward")
print(" - The brain uses similar mechanisms for efficient perception")23.4.2 Exercise 1.2: Experiment with Predictive Coding
Try varying the number of channels and observe how prediction error changes across layers.
# Run this cell to visualize predictive coding behavior
def visualize_prediction_error():
"""Visualize how prediction errors change with hierarchical processing."""
# Create sample input
batch_size = 1
height, width, channels = 32, 32, 16
# Simulated input from lower layer
current_input = tf.random.normal((batch_size, height, width, channels))
# Simulated higher representation (prediction)
higher_repr = tf.random.normal((batch_size, height//2, width//2, channels))
# Upsample prediction to match input size
higher_repr_upsampled = tf.image.resize(higher_repr, (height, width))
prediction = layers.Conv2D(channels, 3, padding='same', activation='relu')(higher_repr_upsampled)
# Compute error
error = tf.nn.relu(current_input - prediction)
# Visualize
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
axes[0].imshow(current_input[0, :, :, 0].numpy(), cmap='viridis')
axes[0].set_title('Input from Below')
axes[0].axis('off')
axes[1].imshow(prediction[0, :, :, 0].numpy(), cmap='viridis')
axes[1].set_title('Top-Down Prediction')
axes[1].axis('off')
axes[2].imshow(error[0, :, :, 0].numpy(), cmap='hot')
axes[2].set_title('Prediction Error (what gets sent up)')
axes[2].axis('off')
plt.tight_layout()
plt.show()
print(f"\nMean prediction error magnitude: {tf.reduce_mean(error).numpy():.4f}")
print("Lower error = better prediction = more efficient coding")
visualize_prediction_error()
23.5 Part 2: Hippocampal Replay - Prioritized Experience Replay
Implementing memory consolidation inspired by the brain’s hippocampal replay mechanism.
23.5.1 Exercise 2.1: Prioritized Replay Buffer
Run this cell to implement a replay buffer that prioritizes surprising experiences.
# Prioritized Experience Replay: Inspired by hippocampal replay
# Concept: Surprising experiences (high TD-error) are replayed more often
class PrioritizedReplayBuffer:
"""
A replay buffer that samples experiences based on their importance.
Similar to how the hippocampus replays surprising memories more frequently
during sleep for consolidation, this buffer prioritizes experiences with
high temporal difference (TD) error.
Parameters:
-----------
capacity : int
Maximum number of experiences to store
alpha : float
How much prioritization to use (0 = uniform, 1 = full prioritization)
"""
def __init__(self, capacity=10000, alpha=0.6):
self.capacity = capacity
self.alpha = alpha # Prioritization exponent
self.buffer = []
self.priorities = []
self.position = 0
def add(self, experience, error):
"""
Add a new experience to the buffer.
Parameters:
-----------
experience : tuple
(state, action, reward, next_state, done)
error : float
TD-error indicating how surprising this experience was
"""
# Priority based on TD-error (add small constant for numerical stability)
priority = (abs(error) + 0.01) ** self.alpha
if len(self.buffer) < self.capacity:
self.buffer.append(experience)
self.priorities.append(priority)
else:
# Overwrite oldest experience
self.buffer[self.position] = experience
self.priorities[self.position] = priority
# Circular buffer
self.position = (self.position + 1) % self.capacity
def sample(self, batch_size):
"""
Sample a batch of experiences based on priority.
Returns:
--------
experiences : list
Batch of sampled experiences
indices : array
Indices of sampled experiences (for updating priorities)
"""
# Convert priorities to probabilities
probs = np.array(self.priorities) / np.sum(self.priorities)
# Sample based on priority
indices = np.random.choice(
len(self.buffer),
batch_size,
p=probs,
replace=False
)
return [self.buffer[i] for i in indices], indices
def update_priorities(self, indices, errors):
"""Update priorities after learning (errors may have changed)."""
for idx, error in zip(indices, errors):
self.priorities[idx] = (abs(error) + 0.01) ** self.alpha
def __len__(self):
return len(self.buffer)
# Test the buffer
print("PrioritizedReplayBuffer implementation complete!")
print("\nThis buffer mimics hippocampal replay by:")
print(" - Storing experiences with associated TD-errors")
print(" - Sampling surprising experiences more frequently")
print(" - Similar to how the brain prioritizes important memories during sleep")23.5.2 Exercise 2.2: Compare Replay Strategies
Compare uniform vs prioritized replay on a simple learning task.
# Run this cell to compare replay strategies
def compare_replay_strategies():
"""Compare uniform vs prioritized experience replay."""
# Simulate experiences with varying TD-errors
num_experiences = 1000
experiences = [
(f"state_{i}", f"action_{i}", np.random.randn(), f"next_state_{i}", False)
for i in range(num_experiences)
]
# Simulated TD-errors (most near zero, some very high - like real RL)
td_errors = np.abs(np.random.exponential(0.5, num_experiences))
# Create prioritized buffer
per_buffer = PrioritizedReplayBuffer(capacity=num_experiences, alpha=0.6)
# Add experiences
for exp, error in zip(experiences, td_errors):
per_buffer.add(exp, error)
# Sample and analyze
batch_size = 50
sampled, _ = per_buffer.sample(batch_size)
# Compare with uniform sampling
uniform_indices = np.random.choice(num_experiences, batch_size, replace=False)
# Visualize
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Distribution of TD-errors
axes[0].hist(td_errors, bins=50, alpha=0.7, label='All experiences', color='gray')
axes[0].hist(
td_errors[np.random.choice(num_experiences, batch_size)],
bins=20, alpha=0.7, label='Uniform sampling', color='blue'
)
axes[0].set_xlabel('TD-Error (Surprise)')
axes[0].set_ylabel('Count')
axes[0].set_title('Uniform Sampling: Random Selection')
axes[0].legend()
# Prioritized sampling distribution
per_sampled_errors = [td_errors[i] for i in
np.random.choice(num_experiences, batch_size,
p=td_errors**0.6/np.sum(td_errors**0.6))]
axes[1].hist(td_errors, bins=50, alpha=0.7, label='All experiences', color='gray')
axes[1].hist(per_sampled_errors, bins=20, alpha=0.7, label='Prioritized sampling', color='red')
axes[1].set_xlabel('TD-Error (Surprise)')
axes[1].set_ylabel('Count')
axes[1].set_title('Prioritized Sampling: Focus on Surprising Events')
axes[1].legend()
plt.tight_layout()
plt.show()
print(f"\nMean TD-error of uniform samples: {np.mean([td_errors[i] for i in uniform_indices]):.4f}")
print(f"Mean TD-error of prioritized samples: {np.mean(per_sampled_errors):.4f}")
print("\nPrioritized replay focuses learning on informative experiences!")
compare_replay_strategies()23.6 Part 3: Vision Transformers - Global Attention Mechanisms
Implementing attention mechanisms inspired by long-range cortical connections.
23.6.1 Exercise 3.1: Transformer Block Implementation
Run this cell to implement a Vision Transformer block with self-attention.
# Vision Transformer: Inspired by brain's long-range connections
# Concept: Self-attention allows any patch to attend to any other patch
class TransformerBlock(layers.Layer):
"""
A transformer block implementing self-attention and feed-forward layers.
This enables global context integration, similar to how the brain uses
long-range horizontal and feedback connections to integrate information
across the entire visual field.
Parameters:
-----------
num_heads : int
Number of attention heads (parallel attention mechanisms)
projection_dim : int
Dimension of the attention projection space
mlp_dim : int
Dimension of the feed-forward network
dropout : float
Dropout rate for regularization
"""
def __init__(self, num_heads, projection_dim, mlp_dim, dropout=0.1):
super(TransformerBlock, self).__init__()
# Multi-head self-attention (global integration)
self.attention = layers.MultiHeadAttention(
num_heads=num_heads,
key_dim=projection_dim
)
# Feed-forward network (local processing)
self.mlp = tf.keras.Sequential([
layers.Dense(mlp_dim, activation="gelu"),
layers.Dropout(dropout),
layers.Dense(projection_dim),
layers.Dropout(dropout),
])
# Layer normalization (stabilizes training)
self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
def call(self, inputs):
"""
Forward pass through transformer block.
The block:
1. Computes self-attention (every patch attends to every other)
2. Adds residual connection
3. Applies feed-forward network
4. Adds residual connection
This mimics the brain's ability to integrate context globally.
"""
# Self-attention with residual connection
x1 = self.layernorm1(inputs)
attention_output = self.attention(x1, x1)
x2 = layers.add([attention_output, inputs])
# Feed-forward with residual connection
x3 = self.layernorm2(x2)
mlp_output = self.mlp(x3)
return layers.add([mlp_output, x2])
# Test the transformer block
print("TransformerBlock implementation complete!")
print("\nThis block implements self-attention which:")
print(" - Allows any image patch to attend to any other patch")
print(" - Enables global context from the first layer")
print(" - Similar to long-range connections in visual cortex")23.6.2 Exercise 3.2: Visualize Attention Patterns
Explore how attention mechanisms learn relationships between image regions.
# Run this cell to visualize attention patterns
def visualize_attention_pattern():
"""Visualize how self-attention connects different image regions."""
# Simulate patch embeddings for a 4x4 grid of patches
num_patches = 16
patch_dim = 64
patches = tf.random.normal((1, num_patches, patch_dim))
# Create attention layer
attention_layer = layers.MultiHeadAttention(
num_heads=4,
key_dim=16
)
# Compute attention
attention_output, attention_weights = attention_layer(
patches, patches, return_attention_scores=True
)
# Visualize attention for one head
fig, axes = plt.subplots(2, 2, figsize=(12, 12))
for head_idx in range(4):
ax = axes[head_idx // 2, head_idx % 2]
# Get attention map for this head (average across query positions)
attention_map = attention_weights[0, head_idx].numpy()
avg_attention = np.mean(attention_map, axis=0).reshape(4, 4)
im = ax.imshow(avg_attention, cmap='viridis')
ax.set_title(f'Attention Head {head_idx + 1}')
ax.set_xlabel('Patch Position (4x4 grid)')
ax.set_ylabel('Patch Position')
# Add grid lines
for i in range(5):
ax.axhline(i - 0.5, color='white', linewidth=0.5)
ax.axvline(i - 0.5, color='white', linewidth=0.5)
plt.suptitle('Attention Patterns: Which Patches Attend to Which?\n'
'(Brighter = more attention)', fontsize=14)
plt.tight_layout()
plt.show()
print("\nKey insight: Different attention heads learn different relationships!")
print(" - Some heads focus on local neighborhoods (like CNNs)")
print(" - Other heads connect distant patches (global context)")
print(" - The brain uses similar context-dependent integration")
visualize_attention_pattern()
23.7 Part 4: Multimodal Fusion - Integrating Multiple Data Sources
Implementing multisensory integration inspired by the brain’s fusion of different modalities.
23.7.1 Exercise 4.1: Multimodal Fusion Model
Run this cell to implement a multimodal fusion model for Alzheimer’s prediction.
# Multimodal Fusion: Inspired by brain's multisensory integration
# Concept: Combine different data types (MRI, clinical, genetic) for better predictions
def create_multimodal_predictor(image_shape, tabular_dim):
"""
Create a multimodal fusion model for medical prediction.
This model fuses information from multiple sources, similar to how
the brain integrates information from vision, audition, and other senses
in regions like the superior temporal sulcus.
Parameters:
-----------
image_shape : tuple
Shape of input images (e.g., MRI scans)
tabular_dim : int
Number of clinical/genetic features
Returns:
--------
model : tf.keras.Model
Multimodal fusion model
"""
# === Image Encoder (for MRI scans) ===
# Processes 3D brain imaging data
image_input = layers.Input(shape=image_shape, name="mri_input")
# 3D convolution for volumetric data
x = layers.Conv3D(32, 3, activation='relu', padding='same')(image_input)
x = layers.MaxPooling3D(2)(x)
x = layers.Conv3D(64, 3, activation='relu', padding='same')(x)
x = layers.MaxPooling3D(2)(x)
x = layers.Conv3D(128, 3, activation='relu', padding='same')(x)
# Global features
cnn_out = layers.GlobalMaxPooling3D()(x)
# === Tabular Encoder (for clinical/genetic data) ===
# Processes structured data (age, test scores, genetic markers)
tabular_input = layers.Input(shape=(tabular_dim,), name="tabular_input")
y = layers.Dense(128, activation='relu')(tabular_input)
y = layers.Dropout(0.3)(y)
y = layers.Dense(64, activation='relu')(y)
tabular_out = layers.Dropout(0.3)(y)
# === Fusion Layer ===
# Combine imaging and clinical features
concatenated = layers.concatenate([cnn_out, tabular_out])
# Joint representation
fused = layers.Dense(256, activation='relu')(concatenated)
fused = layers.Dropout(0.4)(fused)
fused = layers.Dense(128, activation='relu')(fused)
# === Output ===
# Binary classification: Alzheimer's risk
output = layers.Dense(1, activation='sigmoid', name="alzheimers_prediction")(fused)
# Create model
model = Model(
inputs=[image_input, tabular_input],
outputs=output,
name="Multimodal_Alzheimer_Predictor"
)
return model
# Create and summarize the model
print("Creating multimodal fusion model...\n")
# Example dimensions (smaller than real for demonstration)
image_shape = (64, 64, 64, 1) # 3D MRI volume
tabular_dim = 20 # Clinical features
model = create_multimodal_predictor(image_shape, tabular_dim)
model.summary()
print("\n" + "="*70)
print("This model implements multimodal fusion:")
print(" - Image encoder: Processes 3D MRI scans")
print(" - Tabular encoder: Processes clinical/genetic data")
print(" - Fusion: Combines both for better predictions")
print(" - Similar to brain's multisensory integration in STS")
print("="*70)23.7.2 Exercise 4.2: Compare Unimodal vs Multimodal Performance
Demonstrate the advantage of multimodal fusion.
# Run this cell to compare unimodal vs multimodal approaches
def compare_modalities():
"""Compare single modality vs multimodal predictions."""
# Simulated performance metrics (based on research findings)
modalities = ['MRI Only', 'Clinical Only', 'Genetic Only', 'Multimodal Fusion']
accuracies = [78, 72, 65, 94] # Accuracy percentages
aucs = [0.82, 0.76, 0.70, 0.96] # AUC scores
# Visualization
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Accuracy comparison
colors = ['#0066cc', '#0066cc', '#0066cc', '#cc0000']
bars1 = axes[0].bar(modalities, accuracies, color=colors, alpha=0.8)
axes[0].set_ylabel('Accuracy (%)')
axes[0].set_title('Prediction Accuracy: Single Modality vs Fusion')
axes[0].set_ylim([0, 100])
axes[0].axhline(y=80, color='gray', linestyle='--', label='Clinical threshold')
# Add value labels
for bar, val in zip(bars1, accuracies):
axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
f'{val}%', ha='center', fontsize=11, fontweight='bold')
# AUC comparison
bars2 = axes[1].bar(modalities, aucs, color=colors, alpha=0.8)
axes[1].set_ylabel('AUC Score')
axes[1].set_title('AUC Score: Single Modality vs Fusion')
axes[1].set_ylim([0, 1.0])
axes[1].axhline(y=0.9, color='gray', linestyle='--', label='Excellent performance')
# Add value labels
for bar, val in zip(bars2, aucs):
axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
f'{val:.2f}', ha='center', fontsize=11, fontweight='bold')
plt.tight_layout()
plt.show()
print("\n" + "="*70)
print("KEY INSIGHT: Multimodal fusion dramatically improves performance!")
print("="*70)
print(f" - MRI alone: {accuracies[0]}% accuracy")
print(f" - Clinical alone: {accuracies[1]}% accuracy")
print(f" - Genetic alone: {accuracies[2]}% accuracy")
print(f" - Multimodal: {accuracies[3]}% accuracy (+16% improvement!)")
print("\nThis mirrors the brain's strategy of integrating multiple senses")
print("to create a more robust and accurate perception of the world.")
print("="*70)
compare_modalities()23.8 Part 5: Analysis and Integration
23.8.1 Exercise 5.1: The NeuroAI Pipeline
Analyze how biological principles translate to engineering breakthroughs.
# Run this cell to analyze the NeuroAI pipeline
def analyze_neuroai_pipeline():
"""Visualize and analyze the four case studies."""
case_studies = {
'Predictive Coding': {
'biological_principle': 'Brain predicts sensory input',
'ai_implementation': 'PredNet',
'breakthrough': 'Robust video prediction',
'key_metric': '35% better prediction'
},
'Hippocampal Replay': {
'biological_principle': 'Memory replay during sleep',
'ai_implementation': 'Prioritized Replay',
'breakthrough': 'Sample-efficient RL',
'key_metric': '2x faster learning'
},
'Visual Attention': {
'biological_principle': 'Long-range cortical connections',
'ai_implementation': 'Vision Transformer',
'breakthrough': 'Global context integration',
'key_metric': 'SOTA on ImageNet'
},
'Multisensory Integration': {
'biological_principle': 'Fusion of multiple senses',
'ai_implementation': 'Multimodal Fusion',
'breakthrough': 'Early Alzheimer\'s detection',
'key_metric': '94% accuracy'
}
}
# Create summary visualization
fig, ax = plt.subplots(figsize=(14, 8))
ax.axis('off')
# Title
ax.text(0.5, 0.95, 'The NeuroAI Pipeline: Four Case Studies',
ha='center', fontsize=16, fontweight='bold')
# Headers
headers = ['Case Study', 'Biological Principle', 'AI Implementation', 'Breakthrough', 'Key Metric']
x_positions = [0.1, 0.28, 0.48, 0.68, 0.88]
for x, header in zip(x_positions, headers):
ax.text(x, 0.85, header, ha='center', fontsize=11, fontweight='bold',
bbox=dict(boxstyle='round', facecolor='#0066cc', alpha=0.3))
# Case studies
y_positions = [0.68, 0.51, 0.34, 0.17]
colors = ['#cc0000', '#0066cc', '#9966cc', '#cc9900']
for (name, data), y, color in zip(case_studies.items(), y_positions, colors):
ax.text(x_positions[0], y, name, ha='center', fontsize=10, fontweight='bold',
color=color)
ax.text(x_positions[1], y, data['biological_principle'], ha='center', fontsize=9)
ax.text(x_positions[2], y, data['ai_implementation'], ha='center', fontsize=9)
ax.text(x_positions[3], y, data['breakthrough'], ha='center', fontsize=9)
ax.text(x_positions[4], y, data['key_metric'], ha='center', fontsize=9,
fontweight='bold', color=color)
# Separator line
if y > 0.2:
ax.axhline(y=y - 0.08, xmin=0.05, xmax=0.95, color='gray', alpha=0.3)
plt.tight_layout()
plt.show()
print("\n" + "="*70)
print("THE NEUROAI PIPELINE: Biology -> Implementation -> Breakthrough")
print("="*70)
print("\nCommon pattern across all case studies:")
print(" 1. Identify computational principle in neuroscience")
print(" 2. Implement it in an AI system")
print(" 3. Achieve measurable engineering improvement")
print("\nThis demonstrates the immense value of looking to the brain")
print("for inspiration in building intelligent systems.")
print("="*70)
analyze_neuroai_pipeline()
23.9 Challenge Problems
23.9.1 Challenge 1: Build a Complete PredNet
Implement a full multi-layer PredNet and train it on video prediction. - Stack multiple PredNetBlock layers - Train on a simple video dataset - Compare prediction error across layers - Visualize what each layer learns
23.9.2 Challenge 2: Train an RL Agent with PER
Implement a full reinforcement learning agent with prioritized replay. - Build a simple environment - Implement DQN with prioritized replay - Compare to uniform replay - Measure sample efficiency
23.9.3 Challenge 3: Vision Transformer from Scratch
Implement a complete ViT for image classification. - Implement patch embedding - Stack transformer blocks - Train on a small dataset (CIFAR-10) - Visualize learned attention patterns
23.9.4 Challenge 4: Clinical Multimodal Model
Build a multimodal model for a real prediction task. - Find a public medical dataset - Implement multimodal fusion - Compare to single-modality baselines - Analyze which modality contributes most
23.10 Discussion Questions
23.10.1 Question 1: Biology to AI Translation
What aspects of the biological systems are lost when translating to AI? - What simplifications are necessary? - Do these simplifications affect performance? - What future work could make translations more faithful?
23.10.2 Question 2: Scale vs. Biological Inspiration
Do we still need biology-inspired algorithms if scale solves everything? - What problems require biological insights? - How should research resources be allocated?
23.10.3 Question 3: Real-World Impact
Which case study is most impactful in real applications? - What barriers prevent wider adoption? - How could better understanding of brain mechanisms accelerate clinical AI? - What ethical considerations arise when using AI for medical diagnosis?
23.11 Summary
Explored four landmark NeuroAI case studies through hands-on implementation:
Key Takeaways: - Predictive Coding: PredNet demonstrates how prediction-based processing leads to robust representations - Memory Replay: Prioritized replay shows how focusing on surprising experiences improves learning efficiency - Visual Attention: Vision Transformers prove that global context integration outperforms local processing - Multimodal Fusion: Medical AI demonstrates that combining multiple data sources dramatically improves predictions
Core Pattern: Each case study follows the NeuroAI pipeline - identifying a biological principle, implementing it in AI, and achieving an engineering breakthrough.