Upload visualize_micro_world.py with huggingface_hub

visualize_micro_world.py

@@ -0,0 +1,586 @@
+"""
+Micro-World Visualization: Understanding Residual Connections
+
+This script creates intuitive visualizations explaining:
+1. Signal flow through layers (forward pass)
+2. Gradient flow through layers (backward pass)
+3. The "gradient highway" effect of residual connections
+4. Layer-by-layer transformation visualization
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
+import json
+import os
+
+# Set seeds
+torch.manual_seed(42)
+np.random.seed(42)
+
+# Load results from experiment
+with open('results_fair.json', 'r') as f:
+    results = json.load(f)
+
+os.makedirs('plots_micro', exist_ok=True)
+
+# ============================================================
+# VISUALIZATION 1: Signal Flow Diagram (Forward Pass)
+# ============================================================
+def plot_signal_flow():
+    """Visualize how signal magnitude changes through layers"""
+    
+    fig, axes = plt.subplots(1, 2, figsize=(14, 8))
+    
+    plain_stds = results['plain_mlp']['activation_stds']
+    res_stds = results['res_mlp']['activation_stds']
+    
+    # Normalize for visualization (input signal = 1.0)
+    input_std = 0.577  # std of U(-1,1)
+    plain_signal = [input_std] + plain_stds
+    res_signal = [input_std] + res_stds
+    
+    layers = range(len(plain_signal))
+    
+    # Left plot: PlainMLP signal decay
+    ax = axes[0]
+    ax.set_title('PlainMLP: Signal DIES\n(No Residual Connection)', fontsize=14, fontweight='bold', color='#c0392b')
+    
+    # Draw signal as decreasing bars
+    colors_plain = plt.cm.Reds(np.linspace(0.3, 0.9, len(plain_signal)))
+    bars = ax.bar(layers, plain_signal, color=colors_plain, edgecolor='darkred', linewidth=1.5)
+    
+    ax.set_xlabel('Layer (0=Input, 1-20=Hidden)', fontsize=12)
+    ax.set_ylabel('Signal Strength (Activation Std)', fontsize=12)
+    ax.set_ylim(0, 0.7)
+    
+    # Add annotation
+    ax.annotate('Signal\ncollapses!', xy=(15, 0.02), fontsize=12, color='darkred',
+                ha='center', fontweight='bold')
+    ax.axhline(y=0.1, color='gray', linestyle='--', alpha=0.5, label='Healthy threshold')
+    
+    # Right plot: ResMLP signal preservation
+    ax = axes[1]
+    ax.set_title('ResMLP: Signal PRESERVED\n(With Residual Connection)', fontsize=14, fontweight='bold', color='#2980b9')
+    
+    colors_res = plt.cm.Blues(np.linspace(0.3, 0.9, len(res_signal)))
+    bars = ax.bar(layers, res_signal, color=colors_res, edgecolor='darkblue', linewidth=1.5)
+    
+    ax.set_xlabel('Layer (0=Input, 1-20=Hidden)', fontsize=12)
+    ax.set_ylabel('Signal Strength (Activation Std)', fontsize=12)
+    ax.set_ylim(0, 0.7)
+    
+    # Add annotation
+    ax.annotate('Signal stays\nhealthy!', xy=(15, 0.25), fontsize=12, color='darkblue',
+                ha='center', fontweight='bold')
+    ax.axhline(y=0.1, color='gray', linestyle='--', alpha=0.5, label='Healthy threshold')
+    
+    plt.tight_layout()
+    plt.savefig('plots_micro/1_signal_flow.png', dpi=150, bbox_inches='tight')
+    plt.close()
+    print("[Plot 1] Signal flow visualization saved")
+
+
+# ============================================================
+# VISUALIZATION 2: Gradient Flow Diagram (Backward Pass)
+# ============================================================
+def plot_gradient_flow():
+    """Visualize gradient magnitude through layers"""
+    
+    fig, axes = plt.subplots(1, 2, figsize=(14, 8))
+    
+    plain_grads = results['plain_mlp']['gradient_norms']
+    res_grads = results['res_mlp']['gradient_norms']
+    
+    layers = range(1, 21)
+    
+    # Left: PlainMLP gradient vanishing
+    ax = axes[0]
+    ax.set_title('PlainMLP: Gradients VANISH\n(Backward Pass)', fontsize=14, fontweight='bold', color='#c0392b')
+    
+    # Use log scale bar chart
+    colors = plt.cm.Reds(np.linspace(0.9, 0.3, 20))
+    ax.bar(layers, plain_grads, color=colors, edgecolor='darkred', linewidth=1)
+    ax.set_yscale('log')
+    ax.set_xlabel('Layer (1=First, 20=Last)', fontsize=12)
+    ax.set_ylabel('Gradient Magnitude (log scale)', fontsize=12)
+    ax.set_ylim(1e-20, 1e-1)
+    
+    # Annotations
+    ax.annotate(f'Layer 20:\n{plain_grads[-1]:.1e}', xy=(20, plain_grads[-1]), 
+                xytext=(17, 1e-4), fontsize=10, color='darkred',
+                arrowprops=dict(arrowstyle='->', color='darkred'))
+    ax.annotate(f'Layer 1:\n{plain_grads[0]:.1e}\n(DEAD!)', xy=(1, max(plain_grads[0], 1e-20)), 
+                xytext=(4, 1e-15), fontsize=10, color='darkred', fontweight='bold',
+                arrowprops=dict(arrowstyle='->', color='darkred'))
+    
+    # Right: ResMLP healthy gradients
+    ax = axes[1]
+    ax.set_title('ResMLP: Gradients FLOW\n(Backward Pass)', fontsize=14, fontweight='bold', color='#2980b9')
+    
+    colors = plt.cm.Blues(np.linspace(0.9, 0.3, 20))
+    ax.bar(layers, res_grads, color=colors, edgecolor='darkblue', linewidth=1)
+    ax.set_yscale('log')
+    ax.set_xlabel('Layer (1=First, 20=Last)', fontsize=12)
+    ax.set_ylabel('Gradient Magnitude (log scale)', fontsize=12)
+    ax.set_ylim(1e-20, 1e-1)
+    
+    # Annotations
+    ax.annotate(f'Layer 20:\n{res_grads[-1]:.1e}', xy=(20, res_grads[-1]), 
+                xytext=(17, 1e-4), fontsize=10, color='darkblue',
+                arrowprops=dict(arrowstyle='->', color='darkblue'))
+    ax.annotate(f'Layer 1:\n{res_grads[0]:.1e}\n(Healthy!)', xy=(1, res_grads[0]), 
+                xytext=(4, 1e-4), fontsize=10, color='darkblue', fontweight='bold',
+                arrowprops=dict(arrowstyle='->', color='darkblue'))
+    
+    plt.tight_layout()
+    plt.savefig('plots_micro/2_gradient_flow.png', dpi=150, bbox_inches='tight')
+    plt.close()
+    print("[Plot 2] Gradient flow visualization saved")
+
+
+# ============================================================
+# VISUALIZATION 3: The Residual "Highway" Concept
+# ============================================================
+def plot_highway_concept():
+    """Visual diagram showing the gradient highway concept"""
+    
+    fig, axes = plt.subplots(2, 1, figsize=(14, 10))
+    
+    # Top: PlainMLP - no highway
+    ax = axes[0]
+    ax.set_xlim(0, 12)
+    ax.set_ylim(0, 3)
+    ax.set_aspect('equal')
+    ax.axis('off')
+    ax.set_title('PlainMLP: Gradient Must Pass Through EVERY Layer\n(Like a winding mountain road)', 
+                 fontsize=14, fontweight='bold', color='#c0392b', pad=20)
+    
+    # Draw layers as boxes
+    for i in range(6):
+        x = 1 + i * 1.8
+        box = FancyBboxPatch((x, 1), 1.2, 1, boxstyle="round,pad=0.05", 
+                             facecolor='#e74c3c', edgecolor='darkred', linewidth=2)
+        ax.add_patch(box)
+        ax.text(x + 0.6, 1.5, f'L{i+1}', ha='center', va='center', fontsize=11, 
+                color='white', fontweight='bold')
+        
+        # Draw arrows between layers (getting thinner = gradient vanishing)
+        if i < 5:
+            thickness = 3 * (0.5 ** i)  # Exponential decay
+            alpha = max(0.2, 1 - i * 0.18)
+            ax.annotate('', xy=(x + 1.8, 1.5), xytext=(x + 1.2, 1.5),
+                       arrowprops=dict(arrowstyle='->', color='darkred', 
+                                      lw=thickness, alpha=alpha))
+    
+    # Add gradient flow label
+    ax.text(0.3, 1.5, 'Gradient\n→', fontsize=10, ha='center', va='center', color='darkred')
+    ax.text(11.5, 1.5, '→ Loss', fontsize=10, ha='center', va='center', color='darkred')
+    
+    # Add "vanishing" annotation
+    ax.annotate('Gradient shrinks\nat each layer!', xy=(8, 0.5), fontsize=11, 
+                color='darkred', style='italic')
+    
+    # Bottom: ResMLP - with highway
+    ax = axes[1]
+    ax.set_xlim(0, 12)
+    ax.set_ylim(0, 3.5)
+    ax.set_aspect('equal')
+    ax.axis('off')
+    ax.set_title('ResMLP: Gradient Has a Direct HIGHWAY\n(Skip connections = express lane)', 
+                 fontsize=14, fontweight='bold', color='#2980b9', pad=20)
+    
+    # Draw the highway (skip connection) at top
+    ax.plot([1, 11], [2.8, 2.8], color='#27ae60', linewidth=6, alpha=0.8)
+    ax.annotate('', xy=(11, 2.8), xytext=(10.5, 2.8),
+               arrowprops=dict(arrowstyle='->', color='#27ae60', lw=3))
+    ax.text(6, 3.2, '✓ GRADIENT HIGHWAY (Identity Path)', ha='center', fontsize=12, 
+            color='#27ae60', fontweight='bold')
+    
+    # Draw layers as boxes
+    for i in range(6):
+        x = 1 + i * 1.8
+        box = FancyBboxPatch((x, 1), 1.2, 1, boxstyle="round,pad=0.05", 
+                             facecolor='#3498db', edgecolor='darkblue', linewidth=2)
+        ax.add_patch(box)
+        ax.text(x + 0.6, 1.5, f'L{i+1}', ha='center', va='center', fontsize=11, 
+                color='white', fontweight='bold')
+        
+        # Draw arrows between layers (constant thickness = gradient preserved)
+        if i < 5:
+            ax.annotate('', xy=(x + 1.8, 1.5), xytext=(x + 1.2, 1.5),
+                       arrowprops=dict(arrowstyle='->', color='darkblue', lw=2))
+        
+        # Draw skip connections going up to highway
+        ax.plot([x + 0.6, x + 0.6], [2, 2.8], color='#27ae60', linewidth=2, alpha=0.5)
+    
+    ax.text(0.3, 1.5, 'Gradient\n→', fontsize=10, ha='center', va='center', color='darkblue')
+    ax.text(11.5, 1.5, '→ Loss', fontsize=10, ha='center', va='center', color='darkblue')
+    
+    # Add explanation
+    ax.annotate('Gradient flows on highway\neven if layers block it!', xy=(8, 0.3), 
+                fontsize=11, color='#27ae60', style='italic')
+    
+    plt.tight_layout()
+    plt.savefig('plots_micro/3_highway_concept.png', dpi=150, bbox_inches='tight')
+    plt.close()
+    print("[Plot 3] Highway concept visualization saved")
+
+
+# ============================================================
+# VISUALIZATION 4: Mathematical View - Chain Rule
+# ============================================================
+def plot_chain_rule():
+    """Visualize the chain rule multiplication effect"""
+    
+    fig, axes = plt.subplots(1, 2, figsize=(14, 7))
+    
+    # Simulate gradient flow
+    num_layers = 20
+    
+    # PlainMLP: gradient = product of layer gradients (each < 1)
+    plain_layer_grad = 0.7  # Each layer shrinks gradient by 0.7x
+    plain_cumulative = [1.0]
+    for i in range(num_layers):
+        plain_cumulative.append(plain_cumulative[-1] * plain_layer_grad)
+    
+    # ResMLP: gradient = 1 + small_contribution (always >= 1 path)
+    res_layer_contrib = 0.05  # Small contribution from each layer
+    res_cumulative = [1.0]
+    for i in range(num_layers):
+        # The "1" from identity ensures gradient doesn't vanish
+        res_cumulative.append(res_cumulative[-1] * (1.0 + res_layer_contrib * (0.9 ** i)))
+    
+    layers = range(num_layers + 1)
+    
+    # Left: Show the multiplication effect
+    ax = axes[0]
+    ax.semilogy(layers, plain_cumulative, 'o-', color='#e74c3c', linewidth=2, 
+                markersize=8, label='PlainMLP: 0.7 × 0.7 × 0.7 × ...')
+    ax.semilogy(layers, res_cumulative, 's-', color='#3498db', linewidth=2, 
+                markersize=8, label='ResMLP: (1+ε) × (1+ε) × ...')
+    
+    ax.set_xlabel('Layers Traversed (backward from loss)', fontsize=12)
+    ax.set_ylabel('Cumulative Gradient Scale (log)', fontsize=12)
+    ax.set_title('Chain Rule: Why Gradients Vanish\n(Multiplication Effect)', fontsize=14, fontweight='bold')
+    ax.legend(fontsize=11)
+    ax.grid(True, alpha=0.3)
+    ax.set_ylim(1e-8, 10)
+    
+    # Add annotations
+    ax.annotate(f'After 20 layers:\n{plain_cumulative[-1]:.1e}', 
+                xy=(20, plain_cumulative[-1]), xytext=(15, 1e-6),
+                fontsize=10, color='#c0392b',
+                arrowprops=dict(arrowstyle='->', color='#c0392b'))
+    ax.annotate(f'After 20 layers:\n{res_cumulative[-1]:.2f}', 
+                xy=(20, res_cumulative[-1]), xytext=(15, 3),
+                fontsize=10, color='#2980b9',
+                arrowprops=dict(arrowstyle='->', color='#2980b9'))
+    
+    # Right: Show the formula
+    ax = axes[1]
+    ax.axis('off')
+    ax.set_xlim(0, 10)
+    ax.set_ylim(0, 10)
+    
+    ax.text(5, 9, 'The Math Behind It', fontsize=16, fontweight='bold', 
+            ha='center', va='center')
+    
+    # PlainMLP formula
+    ax.text(5, 7.5, 'PlainMLP Gradient:', fontsize=13, fontweight='bold', 
+            ha='center', color='#c0392b')
+    ax.text(5, 6.5, r'$\frac{\partial L}{\partial x_1} = \frac{\partial L}{\partial x_{20}} \times \prod_{i=1}^{20} \frac{\partial x_{i+1}}{\partial x_i}$', 
+            fontsize=14, ha='center', color='#c0392b')
+    ax.text(5, 5.5, '= (small) × (small) × ... × (small) = TINY!', 
+            fontsize=11, ha='center', color='#c0392b', style='italic')
+    
+    # ResMLP formula
+    ax.text(5, 4, 'ResMLP Gradient:', fontsize=13, fontweight='bold', 
+            ha='center', color='#2980b9')
+    ax.text(5, 3, r'$\frac{\partial L}{\partial x_1} = \frac{\partial L}{\partial x_{20}} \times \prod_{i=1}^{20} (1 + \frac{\partial f_i}{\partial x_i})$', 
+            fontsize=14, ha='center', color='#2980b9')
+    ax.text(5, 2, '= (1+ε) × (1+ε) × ... = PRESERVED!', 
+            fontsize=11, ha='center', color='#2980b9', style='italic')
+    
+    # Key insight
+    box = FancyBboxPatch((1, 0.3), 8, 1.2, boxstyle="round,pad=0.1", 
+                         facecolor='#f9e79f', edgecolor='#f39c12', linewidth=2)
+    ax.add_patch(box)
+    ax.text(5, 0.9, '💡 Key Insight: The "+x" in residual adds a "1" to each gradient term,\n'
+                    'preventing the product from shrinking to zero!', 
+            fontsize=11, ha='center', va='center', fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig('plots_micro/4_chain_rule.png', dpi=150, bbox_inches='tight')
+    plt.close()
+    print("[Plot 4] Chain rule visualization saved")
+
+
+# ============================================================
+# VISUALIZATION 5: Layer-by-Layer Transformation
+# ============================================================
+def plot_layer_transformation():
+    """Show what happens to a single input vector through layers"""
+    
+    # Create simple models for visualization
+    class PlainMLP(nn.Module):
+        def __init__(self, dim, num_layers):
+            super().__init__()
+            self.layers = nn.ModuleList()
+            for _ in range(num_layers):
+                layer = nn.Linear(dim, dim)
+                nn.init.kaiming_normal_(layer.weight)
+                layer.weight.data *= 1.0 / np.sqrt(num_layers)
+                nn.init.zeros_(layer.bias)
+                self.layers.append(layer)
+            self.activation = nn.ReLU()
+        
+        def forward_with_intermediates(self, x):
+            intermediates = [x.clone()]
+            for layer in self.layers:
+                x = self.activation(layer(x))
+                intermediates.append(x.clone())
+            return intermediates
+    
+    class ResMLP(nn.Module):
+        def __init__(self, dim, num_layers):
+            super().__init__()
+            self.layers = nn.ModuleList()
+            for _ in range(num_layers):
+                layer = nn.Linear(dim, dim)
+                nn.init.kaiming_normal_(layer.weight)
+                layer.weight.data *= 1.0 / np.sqrt(num_layers)
+                nn.init.zeros_(layer.bias)
+                self.layers.append(layer)
+            self.activation = nn.ReLU()
+        
+        def forward_with_intermediates(self, x):
+            intermediates = [x.clone()]
+            for layer in self.layers:
+                x = x + self.activation(layer(x))
+                intermediates.append(x.clone())
+            return intermediates
+    
+    # Create models
+    dim = 64
+    num_layers = 20
+    plain = PlainMLP(dim, num_layers)
+    res = ResMLP(dim, num_layers)
+    
+    # Single input vector
+    x = torch.randn(1, dim) * 0.5
+    
+    # Get intermediates
+    plain_ints = plain.forward_with_intermediates(x)
+    res_ints = res.forward_with_intermediates(x)
+    
+    # Extract norms and first 2 dimensions for visualization
+    plain_norms = [p.norm().item() for p in plain_ints]
+    res_norms = [r.norm().item() for r in res_ints]
+    
+    plain_2d = [p[0, :2].detach().numpy() for p in plain_ints]
+    res_2d = [r[0, :2].detach().numpy() for r in res_ints]
+    
+    fig, axes = plt.subplots(2, 2, figsize=(14, 12))
+    
+    # Top left: Vector magnitude through layers
+    ax = axes[0, 0]
+    layers = range(len(plain_norms))
+    ax.plot(layers, plain_norms, 'o-', color='#e74c3c', linewidth=2, markersize=6, label='PlainMLP')
+    ax.plot(layers, res_norms, 's-', color='#3498db', linewidth=2, markersize=6, label='ResMLP')
+    ax.set_xlabel('Layer (0=Input)', fontsize=12)
+    ax.set_ylabel('Vector Magnitude (L2 norm)', fontsize=12)
+    ax.set_title('Signal Magnitude Through Network', fontsize=13, fontweight='bold')
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+    
+    # Top right: 2D trajectory visualization
+    ax = axes[0, 1]
+    
+    # PlainMLP trajectory
+    plain_x = [p[0] for p in plain_2d]
+    plain_y = [p[1] for p in plain_2d]
+    ax.plot(plain_x, plain_y, 'o-', color='#e74c3c', linewidth=1.5, markersize=4, 
+            alpha=0.7, label='PlainMLP path')
+    ax.scatter(plain_x[0], plain_y[0], s=100, color='#e74c3c', marker='*', zorder=5)
+    ax.scatter(plain_x[-1], plain_y[-1], s=100, color='#e74c3c', marker='X', zorder=5)
+    
+    # ResMLP trajectory
+    res_x = [r[0] for r in res_2d]
+    res_y = [r[1] for r in res_2d]
+    ax.plot(res_x, res_y, 's-', color='#3498db', linewidth=1.5, markersize=4, 
+            alpha=0.7, label='ResMLP path')
+    ax.scatter(res_x[0], res_y[0], s=100, color='#3498db', marker='*', zorder=5)
+    ax.scatter(res_x[-1], res_y[-1], s=100, color='#3498db', marker='X', zorder=5)
+    
+    ax.set_xlabel('Dimension 1', fontsize=12)
+    ax.set_ylabel('Dimension 2', fontsize=12)
+    ax.set_title('2D Projection of Vector Path\n(★=start, ✕=end)', fontsize=13, fontweight='bold')
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+    ax.axhline(y=0, color='gray', linestyle='-', alpha=0.3)
+    ax.axvline(x=0, color='gray', linestyle='-', alpha=0.3)
+    
+    # Bottom left: PlainMLP heatmap of activations
+    ax = axes[1, 0]
+    plain_acts = np.array([p[0, :32].detach().numpy() for p in plain_ints])  # First 32 dims
+    im = ax.imshow(plain_acts.T, aspect='auto', cmap='Reds', interpolation='nearest')
+    ax.set_xlabel('Layer', fontsize=12)
+    ax.set_ylabel('Dimension (first 32)', fontsize=12)
+    ax.set_title('PlainMLP: Activations Die Out', fontsize=13, fontweight='bold', color='#c0392b')
+    plt.colorbar(im, ax=ax, label='Activation Value')
+    
+    # Bottom right: ResMLP heatmap of activations
+    ax = axes[1, 1]
+    res_acts = np.array([r[0, :32].detach().numpy() for r in res_ints])  # First 32 dims
+    im = ax.imshow(res_acts.T, aspect='auto', cmap='Blues', interpolation='nearest')
+    ax.set_xlabel('Layer', fontsize=12)
+    ax.set_ylabel('Dimension (first 32)', fontsize=12)
+    ax.set_title('ResMLP: Activations Stay Alive', fontsize=13, fontweight='bold', color='#2980b9')
+    plt.colorbar(im, ax=ax, label='Activation Value')
+    
+    plt.tight_layout()
+    plt.savefig('plots_micro/5_layer_transformation.png', dpi=150, bbox_inches='tight')
+    plt.close()
+    print("[Plot 5] Layer transformation visualization saved")
+
+
+# ============================================================
+# VISUALIZATION 6: Before/After Training Comparison
+# ============================================================
+def plot_learning_comparison():
+    """Show what each model learned (or didn't learn)"""
+    
+    fig, axes = plt.subplots(2, 2, figsize=(14, 12))
+    
+    plain_losses = results['plain_mlp']['loss_history']
+    res_losses = results['res_mlp']['loss_history']
+    
+    # Top left: Loss curves with annotations
+    ax = axes[0, 0]
+    steps = range(len(plain_losses))
+    ax.plot(steps, plain_losses, color='#e74c3c', linewidth=2, label='PlainMLP')
+    ax.plot(steps, res_losses, color='#3498db', linewidth=2, label='ResMLP')
+    ax.set_xlabel('Training Steps', fontsize=12)
+    ax.set_ylabel('MSE Loss', fontsize=12)
+    ax.set_title('Learning Progress', fontsize=13, fontweight='bold')
+    ax.set_yscale('log')
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+    
+    # Add phase annotations
+    ax.axvspan(0, 50, alpha=0.1, color='gray')
+    ax.text(25, 5, 'Early\nTraining', ha='center', fontsize=9, color='gray')
+    ax.axvspan(450, 500, alpha=0.1, color='green')
+    ax.text(475, 5, 'Final', ha='center', fontsize=9, color='gray')
+    
+    # Top right: Loss reduction bar chart
+    ax = axes[0, 1]
+    
+    plain_initial = plain_losses[0]
+    plain_final = plain_losses[-1]
+    res_initial = res_losses[0]
+    res_final = res_losses[-1]
+    
+    plain_reduction = (1 - plain_final / plain_initial) * 100
+    res_reduction = (1 - res_final / res_initial) * 100
+    
+    bars = ax.bar(['PlainMLP', 'ResMLP'], [plain_reduction, res_reduction], 
+                  color=['#e74c3c', '#3498db'], edgecolor='black', linewidth=2)
+    ax.set_ylabel('Loss Reduction (%)', fontsize=12)
+    ax.set_title('How Much Did Each Model Learn?', fontsize=13, fontweight='bold')
+    ax.set_ylim(0, 110)
+    
+    # Add value labels
+    ax.text(0, plain_reduction + 3, f'{plain_reduction:.1f}%', ha='center', fontsize=14, fontweight='bold')
+    ax.text(1, res_reduction + 3, f'{res_reduction:.1f}%', ha='center', fontsize=14, fontweight='bold')
+    
+    # Add verdict
+    ax.text(0, plain_reduction/2, 'FAILED\nTO LEARN', ha='center', va='center', 
+            fontsize=11, color='white', fontweight='bold')
+    ax.text(1, res_reduction/2, 'LEARNED\nSUCCESSFULLY', ha='center', va='center', 
+            fontsize=11, color='white', fontweight='bold')
+    
+    # Bottom: Gradient comparison at different training stages
+    ax = axes[1, 0]
+    
+    plain_grads = results['plain_mlp']['gradient_norms']
+    res_grads = results['res_mlp']['gradient_norms']
+    
+    layers = range(1, 21)
+    width = 0.35
+    
+    ax.bar([l - width/2 for l in layers], plain_grads, width, label='PlainMLP', 
+           color='#e74c3c', alpha=0.8)
+    ax.bar([l + width/2 for l in layers], res_grads, width, label='ResMLP', 
+           color='#3498db', alpha=0.8)
+    
+    ax.set_xlabel('Layer', fontsize=12)
+    ax.set_ylabel('Gradient Magnitude', fontsize=12)
+    ax.set_title('Final Gradient Distribution by Layer', fontsize=13, fontweight='bold')
+    ax.set_yscale('log')
+    ax.legend()
+    ax.grid(True, alpha=0.3, axis='y')
+    
+    # Bottom right: Summary diagram
+    ax = axes[1, 1]
+    ax.axis('off')
+    ax.set_xlim(0, 10)
+    ax.set_ylim(0, 10)
+    
+    ax.text(5, 9.5, '📊 Summary: Why Residuals Work', fontsize=16, fontweight='bold', ha='center')
+    
+    # PlainMLP box
+    box1 = FancyBboxPatch((0.5, 5), 4, 3.5, boxstyle="round,pad=0.1", 
+                          facecolor='#fadbd8', edgecolor='#c0392b', linewidth=2)
+    ax.add_patch(box1)
+    ax.text(2.5, 8, 'PlainMLP ❌', fontsize=13, fontweight='bold', ha='center', color='#c0392b')
+    ax.text(2.5, 7, f'• Loss: {plain_final:.3f}', fontsize=11, ha='center')
+    ax.text(2.5, 6.3, f'• Gradient L1: {plain_grads[0]:.1e}', fontsize=11, ha='center')
+    ax.text(2.5, 5.6, '• Status: UNTRAINABLE', fontsize=11, ha='center', color='#c0392b')
+    
+    # ResMLP box
+    box2 = FancyBboxPatch((5.5, 5), 4, 3.5, boxstyle="round,pad=0.1", 
+                          facecolor='#d4e6f1', edgecolor='#2980b9', linewidth=2)
+    ax.add_patch(box2)
+    ax.text(7.5, 8, 'ResMLP ✓', fontsize=13, fontweight='bold', ha='center', color='#2980b9')
+    ax.text(7.5, 7, f'• Loss: {res_final:.3f}', fontsize=11, ha='center')
+    ax.text(7.5, 6.3, f'• Gradient L1: {res_grads[0]:.1e}', fontsize=11, ha='center')
+    ax.text(7.5, 5.6, '• Status: TRAINED', fontsize=11, ha='center', color='#2980b9')
+    
+    # Key insight box
+    box3 = FancyBboxPatch((1, 0.5), 8, 3.5, boxstyle="round,pad=0.1", 
+                          facecolor='#fef9e7', edgecolor='#f39c12', linewidth=2)
+    ax.add_patch(box3)
+    ax.text(5, 3.5, '💡 The Residual Connection:', fontsize=13, fontweight='bold', ha='center')
+    ax.text(5, 2.6, '1. Creates a "gradient highway" for backpropagation', fontsize=11, ha='center')
+    ax.text(5, 1.9, '2. Preserves signal magnitude through forward pass', fontsize=11, ha='center')
+    ax.text(5, 1.2, '3. Allows training of very deep networks', fontsize=11, ha='center')
+    
+    plt.tight_layout()
+    plt.savefig('plots_micro/6_learning_comparison.png', dpi=150, bbox_inches='tight')
+    plt.close()
+    print("[Plot 6] Learning comparison visualization saved")
+
+
+# ============================================================
+# MAIN
+# ============================================================
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Creating Micro-World Visualizations")
+    print("=" * 60)
+    
+    plot_signal_flow()
+    plot_gradient_flow()
+    plot_highway_concept()
+    plot_chain_rule()
+    plot_layer_transformation()
+    plot_learning_comparison()
+    
+    print("\n" + "=" * 60)
+    print("All visualizations saved to plots_micro/")
+    print("=" * 60)