Files
TinyTorch/tests/milestones/README.md
Vijay Janapa Reddi d05daeb83b Add comprehensive milestone learning verification tests
- Created test suite that verifies actual learning (gradient flow, weight updates, loss convergence)
- Fixed MLP Digits (1986): increased training epochs from 15 to 25
- Added requires_grad=True to Conv2d weights (partial fix)
- Identified gradient flow issues in Conv2d, Embedding, and Attention layers
- Comprehensive documentation of issues and fixes needed
2025-11-22 17:02:10 -05:00

11 KiB

TinyTorch Milestone Learning Verification Tests

Overview

This test suite verifies that actual LEARNING is happening in TinyTorch milestones, not just that code runs without errors. We check:

  1. Loss Convergence: Loss decreases significantly over training
  2. Gradient Flow: All parameters receive non-zero gradients
  3. Weight Updates: Parameters actually change during training
  4. Performance: Models achieve expected accuracy/performance

This is the "trust but verify" approach to ML systems - we don't just hope learning happens, we prove it with rigorous tests.

Test Suite Structure

Main Test File

test_learning_verification.py - Comprehensive learning verification for all milestones

Tests Included

Test Milestone What It Verifies
test_perceptron_learning() 1957 Perceptron Linear classification with gradient descent
test_xor_learning() 1969 XOR Multi-layer network solves non-linear problem
test_mlp_digits_learning() 1986 MLP Real-world digit classification
test_cnn_learning() 1998 CNN Convolutional learning on images
test_transformer_learning() 2017 Transformer Attention-based sequence modeling

Running the Tests

Run All Tests

cd /Users/VJ/GitHub/TinyTorch
python tests/milestones/test_learning_verification.py

Run with pytest

pytest tests/milestones/test_learning_verification.py -v

Run Individual Tests

from tests.milestones.test_learning_verification import test_perceptron_learning
test_perceptron_learning()

What Each Test Checks

1. Gradient Flow Verification

def check_gradient_flow(parameters):
    """
    Verifies gradients are flowing properly:
    - All parameters have gradients
    - Gradients are non-zero
    - Gradients have reasonable magnitude (not exploding/vanishing)
    - No parameters stuck with zero gradients
    """

Why it matters: If gradients don't flow, training is broken. This catches the most common training failures.

2. Weight Update Verification

def check_weight_updates(params_before, params_after):
    """
    Verifies weights actually changed during training:
    - Parameters before vs after training differ
    - Updates have reasonable magnitude
    - No parameters frozen/unchanged
    """

Why it matters: Weights not updating = optimizer not working. Catches broken optimizer step() or zero learning rates.

3. Loss Convergence Verification

def verify_loss_convergence(loss_history, min_decrease=0.1):
    """
    Verifies loss is decreasing (learning is happening):
    - Initial loss > Final loss
    - Decrease is significant (not just noise)
    - Loss generally decreases over time
    """

Why it matters: Loss not decreasing = model not learning. This is the ultimate test of whether learning actually happens.

Test Output

Successful Test

🔬 Training perceptron...
  Epoch  0: Loss = 0.6129
  Epoch 10: Loss = 0.5530
  Epoch 20: Loss = 0.5214

📊 Learning Verification Results:
┌───────────────────────┬──────────┬─────────┐
│ Metric                │ Value    │ Status  │
├───────────────────────┼──────────┼─────────┤
│ Final Accuracy        │ 92.0%    │ ✅ PASS │
│ Loss Decrease         │ 52.3%    │ ✅ PASS │
│ Gradients Flowing     │ 2/2      │ ✅ PASS │
│ Mean Gradient Mag     │ 0.208659 │ ✅ PASS │
│ Weights Updated       │ 2/2      │ ✅ PASS │
│ Mean Weight Change    │ 0.468087 │ ✅ PASS │
└───────────────────────┴──────────┴─────────┘

✅ PERCEPTRON LEARNING VERIFIED
   • Loss decreased significantly
   • Gradients flow properly
   • Weights updated correctly
   • Model converged to high accuracy

Failed Test

🔬 Training CNN on TinyDigits...
  Epoch  0: Loss = 2.3525
  Epoch  3: Loss = 2.2526
  Epoch  6: Loss = 2.2015

📊 Learning Verification Results:
┌───────────────────────┬──────────┬─────────┐
│ Metric                │ Value    │ Status  │
├───────────────────────┼──────────┼─────────┤
│ Final Accuracy        │ 45.0%    │ ❌ FAIL │
│ Loss Decrease         │ 8.3%     │ ❌ FAIL │
│ Gradients Flowing     │ 4/6      │ ❌ FAIL │
│ Conv Gradients        │ 0.000000 │ ❌ FAIL │
│ Weights Updated       │ 4/6      │ ❌ FAIL │
└───────────────────────┴──────────┴─────────┘

❌ CNN LEARNING FAILED
   • Convolutional gradients not flowing
   • Check Conv2d backward() implementation

Understanding the Metrics

Gradient Metrics

  • Gradients Flowing: X/Y means X out of Y parameters received gradients

    • Should be Y/Y (all parameters)
    • If less, some parameters aren't being trained
  • Mean Gradient Magnitude: Average absolute gradient value

    • Should be > 1e-6 (gradients exist and are meaningful)
    • If ~0, gradients vanishing or not flowing
    • If very large (>100), gradients exploding

