TinyTorch/tests/milestones/FIXES_NEEDED.md

TinyTorch Milestone Fixes - Complete Analysis

Executive Summary

Created comprehensive learning verification tests that check actual learning (not just "code runs"). Found and fixed some issues, identified others that need deeper architectural fixes.

Status Dashboard

Milestone	Status	Issue	Fix Complexity
✅ Perceptron (1957)	PASSING	None	N/A
✅ XOR (1969)	PASSING	None	N/A
✅ MLP Digits (1986)	FIXED	Variable performance	✅ Simple (more epochs)
⚠️ CNN (1998)	BROKEN	No conv gradients	🔴 Complex (autograd integration)
⚠️ Transformer (2017)	BROKEN	No attention/embedding gradients	🔴 Complex (autograd integration)

---

✅ FIXED: MLP Digits (1986)

Problem

• Variable test results: sometimes 75% (pass), sometimes 63.5% (fail)
• Root cause: Random initialization + small dataset (1000 samples)

Solution Applied

Increased training epochs from 15 → 25

# Before:
epochs = 15  # Too few for small dataset

# After:
epochs = 25  # Sufficient for convergence

Results

• ✅ All 3 test runs now pass consistently
• ✅ Achieves 75-87.5% accuracy reliably
• ✅ Loss decreases 30%+
• ✅ All gradients flow correctly

Status: FIXED AND VERIFIED ✅

---

🔴 BROKEN: CNN (1998) - Critical Autograd Issue

Problem

Conv2d doesn't integrate with autograd at all

Symptoms

🔬 Training CNN...
  Loss: 2.46 → 2.00 (barely decreasing)
  Accuracy: 8.5% → 34.5% (random guessing)
  
  ❌ Gradients Flowing: 2/6 (only FC layer, NOT conv layers)
  ❌ Conv Gradients: 0.000000 (completely broken)

Root Cause Analysis

File: tinytorch/core/spatial.py

Issue 1: Missing requires_grad (FIXED BUT INSUFFICIENT)

# Line 87-88: Weights created without gradient tracking
self.weight = Tensor(np.random.normal(...))  # ❌ No requires_grad
self.bias = Tensor(np.zeros(...))            # ❌ No requires_grad

Fix applied:

self.weight = Tensor(np.random.normal(...), requires_grad=True)  # ✅
self.bias = Tensor(np.zeros(...), requires_grad=True)             # ✅

Issue 2: Forward Pass Bypasses Autograd Entirely (FUNDAMENTAL PROBLEM)

Line 188: return Tensor(output)

The entire forward() implementation uses raw numpy operations and .data access:

def forward(self, x):
    # Line 147-151: Uses x.data directly (no gradient tracking)
    padded_input = np.pad(x.data, ...)
    
    # Line 154: Creates raw numpy array
    output = np.zeros((batch_size, ...))
    
    # Line 171-177: All operations on .data (bypasses autograd)
    input_val = padded_input[b, in_ch, ...]
    weight_val = self.weight.data[out_ch, ...]  # ❌ Uses .data!
    conv_sum += input_val * weight_val
    
    # Line 186: Bias also uses .data
    output[:, out_ch, :, :] += self.bias.data[out_ch]
    
    # Line 188: Returns Tensor WITHOUT gradient function attached
    return Tensor(output)  # ❌ No computation graph!

Why This Breaks Learning

1. No Computation Graph: Forward pass doesn't build a graph for backward()
2. .data Access Everywhere: Breaks gradient flow by accessing raw arrays
3. Missing Gradient Function: No Conv2dBackward attached to output Tensor
4. Manual numpy Operations: Autograd can't track manual loops and accumulations

What's Needed to Fix

Option 1: Implement Conv2dBackward (Recommended)

class Conv2dBackward:
    """Gradient function for Conv2d"""
    def __init__(self, x, weight, bias, stride, padding):
        self.x = x
        self.weight = weight
        # ... store context for backward
    
    def backward(self, grad_output):
        # Compute grad_input (deconvolution)
        # Compute grad_weight (correlation)
        # Compute grad_bias (sum over spatial dims)
        return grad_input

def forward(self, x):
    # ... existing convolution code ...
    result = Tensor(output, requires_grad=(x.requires_grad or self.weight.requires_grad))
    if result.requires_grad:
        result._grad_fn = Conv2dBackward(x, self.weight, self.bias, ...)
    return result

