First Training Loop - Entrenar - Training & Optimization Library
Description: A comprehensive guide to building neural network training systems with autograd, optimizers, LoRA/QLoRA, and quantization using EXTREME TDD methodology

# First Training Loop

This guide will walk you through building a complete, production-ready training pipeline with validation, checkpointing, and early stopping.

## Complete Training Example

We'll train a multi-layer perceptron (MLP) on a simple classification task with all best practices included.

### Project Structure

    first-training-loop/
    ├── Cargo.toml
    └── src/
        ├── main.rs          # Training script
        ├── model.rs         # Model definition
        └── data.rs          # Data loading

### Model Definition

Create `src/model.rs`:

    #![allow(unused)]
    fn main() {
    use entrenar::{Tensor, autograd::ops::{matmul, relu}};

    pub struct MLP {
        pub w1: Tensor,
        pub b1: Tensor,
        pub w2: Tensor,
        pub b2: Tensor,
    }

    impl MLP {
        /// Create a new 2-layer MLP: input_dim -> hidden_dim -> output_dim
        pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize) -> Self {
            // Xavier/Glorot initialization
            let scale1 = (2.0 / (input_dim + hidden_dim) as f32).sqrt();
            let scale2 = (2.0 / (hidden_dim + output_dim) as f32).sqrt();

            Self {
                w1: Tensor::randn(vec![hidden_dim * input_dim], true) * scale1,
                b1: Tensor::zeros(vec![hidden_dim], true),
                w2: Tensor::randn(vec![output_dim * hidden_dim], true) * scale2,
                b2: Tensor::zeros(vec![output_dim], true),
            }
        }

        /// Forward pass
        pub fn forward(&self, x: &Tensor, input_dim: usize, hidden_dim: usize, output_dim: usize, batch_size: usize) -> Tensor {
            // Layer 1: h = relu(W1 * x + b1)
            let h = relu(&(
                &matmul(&self.w1, x, hidden_dim, input_dim, batch_size) + &self.b1
            ));

            // Layer 2: y = W2 * h + b2
            let y = &matmul(&self.w2, &h, output_dim, hidden_dim, batch_size) + &self.b2;

            y
        }

        /// Get all trainable parameters
        pub fn parameters(&mut self) -> Vec<&mut Tensor> {
            vec![&mut self.w1, &mut self.b1, &mut self.w2, &mut self.b2]
        }

        /// Zero all gradients
        pub fn zero_grad(&mut self) {
            for param in self.parameters() {
                param.zero_grad();
            }
        }
    }
    }

### Data Loading

Create `src/data.rs`:

    #![allow(unused)]
    fn main() {
    use entrenar::Tensor;

    /// Generate synthetic XOR dataset
    pub fn generate_xor_data(n_samples: usize) -> (Vec<Vec<f32>>, Vec<f32>) {
        let mut x_data = Vec::new();
        let mut y_data = Vec::new();

        for _ in 0..n_samples {
            let x1 = if rand::random::<f32>() > 0.5 { 1.0 } else { 0.0 };
            let x2 = if rand::random::<f32>() > 0.5 { 1.0 } else { 0.0 };

            // XOR: output is 1 if inputs differ
            let y = if (x1 > 0.5) != (x2 > 0.5) { 1.0 } else { 0.0 };

            x_data.push(vec![x1, x2]);
            y_data.push(y);
        }

        (x_data, y_data)
    }

    /// Split data into train/validation sets
    pub fn train_val_split(
        x: Vec<Vec<f32>>,
        y: Vec<f32>,
        val_ratio: f32,
    ) -> ((Vec<Vec<f32>>, Vec<f32>), (Vec<Vec<f32>>, Vec<f32>)) {
        let n = x.len();
        let n_val = (n as f32 * val_ratio) as usize;
        let n_train = n - n_val;

        let x_train = x[..n_train].to_vec();
        let y_train = y[..n_train].to_vec();
        let x_val = x[n_train..].to_vec();
        let y_val = y[n_train..].to_vec();

        ((x_train, y_train), (x_val, y_val))
    }

    /// Create mini-batches
    pub fn create_batches(
        x: &[Vec<f32>],
        y: &[f32],
        batch_size: usize,
    ) -> Vec<(Tensor, Tensor)> {
        let mut batches = Vec::new();

        for i in (0..x.len()).step_by(batch_size) {
            let end = (i + batch_size).min(x.len());
            let batch_x: Vec<f32> = x[i..end].iter().flatten().copied().collect();
            let batch_y: Vec<f32> = y[i..end].to_vec();

            batches.push((
                Tensor::from_vec(batch_x, false),
                Tensor::from_vec(batch_y, false),
            ));
        }

        batches
    }
    }

### Training Script