Weight Metrics

  • Weights Updated: How many parameters actually changed

    • Should equal total parameters
    • If less, optimizer not updating or LR too small
  • Mean Weight Change: Average change in parameter values

    • Should be > 1e-4 (parameters actually moving)
    • If ~0, learning rate too small or optimizer broken

Loss Metrics

  • Loss Decrease: (initial_loss - final_loss) / initial_loss * 100%
    • Should be > 30% for simple tasks
    • Should be > 10% for complex tasks
    • If < 10%, model not learning effectively

Common Failure Modes

Gradients Not Flowing

Symptoms:

  • Gradients Flowing: X/Y where X < Y
  • Some parameters show "Gradients: No"

Causes:

  • Missing .backward() call
  • Incorrect autograd implementation
  • Parameters not connected to loss (dead branches)
  • .data access breaking computation graph

Fix: Check backward() implementation for each layer

Weights Not Updating

Symptoms:

  • Weights Updated: X/Y where X < Y
  • Mean Weight Change: 0.000000

Causes:

  • Optimizer not calling step()
  • Learning rate = 0
  • Parameters don't have requires_grad=True
  • Gradients being cleared before step()

Fix: Check optimizer step() and learning rate

Loss Not Decreasing

Symptoms:

  • Loss Decrease: 5.2% (very small)
  • Loss stays roughly constant

Causes:

  • Learning rate too small
  • Learning rate too large (diverging)
  • Wrong loss function for task
  • Data/label mismatch
  • Architecture too weak for task

Fix: Try different learning rates, check data/labels

Integration with TinyTorch Development

When to Run These Tests

  1. After implementing new modules: Verify learning still works
  2. Before major releases: Ensure all milestones pass
  3. When debugging training: Identify where learning breaks
  4. After autograd changes: Verify gradient flow still works

Adding New Milestone Tests

Template for new tests:

def test_new_milestone_learning():
    """
    Verify [milestone name] learns on [task description].
    
    Expected behavior:
      - Loss should decrease by >X%
      - All Y parameters should receive gradients
      - Final performance should be >Z%
    """
    console.print("\\n" + "="*70)
    console.print(Panel.fit(
        "[bold cyan]TEST N: [Milestone Name][/bold cyan]\\n"
        "[dim][Year] - [Key Paper/Researcher][/dim]",
        border_style="cyan"
    ))
    
    # 1. Create data
    X, y = create_data()
    
    # 2. Build model
    model = build_model()
    params = model.parameters()
    params_before = [Tensor(p.data.copy()) for p in params]
    
    # 3. Train
    loss_fn = SomeLoss()
    optimizer = SomeOptimizer(params, lr=0.01)
    loss_history = []
    
    for epoch in range(epochs):
        predictions = model(X)
        loss = loss_fn(predictions, y)
        loss.backward()
        
        if epoch == 0:
            grad_stats = check_gradient_flow(params)
        
        optimizer.step()
        optimizer.zero_grad()
        loss_history.append(loss.data.item())
    
    # 4. Verify learning
    weight_stats = check_weight_updates(params_before, params)
    convergence_stats = verify_loss_convergence(loss_history, min_decrease=0.3)
    
    # 5. Display results
    # ... create table with metrics ...
    
    # 6. Return pass/fail
    passed = (
        convergence_stats['converged'] and
        grad_stats['params_with_grad'] == grad_stats['total_params'] and
        weight_stats['params_updated'] == weight_stats['total_params']
    )
    
    return passed

Philosophy

Why Test Learning, Not Just Code?

Traditional Unit Tests: "Does the function return the right shape?"
Learning Verification Tests: "Does the model actually learn?"

Example:

  • Unit test: assert output.shape == (batch_size, num_classes)
  • 🔥 Learning test: assert final_accuracy > 90% and loss_decreased > 50%

The "Real Learning" Standard

A milestone passes if:

  1. Loss decreases significantly (not just random fluctuations)
  2. Gradients flow to ALL parameters (no dead weights)
  3. Weights actually update (optimizer working)
  4. Final performance meets expectations (model converges)

If any of these fail, learning is broken - even if the code "works".

Results Summary

Current status of TinyTorch milestones:

Milestone Status Notes
1957 Perceptron PASS Learns linear classification perfectly
1969 XOR PASS Solves XOR with multi-layer network
1986 MLP Digits ⚠️ VARIABLE Sometimes passes (depends on init)
1998 CNN ⚠️ NEEDS WORK Gradient flow issues in Conv2d
2017 Transformer ⚠️ NEEDS WORK Attention/embedding gradient flow

Next Steps

For failing tests:

  1. CNN: Debug Conv2d backward() - gradients not flowing properly
  2. Transformer: Debug attention backward() - only 4/19 params get gradients
  3. MLP Digits: Improve initialization or increase training epochs

Files

  • test_learning_verification.py - Main test suite
  • README.md - This file
  • INTERMODULE_TEST_COVERAGE.md - Related integration tests
  • /tests/integration/INTERMODULE_TEST_COVERAGE.md - Integration tests
  • /milestones/*/GRADIENT_FLOW_VERIFICATION.md - Milestone-specific docs
  • /docs/development/REAL_DATA_REAL_SYSTEMS.md - Development philosophy

Remember: Code that runs is not the same as code that learns. These tests verify the latter.