Option 2: Rewrite Using Tensor Operations (Cleaner)

def forward(self, x):
    # Use tensor operations that autograd can track:
    # - Use im2col to convert convolution to matrix multiplication
    # - Use Tensor.matmul() instead of raw numpy
    # - Autograd automatically handles gradients
    pass

Option 3: Use PyTorch/JAX backend (Not educational)

Current Status

• ⚠️ requires_grad=True added to weights (partial fix)
• 🔴 Conv2d forward() still bypasses autograd completely
• 🔴 No backward() implementation
• 🔴 CNN milestones don't actually learn from convolutions

Estimated Fix Time: 4-6 hours (implement Conv2dBackward + test thoroughly)

---

🔴 BROKEN: Transformer (2017) - Similar Autograd Issues

Problem

Attention and Embedding layers don't propagate gradients

Symptoms

🔬 Training transformer...
  Loss: 3.43 → 3.22 (minimal decrease)
  
  ❌ Gradients Flowing: 4/19 (only 21% of parameters!)
  ❌ Attention Gradients: No
  ❌ Embedding Gradients: No

Root Cause

Same as Conv2d - These layers likely:

1. Use .data access in forward()
2. Return Tensors without gradient functions
3. Don't integrate with autograd

Files to Check

• tinytorch/text/embeddings.py - Embedding layer
• tinytorch/core/attention.py - MultiHeadAttention layer
• tinytorch/models/transformer.py - LayerNorm, TransformerBlock

What's Likely Broken

# Embedding.forward() probably does:
def forward(self, indices):
    embedded = self.weight.data[indices]  # ❌ Uses .data
    return Tensor(embedded)                # ❌ No grad_fn

# Should do:
def forward(self, indices):
    embedded = self.weight.data[indices]
    result = Tensor(embedded, requires_grad=self.weight.requires_grad)
    if result.requires_grad:
        result._grad_fn = EmbeddingBackward(self.weight, indices)
    return result

Note: There was a fix for embedding gradients mentioned in GRADIENT_FLOW_VERIFICATION.md, but it may not be applied or may be insufficient.

Current Status

• 🔴 Only 4/19 transformer parameters receive gradients
• 🔴 Attention mechanism doesn't backprop
• 🔴 Embeddings don't learn
• 🔴 Transformer milestones don't actually learn from attention

Estimated Fix Time: 3-5 hours (implement EmbeddingBackward + AttentionBackward)

---

The Fundamental Pattern

The Problem

All custom layers that use manual numpy operations have the same issue:

# BROKEN PATTERN (current):
def forward(self, x):
    # Manual numpy operations
    result_data = np.some_operation(x.data)  # ❌ Uses .data
    return Tensor(result_data)                # ❌ No grad tracking

# Gradient never flows backward!

The Solution

Two options:

Option A: Attach Gradient Functions (More control, educational)

def forward(self, x):
    result_data = np.some_operation(x.data)
    result = Tensor(result_data, requires_grad=True)
    if x.requires_grad or self.param.requires_grad:
        result._grad_fn = CustomBackward(x, self.param, ...)
    return result

class CustomBackward:
    def backward(self, grad_output):
        # Compute gradients manually
        return grad_input

Option B: Use Autograd-Tracked Operations (Less work, less control)

def forward(self, x):
    # Use operations autograd already tracks
    result = x.matmul(self.weight)  # Autograd tracks this
    result = result + self.bias      # Autograd tracks this
    return result  # Gradient functions attached automatically

---

Layers That Need Fixing

Priority 1: Core Learning Blocks (CRITICAL)

1. Conv2d - Breaks all CNN milestones
2. Embedding - Breaks all NLP milestones
3. MultiHeadAttention - Breaks transformer milestone

Priority 2: Supporting Layers (IMPORTANT)

4. LayerNorm - May break transformer training stability
5. MaxPool2d - If used in training (usually not trainable, but needs grad flow)
6. AvgPool2d - Same as MaxPool2d

Priority 3: Optional Enhancements (NICE TO HAVE)

7. Dropout - Usually handled correctly if using mask multiplication
8. Other activations - Check ReLU, Sigmoid, etc. (likely fine)

---

Testing Strategy

What We Built

Comprehensive learning verification tests in test_learning_verification.py:

