Training a Neural Network

Putting It Together

Now that we have covered forward propagation, backpropagation, and gradient checking, let’s combine everything into a complete training pipeline.

1. 🔀 Choose a Network Architecture

First, decide the structure of your neural network:

• Number of layers $L$
• Number of hidden units per layer $j$
• Number of Outputs $y$

How to choose Network

• Input layer size = dimension of feature vector $x^{(i)}$
• Output layer size = number of output classes
• Hidden units:
  • More units usually perform better
  • But increase computational cost
• Default choice:
  • Use 1 hidden layer
  • If using multiple hidden layers, use the same number of units in each layer

---

2. 📚 Training a Neural Network

2.1 🎲 Randomly Initialize Weights

Initialize each $\Theta^{(l)}$ randomly (not to zero).

This breaks symmetry and allows learning.

2.2 ⏩ Forward Propagation (FP)

For each training example $x^{(i)}$, compute:

$$h_\Theta(x^{(i)})$$

This gives the network’s prediction.

2.3 💰 Implement the Cost Function

Compute:

$$J(\Theta)$$

This includes:

• Logistic loss over all output units
• Regularization term

2.4 ⏪ Backpropagation (BP)

Use backpropagation to compute:

$$\frac{\partial}{\partial \Theta_{i,j}^{(l)}} J(\Theta)$$

This gives the gradients needed for optimization.

2.5 🎢 Gradient Checking

Use numerical approximation to verify backpropagation:

$$\frac{\partial}{\partial \Theta} J(\Theta) \approx \frac{J(\Theta + \epsilon) - J(\Theta - \epsilon)}{2\epsilon}$$

⚠️ Once verified:

• Disable gradient checking
• It is computationally expensive

2.6 ⚖️ Minimize the Cost Function

Use:

• Gradient descent, or
• A built-in optimization algorithm (e.g., advanced optimizers)

to minimize $J(\Theta)$.

---

Training Loop

During training, we iterate over all examples:

for i = 1:m
    % Forward propagation
    % Compute activations a^(l)

    % Backpropagation
    % Compute delta terms d^(l) for l = 2,...,L
end

For each example:

• Perform forward pass
• Compute errors
• Accumulate gradients

---

Final Insight

Neural network training is simply:

• Forward propagation
• Backpropagation
• Gradient-based optimization

All of deep learning is built on this foundation.

Complete Neural Network Workflow

1. Choose architecture
2. Initialize weights randomly
3. Implement forward propagation
4. Implement cost function
5. Implement backpropagation
6. Perform gradient checking
7. Optimize using gradient descent
8. Train until convergence

---

Modern Relevance

The 8-step procedure here maps directly to a production PyTorch training loop — expanded to run across hundreds of GPUs on a DGX cluster.

Step-by-step modern mapping:

Step (this post)	Modern equivalent
1. Choose architecture	Model card selection: 7B, 13B, 70B, or custom; transformer depth/width config
2. Randomly initialize weights	kaiming_normal_ / xavier_uniform_ in PyTorch; or load pre-trained checkpoint
2.1 Forward propagation	output = model(input_ids) — runs on 8 H100s with tensor parallelism
2.2 Cost function	loss = F.cross_entropy(logits, target_ids)
2.3 Backpropagation	loss.backward() — PyTorch autograd traverses computation graph
3. Gradient descent	optimizer.step() — AdamW with cosine LR schedule and weight decay
4. Gradient checking	torch.autograd.gradcheck() — used when writing custom CUDA kernels only
Repeat until convergence	Training loop over epochs; early stopping on validation loss plateau

At DGX scale, steps 2.1–2.3 are split across multiple GPUs:

• Data parallelism: each GPU processes a different mini-batch, runs its own forward+backward, then NCCL AllReduce synchronizes gradients before the optimizer step
• Tensor parallelism: the weight matrix in each layer is split across GPUs; forward and backward both require NCCL AllReduce within each layer
• Gradient accumulation: for effective batch sizes larger than GPU memory allows, the optimizer.step() is deferred until N mini-batches have accumulated gradients

The output of this training procedure — the trained weight tensors — is what TensorRT-LLM compiles to a GPU-optimized inference engine, which NIM then packages and serves behind an OpenAI-compatible API.

---

Related Posts

• Backpropagation Algorithm — the core algorithm in step 4 of this training procedure; how gradients propagate backward through each layer
• Bias-Variance Dilemma — after training, diagnose whether the network is underfitting (high bias) or overfitting (high variance) using training vs validation error
• What are Transformer Models? — transformers follow the same architecture selection → forward propagation → cost → backprop → gradient descent procedure described here
• TensorRT and High-Performance AI Inference — the network trained by this procedure is what TensorRT compiles into a GPU-optimized engine for production inference
• Multi-Node Distributed Training on Kubernetes — scaling this training procedure across multiple GPU nodes using PyTorchJob and data parallelism