Create `src/main.rs`:

    mod model;
    mod data;

    use entrenar::{backward, optim::Adam};
    use model::MLP;
    use data::{generate_xor_data, train_val_split, create_batches};

    fn main() {
        println!("=== Entrenar Training Example: XOR Problem ===\n");

        // Hyperparameters
        let input_dim = 2;
        let hidden_dim = 8;
        let output_dim = 1;
        let learning_rate = 0.01;
        let batch_size = 32;
        let n_epochs = 100;
        let val_ratio = 0.2;
        let patience = 10;  // Early stopping patience

        // Generate data
        let (x_data, y_data) = generate_xor_data(1000);
        let ((x_train, y_train), (x_val, y_val)) = train_val_split(x_data, y_data, val_ratio);

        println!("Dataset:");
        println!("  Training samples: {}", x_train.len());
        println!("  Validation samples: {}", x_val.len());
        println!();

        // Create model and optimizer
        let mut model = MLP::new(input_dim, hidden_dim, output_dim);
        let mut optimizer = Adam::default_params(learning_rate);

        // Early stopping tracker
        let mut best_val_loss = f32::INFINITY;
        let mut patience_counter = 0;

        // Training loop
        for epoch in 0..n_epochs {
            // Training phase
            let train_batches = create_batches(&x_train, &y_train, batch_size);
            let mut train_loss = 0.0;

            for (batch_x, batch_y) in &train_batches {
                // Forward pass
                let y_pred = model.forward(
                    batch_x,
                    input_dim,
                    hidden_dim,
                    output_dim,
                    batch_x.data().len() / input_dim,
                );

                // Binary cross-entropy loss
                let loss = binary_cross_entropy(&y_pred, batch_y);
                train_loss += loss.data()[0];

                // Backward pass
                backward(&loss);

                // Update parameters
                optimizer.step(&mut model.parameters());

                // Zero gradients
                model.zero_grad();
            }

            train_loss /= train_batches.len() as f32;

            // Validation phase
            let val_batches = create_batches(&x_val, &y_val, batch_size);
            let mut val_loss = 0.0;

            for (batch_x, batch_y) in &val_batches {
                let y_pred = model.forward(
                    batch_x,
                    input_dim,
                    hidden_dim,
                    output_dim,
                    batch_x.data().len() / input_dim,
                );

                let loss = binary_cross_entropy(&y_pred, batch_y);
                val_loss += loss.data()[0];
            }

            val_loss /= val_batches.len() as f32;

            // Early stopping check
            if val_loss < best_val_loss {
                best_val_loss = val_loss;
                patience_counter = 0;
                println!("Epoch {:3}: train_loss={:.4}, val_loss={:.4} ✓ (best)", epoch, train_loss, val_loss);
            } else {
                patience_counter += 1;
                println!("Epoch {:3}: train_loss={:.4}, val_loss={:.4}   (patience: {}/{})",
                         epoch, train_loss, val_loss, patience_counter, patience);

                if patience_counter >= patience {
                    println!("\nEarly stopping triggered!");
                    break;
                }
            }
        }

        println!("\n=== Training Complete ===");
        println!("Best validation loss: {:.4}", best_val_loss);
    }

    /// Binary cross-entropy loss: -[y*log(p) + (1-y)*log(1-p)]
    fn binary_cross_entropy(y_pred: &Tensor, y_true: &Tensor) -> Tensor {
        // Sigmoid activation
        let sigmoid = |x: f32| 1.0 / (1.0 + (-x).exp());

        let pred_data: Vec<f32> = y_pred.data().iter().map(|&x| sigmoid(x)).collect();
        let true_data = y_true.data();

        let mut loss = 0.0;
        for (p, y) in pred_data.iter().zip(true_data.iter()) {
            let p_clamped = p.clamp(1e-7, 1.0 - 1e-7);  // Numerical stability
            loss += -y * p_clamped.ln() - (1.0 - y) * (1.0 - p_clamped).ln();
        }

        Tensor::from_vec(vec![loss / pred_data.len() as f32], false)
    }

### Running the Training

    cargo run --release

Expected output:

    === Entrenar Training Example: XOR Problem ===

    Dataset:
      Training samples: 800
      Validation samples: 200

    Epoch   0: train_loss=0.7123, val_loss=0.7001 ✓ (best)
    Epoch   1: train_loss=0.6845, val_loss=0.6723 ✓ (best)
    Epoch   2: train_loss=0.6234, val_loss=0.6102 ✓ (best)
    ...
    Epoch  42: train_loss=0.0523, val_loss=0.0498 ✓ (best)
    Epoch  43: train_loss=0.0501, val_loss=0.0512   (patience: 1/10)
    ...
    Epoch  52: train_loss=0.0412, val_loss=0.0556   (patience: 10/10)

    Early stopping triggered!

    === Training Complete ===
    Best validation loss: 0.0498

## Key Components Explained

### 1. Xavier Initialization

    #![allow(unused)]
    fn main() {
    let scale = (2.0 / (input_dim + output_dim) as f32).sqrt();
    let w = Tensor::randn(shape, true) * scale;
    }

- Prevents vanishing/exploding gradients
- Scales weights based on layer dimensions

### 2. Mini-Batch Training

    #![allow(unused)]
    fn main() {
    let batches = create_batches(&x_train, &y_train, batch_size=32);
    }

- Processes multiple samples together
- Reduces training time via batched operations
- Provides gradient noise for better generalization