def test_cnn_learning():
    """Verifies CNN ACTUALLY LEARNS"""
    model = build_cnn()
    
    # Train the model
    for epoch in range(epochs):
        train_step(model, X, y)
    
    # Verify learning happened:
    ✅ check_gradient_flow(params)      # All params get gradients?
    ✅ check_weight_updates(before, after)  # Weights changed?
    ✅ verify_loss_convergence(history)     # Loss decreased?
    ✅ check_final_accuracy(model)          # Model converged?

How to Use for Debugging

1. Run test for broken layer:

   python tests/milestones/test_learning_verification.py
   

2. Check gradient flow:

   Gradients Flowing: 4/19  ← Only 4 params get gradients!
   Conv Gradients: 0.000000  ← Conv layer completely dead!
   

3. Fix the layer (add gradient function)

4. Re-run test to verify fix

5. Iterate until all checks pass

---

Recommended Fix Order

Phase 1: CNN Fix (Highest Impact)

Time: 4-6 hours Impact: Enables all image processing milestones

1. Implement Conv2dBackward gradient function
2. Modify Conv2d.forward() to attach gradient function
3. Test with test_cnn_learning()
4. Verify actual CNN milestone scripts work

Phase 2: Embedding Fix (High Impact)

Time: 2-3 hours Impact: Enables all NLP milestones

1. Check if EmbeddingBackward exists (may already be implemented)
2. Verify Embedding.forward() attaches gradient function
3. Test with test_transformer_learning()

Phase 3: Attention Fix (High Impact)

Time: 3-4 hours Impact: Completes transformer support

1. Implement AttentionBackward gradient function
2. Modify MultiHeadAttention.forward() to attach gradient function
3. Test with test_transformer_learning()
4. Verify all 19 params get gradients

Phase 4: Verification (Critical)

Time: 2-3 hours Impact: Ensures all fixes work end-to-end

1. Run all learning verification tests
2. Run actual milestone scripts (not just tests)
3. Verify students can complete assignments
4. Update documentation

---

Files Modified So Far

Test Files (Created/Modified)

• ✅ tests/milestones/test_learning_verification.py - Comprehensive learning tests
• ✅ tests/milestones/README.md - Complete documentation
• ✅ tests/milestones/VERIFICATION_SUMMARY.md - Quick overview
• ✅ tests/milestones/FIXES_NEEDED.md - This file

Source Files (Modified)

• ⚠️ tinytorch/core/spatial.py - Added requires_grad=True (insufficient fix)

Source Files (Need Modification)

• 🔴 tinytorch/core/spatial.py - Needs Conv2dBackward implementation
• 🔴 tinytorch/text/embeddings.py - Check/fix gradient flow
• 🔴 tinytorch/core/attention.py - Needs AttentionBackward implementation

---

Summary for User

What Works ✅

1. Perceptron (1957) - Perfect learning, all tests pass
2. XOR (1969) - Perfect learning, all tests pass
3. MLP Digits (1986) - Fixed and verified, passes consistently

What's Broken 🔴

1. CNN (1998) - Conv2d doesn't integrate with autograd

   • Conv layers don't receive gradients
   • Model barely learns (random guessing)
   • Needs Conv2dBackward implementation

2. Transformer (2017) - Attention/Embedding don't integrate with autograd

   • Only 21% of parameters receive gradients
   • Attention and embeddings don't learn
   • Needs EmbeddingBackward + AttentionBackward

The Core Issue

Custom layers use manual numpy operations and bypass autograd entirely.

They need to either:

1. Attach gradient functions to returned Tensors (more work, more control)
2. Use tensor operations that autograd already tracks (less work)

This is a fundamental architectural issue that affects multiple modules.

Next Steps

1. Decision needed: Fix Conv2d first (enables image processing) or Transformer first (enables NLP)?
2. Implementation: Add backward() methods to custom layers
3. Testing: Verify with learning verification tests
4. Validation: Run actual milestone scripts end-to-end

Estimated Total Time

• Conv2d fix: 4-6 hours
• Embedding fix: 2-3 hours
• Attention fix: 3-4 hours
• Testing/validation: 2-3 hours
• Total: 11-16 hours of focused development

---

References

• Learning verification tests: tests/milestones/test_learning_verification.py
• Test documentation: tests/milestones/README.md
• Gradient flow guide: tests/integration/INTERMODULE_TEST_COVERAGE.md
• Transformer gradient notes: milestones/05_2017_transformer/GRADIENT_FLOW_VERIFICATION.md