Files
TinyTorch/modules/18_compression/README.md
Vijay Janapa Reddi 2f23f757e7 MAJOR: Implement beautiful module progression through strategic reordering
This commit implements the pedagogically optimal "inevitable discovery" module progression based on expert validation and educational design principles.

## Module Reordering Summary

**Previous Order (Problems)**:
- 05_losses → 06_autograd → 07_dataloader → 08_optimizers → 09_spatial → 10_training
- Issues: Autograd before optimizers, DataLoader before training, scattered dependencies

**New Order (Beautiful Progression)**:
- 05_losses → 06_optimizers → 07_autograd → 08_training → 09_spatial → 10_dataloader
- Benefits: Each module creates inevitable need for the next

## Pedagogical Flow Achieved

**05_losses** → "Need systematic weight updates" → **06_optimizers**
**06_optimizers** → "Need automatic gradients" → **07_autograd**
**07_autograd** → "Need systematic training" → **08_training**
**08_training** → "MLPs hit limits on images" → **09_spatial**
**09_spatial** → "Training is too slow" → **10_dataloader**

## Technical Changes

### Module Directory Renaming
- `06_autograd` → `07_autograd`
- `07_dataloader` → `10_dataloader`
- `08_optimizers` → `06_optimizers`
- `10_training` → `08_training`
- `09_spatial` → `09_spatial` (no change)

### System Integration Updates
- **MODULE_TO_CHECKPOINT mapping**: Updated in tito/commands/export.py
- **Test directories**: Renamed module_XX directories to match new numbers
- **Documentation**: Updated all references in MD files and agent configurations
- **CLI integration**: Updated next-steps suggestions for proper flow

### Agent Configuration Updates
- **Quality Assurance**: Updated module audit status with new numbers
- **Module Developer**: Updated work tracking with new sequence
- **Documentation**: Updated MASTER_PLAN_OF_RECORD.md with beautiful progression

## Educational Benefits

1. **Inevitable Discovery**: Each module naturally leads to the next
2. **Cognitive Load**: Concepts introduced exactly when needed
3. **Motivation**: Students understand WHY each tool is necessary
4. **Synthesis**: Everything flows toward complete ML systems understanding
5. **Professional Alignment**: Matches real ML engineering workflows

## Quality Assurance

-  All CLI commands still function
-  Checkpoint system mappings updated
-  Documentation consistency maintained
-  Test directory structure aligned
-  Agent configurations synchronized

**Impact**: This reordering transforms TinyTorch from a collection of modules into a coherent educational journey where each step naturally motivates the next, creating optimal conditions for deep learning systems understanding.
2025-09-24 15:56:47 -04:00

3.2 KiB

Module 18: Compression - Model Size Optimization

Overview

Reduce model size by 90% while maintaining accuracy through pruning and distillation. Learn how production systems deploy efficient models at scale.

What You'll Build

  • Magnitude Pruner: Remove unimportant weights
  • Structured Pruning: Remove entire channels/layers
  • Knowledge Distillation: Transfer knowledge to smaller models
  • Sparse Inference: Efficient computation with pruned models

Learning Objectives

  1. Sparsity Patterns: Structured vs unstructured pruning
  2. Pruning Strategies: Magnitude, gradient, lottery ticket
  3. Distillation: Teacher-student knowledge transfer
  4. Deployment: Optimize sparse models for production

Prerequisites

  • Module 10: Training (models to compress)
  • Module 17: Precision (understanding of optimization tradeoffs)

Key Concepts

Magnitude-Based Pruning

# Remove 90% of smallest weights
def prune_magnitude(model, sparsity=0.9):
    for layer in model.layers:
        threshold = torch.quantile(abs(layer.weight), sparsity)
        mask = abs(layer.weight) > threshold
        layer.weight *= mask  # Zero out small weights

Structured Pruning

# Remove entire filters/channels
def prune_structured(conv_layer, num_filters_to_remove):
    # Compute filter importance (L2 norm)
    importance = conv_layer.weight.norm(dim=(1,2,3))
    
    # Keep only important filters
    keep_indices = importance.topk(n_keep).indices
    conv_layer.weight = conv_layer.weight[keep_indices]

Knowledge Distillation

# Small student learns from large teacher
teacher = LargeModel()  # 100M parameters
student = SmallModel()  # 10M parameters

# Student learns both from labels and teacher
loss = alpha * cross_entropy(student(x), y) + \
       beta * kl_divergence(student(x), teacher(x))

Performance Impact

  • Model Size: 10x reduction with pruning
  • Inference Speed: 3-5x faster with structured pruning
  • Accuracy: Maintain 95%+ of original performance
  • Memory: Deploy large models on edge devices

Real-World Applications

  • MobileNet: Designed for mobile deployment
  • DistilBERT: 60% faster, 97% performance
  • Lottery Ticket Hypothesis: Finding efficient subnetworks
  • Neural Architecture Search: Automated compression

Module Structure

  1. Sparsity Theory: Why neural networks are compressible
  2. Magnitude Pruning: Simple but effective compression
  3. Structured Pruning: Hardware-friendly sparsity
  4. Knowledge Distillation: Learning from larger models
  5. Deployment: Optimizing sparse models

Hands-On Projects

# Project 1: Prune your CNN
cnn = load_model("cifar10_cnn.pt")
pruned = progressive_prune(cnn, target_sparsity=0.9)
print(f"Parameters: {count_params(cnn)}{count_params(pruned)}")
print(f"Accuracy: {evaluate(cnn)}% → {evaluate(pruned)}%")

# Project 2: Distill transformer to CNN
teacher = TinyTransformer()  
student = SimpleCNN()
distilled = distill(teacher, student, data_loader)

Success Criteria

  • Achieve 90% sparsity with <5% accuracy loss
  • 3x inference speedup with structured pruning
  • Successfully distill large models to small ones
  • Deploy compressed models efficiently