### 3. Train/Validation Split

    #![allow(unused)]
    fn main() {
    let ((x_train, y_train), (x_val, y_val)) = train_val_split(data, 0.2);
    }

- 80% training, 20% validation
- Validation set detects overfitting
- Never use validation data for gradient updates

### 4. Early Stopping

    #![allow(unused)]
    fn main() {
    if val_loss < best_val_loss {
        best_val_loss = val_loss;
        patience_counter = 0;
    } else {
        patience_counter += 1;
        if patience_counter >= patience {
            break;  // Stop training
        }
    }
    }

- Prevents overfitting
- Stops when validation loss stops improving
- Saves computational resources

### 5. Gradient Flow

    #![allow(unused)]
    fn main() {
    backward(&loss);             // Compute gradients
    optimizer.step(&mut params); // Update parameters
    model.zero_grad();           // Clear gradients for next iteration
    }

- Critical: Zero gradients after each step
- Gradients accumulate by default in Entrenar

## Advanced Features

### Checkpointing

Save model state periodically:

    #![allow(unused)]
    fn main() {
    use std::fs::File;
    use std::io::Write;

    if epoch % 10 == 0 {
        let checkpoint = serde_json::json!({
            "epoch": epoch,
            "w1": model.w1.data(),
            "b1": model.b1.data(),
            "w2": model.w2.data(),
            "b2": model.b2.data(),
            "best_val_loss": best_val_loss,
        });

        let mut file = File::create(format!("checkpoint_epoch_{}.json", epoch))?;
        file.write_all(checkpoint.to_string().as_bytes())?;
    }
    }

### Learning Rate Scheduling

Decay learning rate over time:

    #![allow(unused)]
    fn main() {
    use entrenar::optim::schedulers::CosineScheduler;

    let scheduler = CosineScheduler::new(0.01, 0.0001, n_epochs * batches_per_epoch);

    for step in 0.. {
        let lr = scheduler.get_lr(step);
        optimizer.set_lr(lr);

        // ... training step ...
    }
    }

### Gradient Clipping

Prevent exploding gradients:

    #![allow(unused)]
    fn main() {
    use entrenar::optim::clip_grad_norm;

    backward(&loss);

    // Clip gradients to max norm of 1.0
    clip_grad_norm(&mut model.parameters(), 1.0);

    optimizer.step(&mut model.parameters());
    }

### Logging and Metrics

Track additional metrics:

    #![allow(unused)]
    fn main() {
    struct Metrics {
        train_losses: Vec<f32>,
        val_losses: Vec<f32>,
        train_accuracies: Vec<f32>,
        val_accuracies: Vec<f32>,
    }

    impl Metrics {
        fn log(&mut self, epoch: usize, train_loss: f32, val_loss: f32, train_acc: f32, val_acc: f32) {
            self.train_losses.push(train_loss);
            self.val_losses.push(val_loss);
            self.train_accuracies.push(train_acc);
            self.val_accuracies.push(val_acc);

            println!("Epoch {}: train_loss={:.4} train_acc={:.2}% | val_loss={:.4} val_acc={:.2}%",
                     epoch, train_loss, train_acc * 100.0, val_loss, val_acc * 100.0);
        }

        fn save(&self, path: &str) -> std::io::Result<()> {
            let json = serde_json::to_string_pretty(&self)?;
            std::fs::write(path, json)?;
            Ok(())
        }
    }
    }

## Best Practices

### ✅ Do's

1. Always use release mode for training: `cargo run --release`
2. Validate hyperparameters on a small dataset first
3. Monitor both training and validation loss to detect overfitting
4. Use early stopping to prevent unnecessary computation
5. Zero gradients after each optimizer step
6. Checkpoint regularly to resume interrupted training

### ❌ Don'ts

1. Don't train in debug mode (10-100x slower)
2. Don't use validation data for training (data leakage)
3. Don't forget to zero gradients (leads to incorrect updates)
4. Don't use tiny learning rates (<1e-6) without a good reason
5. Don't ignore validation loss (only watching training loss hides overfitting)

## Troubleshooting

### Loss is NaN

Causes:

- Learning rate too high
- Numerical instability in loss function

Solutions:

- Reduce learning rate (try 0.001, 0.0001)
- Add gradient clipping: `clip_grad_norm(&mut params, 1.0)`
- Clamp predictions: `p.clamp(1e-7, 1.0 - 1e-7)`

### Training is Slow

Causes:

- Running in debug mode
- Batch size too small
- SIMD not activating

Solutions:

- Use `cargo run --release`
- Increase batch size (32, 64, 128)
- Ensure tensors are ≥16 elements for SIMD

### Validation Loss Increases

Cause: Overfitting

Solutions:

- Enable early stopping
- Reduce model size (fewer parameters)
- Add regularization (L2 weight decay)
- Increase dataset size

## What's Next?

- Core Concepts - Understand Entrenar's architecture
- Autograd Engine - Learn how automatic differentiation works
- Optimizers - Explore SGD, Adam, AdamW, and schedulers

Ready to dive deeper? Continue to Core Concepts →