Training Neural Networks

How Networks Learn

Neural networks learn by minimizing error. The process:

1. Make a prediction (forward pass)
2. Calculate error (how wrong was it?)
3. Adjust weights to reduce error (backward pass)
4. Repeat thousands of times

Loss Functions

Loss measures how wrong the network's predictions are. Lower loss = better predictions.

// Mean Squared Error (regression)
function mseLoss(predictions, targets) {
    let sum = 0;
    for (let i = 0; i < predictions.length; i++) {
        const diff = predictions[i] - targets[i];
        sum += diff * diff;
    }
    return sum / predictions.length;
}

// Cross-entropy loss (classification)
function crossEntropyLoss(predictions, targets) {
    let loss = 0;
    for (let i = 0; i < predictions.length; i++) {
        loss -= targets[i] * Math.log(predictions[i]);
    }
    return loss;
}

// For stem separation: L1 loss (mean absolute error)
function l1Loss(prediction, target) {
    let sum = 0;
    for (let i = 0; i < prediction.length; i++) {
        sum += Math.abs(prediction[i] - target[i]);
    }
    return sum / prediction.length;
}

Backpropagation

Backpropagation calculates how much each weight contributed to the error, then adjusts them.

// Simplified backpropagation
function backpropagate(network, predictions, targets) {
    // 1. Calculate loss
    const loss = mseLoss(predictions, targets);
    
    // 2. Calculate gradients (how much each weight affects loss)
    // Using chain rule from calculus
    const gradients = calculateGradients(network, predictions, targets);
    
    // 3. Update weights
    const learningRate = 0.001;
    for (const layer of network.layers) {
        for (const neuron of layer.neurons) {
            for (let i = 0; i < neuron.weights.length; i++) {
                neuron.weights[i] -= learningRate * gradients[i];
            }
            neuron.bias -= learningRate * gradients.bias;
        }
    }
    
    return loss;
}

// Repeat this process thousands of times
// Each iteration = one "training step"
// Model gradually improves

Training Loop

// Full training loop
async function trainModel(model, trainingData, epochs = 100) {
    for (let epoch = 0; epoch < epochs; epoch++) {
        let totalLoss = 0;
        
        // Shuffle data
        shuffle(trainingData);
        
        // Process in batches
        const batchSize = 32;
        for (let i = 0; i < trainingData.length; i += batchSize) {
            const batch = trainingData.slice(i, i + batchSize);
            
            // Forward pass
            const predictions = batch.map(item => model.forward(item.input));
            
            // Calculate loss
            const loss = batch.reduce((sum, item, idx) => {
                return sum + mseLoss(predictions[idx], item.target);
            }, 0) / batch.length;
            
            // Backward pass
            backpropagate(model, predictions, batch.map(b => b.target));
            
            totalLoss += loss;
        }
        
        console.log(`Epoch ${epoch}: Loss = ${totalLoss}`);
    }
    
    return model;
}

// For stem separation:
// Training data: songs with separated stems
// Input: mixed audio spectrogram
// Target: 4 separate spectrograms (vocals, drums, bass, other)
// Loss: difference between predicted and actual stems

// Training pipeline (done offline, not in browser)
class ModelTrainer {
    async train() {
        // 1. Prepare dataset
        const dataset = await this.loadDataset('musdb18'); // 150 songs with stems
        
        // 2. Create model architecture
        const model = new UNetModel({
            inputChannels: 1, // Mono spectrogram
            outputChannels: 4, // 4 stems
            filters: [16, 32, 64, 128, 256],
            kernelSize: [7, 5],
            strides: [2, 2]
        });
        
        // 3. Training configuration
        const config = {
            learningRate: 0.001,
            batchSize: 16,
            epochs: 100,
            optimizer: 'adam', // Adaptive learning rate
            loss: 'l1' // Mean absolute error
        };
        
        // 4. Training loop
        for (let epoch = 0; epoch < config.epochs; epoch++) {
            for (const batch of this.createBatches(dataset, config.batchSize)) {
                // Extract spectrograms
                const inputs = batch.map(item => this.audioToSpectrogram(item.mix));
                const targets = batch.map(item => ({
                    vocals: this.audioToSpectrogram(item.vocals),
                    drums: this.audioToSpectrogram(item.drums),
                    bass: this.audioToSpectrogram(item.bass),
                    other: this.audioToSpectrogram(item.other)
                }));
                
                // Train on batch
                const loss = await model.trainBatch(inputs, targets, config);
            }
            
            // Validate
            const valLoss = await this.validate(model, validationData);
            console.log(`Epoch ${epoch}: Train Loss = ${trainLoss}, Val Loss = ${valLoss}`);
            
            // Save best model
            if (valLoss < bestValLoss) {
                await model.save('best-model.onnx');
            }
        }
    }
}

// After training: export to ONNX format for web deployment
// ONNX = Open Neural Network Exchange (standard format)