Neural Networks: Gradients, Losses, and Generalization

TL;DR

A neural network composes parameterized functions. An affine transformation mixes features, nonlinearities increase the class of functions that can be represented, and optimization adjusts parameters using gradients of a chosen objective. Backpropagation is an efficient application of the chain rule; it is not a biological simulation or a guarantee of understanding.

The architecture, loss, data split, initialization, optimizer, regularization, and evaluation protocol form one system. A lower training loss can coexist with worse calibration, robustness, fairness, or held-out performance because of overfitting. Always compare against a simple baseline on a held-out test set.

From Neurons to Functions

The biological-neuron analogy is useful for intuition but not identity. Biological neurons are electrochemical, recurrent, adaptive systems; an artificial layer is a numerical operation implemented by software and hardware. Claims about “nanosecond neurons,” brain equivalence, or energy comparisons require a common workload and measurement boundary and should not be inferred from a diagram.

For a batch of inputs X, a dense layer computes:

Z = X Wᵀ + b
A = φ(Z)

W and b are learned parameters, and φ is a nonlinearity. Without nonlinearities, stacking affine layers is still one affine function. The input and output dimensions are properties of the data contract; “input neurons” do not themselves perform a separate biological-style computation.

Network Shapes and Output Semantics

Write tensor shapes next to equations and tests:

X: [batch, features]
W: [outputs, features]
b: [outputs]
Z: [batch, outputs]

The output head depends on the target:

Pass logits to numerically stable framework losses. Applying softmax before a cross-entropy function that expects logits can double-apply normalization and reduce numerical stability.

Activation Functions

ReLU and Variants

ReLU(x) = max(0, x)

ReLU is cheap and often works well, but units can become inactive for all inputs. Leaky ReLU, GELU, SiLU, and other variants change the optimization and compute trade-offs; none is universally best.

Sigmoid and Tanh

Sigmoid maps to (0, 1) and is useful as a probability interpretation at a binary output boundary, but it can saturate. Tanh is zero-centered and maps to (-1, 1), yet also saturates. Hidden-layer choices should be tied to architecture, normalization, initialization, and measured gradients.

Softmax

Softmax maps a vector of logits to a categorical distribution:

softmax(zᵢ) = exp(zᵢ - max(z)) / Σⱼ exp(zⱼ - max(z))

The subtraction is a numerical-stability trick. A probability vector is not automatically calibrated; calibration must be evaluated separately.

Forward Pass and Backpropagation

For a two-layer network:

z₁ = xW₁ᵀ + b₁
a₁ = φ(z₁)
z₂ = a₁W₂ᵀ + b₂
loss = L(z₂, y)

Backpropagation applies the chain rule from the loss toward earlier operations. For a cross-entropy loss with class logits, the output gradient has a convenient form p - y under the usual one-hot convention. Automatic differentiation computes these vector-Jacobian products; it does not require constructing a full Jacobian.

Parameter updates take the form:

θ ← θ - η ∇θ L

where η is a gradient-descent learning-rate policy. Gradients are estimates when using a minibatch, and optimizer state changes the update trajectory.

A Shape-Safe Go Forward Pass

The following dependency-free Go program evaluates a two-layer MLP and computes stable multiclass cross-entropy directly from logits. It validates matrix shapes instead of relying on implicit broadcasting. The fixed weights make the output reproducible; they are not trained parameters or recommended hyperparameters.

package main

import (
	"errors"
	"fmt"
	"math"
)

func dense(input, weights [][]float64, bias []float64) ([][]float64, error) {
	if len(input) == 0 || len(input[0]) == 0 {
		return nil, errors.New("input must be a non-empty [batch, features] matrix")
	}
	if len(weights) == 0 || len(weights) != len(bias) {
		return nil, errors.New("weights and bias output dimensions must match")
	}

	inputFeatures := len(input[0])
	for i, row := range input {
		if len(row) != inputFeatures {
			return nil, fmt.Errorf("input row %d has inconsistent width", i)
		}
	}
	for i, row := range weights {
		if len(row) != inputFeatures {
			return nil, fmt.Errorf("weight row %d has width %d, want %d", i, len(row), inputFeatures)
		}
	}

	output := make([][]float64, len(input))
	for batch, sample := range input {
		output[batch] = make([]float64, len(weights))
		for unit, row := range weights {
			sum := bias[unit]
			for feature, value := range sample {
				sum += value * row[feature]
			}
			output[batch][unit] = sum
		}
	}
	return output, nil
}

func relu(matrix [][]float64) {
	for i := range matrix {
		for j := range matrix[i] {
			matrix[i][j] = math.Max(0, matrix[i][j])
		}
	}
}

func crossEntropy(logits [][]float64, labels []int) (float64, error) {
	if len(logits) == 0 || len(logits) != len(labels) {
		return 0, errors.New("logits and labels must have the same batch size")
	}

	var loss float64
	for i, row := range logits {
		if labels[i] < 0 || labels[i] >= len(row) {
			return 0, fmt.Errorf("label %d is outside [0, %d)", labels[i], len(row))
		}
		maxLogit := row[0]
		for _, value := range row[1:] {
			maxLogit = math.Max(maxLogit, value)
		}
		var denominator float64
		for _, value := range row {
			denominator += math.Exp(value - maxLogit)
		}
		logProbability := row[labels[i]] - maxLogit - math.Log(denominator)
		loss -= logProbability
	}
	return loss / float64(len(logits)), nil
}

func main() {
	features := [][]float64{
		{1, 0, 0, 1},
		{0, 1, 1, 0},
	}
	hiddenWeights := [][]float64{
		{0.5, -0.5, 0.25, 0.75},
		{-0.25, 0.5, 0.75, -0.5},
		{0.4, 0.1, -0.3, 0.2},
	}
	hidden, err := dense(features, hiddenWeights, []float64{0.1, -0.1, 0})
	if err != nil {
		panic(err)
	}
	relu(hidden)

	outputWeights := [][]float64{
		{0.8, -0.2, 0.3},
		{-0.4, 0.7, 0.1},
		{0.1, -0.3, 0.9},
	}
	logits, err := dense(hidden, outputWeights, []float64{0, 0.1, -0.1})
	if err != nil {
		panic(err)
	}
	loss, err := crossEntropy(logits, []int{0, 1})
	if err != nil {
		panic(err)
	}

	fmt.Printf("hidden shape: %d x %d\n", len(hidden), len(hidden[0]))
	fmt.Printf("logit shape: %d x %d\n", len(logits), len(logits[0]))
	fmt.Printf("mean cross-entropy: %.6f\n", loss)
}

Expected output:

hidden shape: 2 x 3
logit shape: 2 x 3
mean cross-entropy: 0.493668

Production training also needs backpropagation or automatic differentiation, an optimizer, data validation, train/validation/test separation, dtype and device policy, checkpoints, and deliberate gradient accumulation. Framework losses such as PyTorch CrossEntropyLoss accept unnormalized logits; applying Softmax first would duplicate part of the stable loss calculation.

Losses and Optimizers

Choose a loss that matches the label semantics. MSE treats deviations symmetrically; MAE is less sensitive to outliers; Huber interpolates behaviors. Cross-entropy compares a categorical target with logits, while binary cross-entropy with logits handles independent binary labels.

SGD, momentum, Adam, and AdamW are update rules, not guarantees of finding a global optimum. AdamW's decoupled weight decay is not identical to adding an L2 penalty in every parameterization. Learning-rate warmup, schedules, gradient clipping, batch size, and normalization can matter as much as the optimizer name.

Record optimizer type, learning rate, betas/momentum, weight decay, clipping, scheduler, effective batch size, precision, and seed. A statement such as “Adam is always the default choice” is an empirical shortcut, not a theorem.

Initialization, Normalization, and Gradient Health

Poor initialization can produce exploding or vanishing activations and gradients. Xavier/Glorot and He/Kaiming schemes match common assumptions about fan-in and activation families, but architecture and normalization alter the appropriate choice.

Monitor:

• activation and gradient distributions;
• NaN/Inf counts and loss scale;
• update-to-weight ratios;
• dead or saturated units;
• train/validation loss and metric gaps.

Batch normalization, layer normalization, residual connections, and careful initialization change optimization dynamics. They are not interchangeable “regularization switches.”

CNNs, RNNs, and Transformers

Convolutional Networks

Convolutions use local connectivity and shared weights, creating useful translation-related inductive biases for grid-like inputs. Padding, stride, dilation, channel order, normalization, and input resolution determine tensor shapes. Avoid hard-coding a flatten size unless the input shape is fixed and asserted; adaptive pooling or a shape test is safer.

Recurrent Networks

RNNs carry a hidden state across sequence steps. Vanilla recurrence can struggle with long dependencies; LSTM and GRU gates change that trade-off. Mask padding, place hidden states on the same device and dtype as inputs, and define whether the task uses the final state or every timestep. A recurrent model's wall-clock latency also depends on sequential dependence and batching.

Transformers

Self-attention lets tokens interact through query, key, and value projections. A standard attention layer has quadratic pairwise interaction in sequence length, although kernels, sparsity, recurrence, and other designs change practical scaling. Transformers still require positional information or another order mechanism, attention masks, residual paths, normalization, and a task-specific head.

PyTorch modules have explicit shape conventions. Configure batch_first=True when supplying [batch, sequence, embedding], pass padding/causal masks deliberately, and test the output shape. Do not present a block without masks or positional treatment as a complete language model.

Training Protocol and Generalization

Split data before augmentation or repeated tuning. Keep validation for model and hyperparameter decisions and reserve the test set for final reporting. Deduplicate near-identical examples and check temporal, entity, and benchmark contamination. The scikit-learn guide on data leakage illustrates why preprocessing must be fitted only on training data.

Regularization options include weight decay, dropout, data augmentation, label smoothing, early stopping, smaller models, and better data. Each changes the bias/variance trade-off; “add dropout” is not a diagnosis.

Report:

• data source, license, preprocessing, and split policy;
• model revision, tokenizer/features, architecture, and parameter count;
• optimizer, schedule, batch, precision, hardware, seeds, and runtime;
• baseline, metrics, confidence intervals, and error slices;
• calibration, robustness, fairness, privacy, and out-of-distribution behavior where relevant.

Reproducibility is not the same as determinism. GPU kernels, parallelism, library versions, and hardware can produce small differences even with a fixed seed.

Evaluation and Debugging

Use a minimal baseline before increasing model size. Check in this order:

1. data labels, ranges, missingness, leakage, and class balance;
2. tensor shapes, dtypes, masks, and label alignment;
3. loss/activation pairing and gradient flow;
4. train/validation curves and a deliberately tiny overfit test;
5. held-out slices, calibration, and failure examples;
6. deployment preprocessing parity and monitoring.

Accuracy alone can hide class imbalance. Add precision/recall, F1, AUROC/PR-AUC where appropriate, calibration, confusion matrices, cost-weighted errors, or task-specific metrics. For generative systems, evaluate factuality, refusal, toxicity, tool calls, citations, and human-rated usefulness with a fixed protocol.

Frequently Asked Questions

Are neural networks modeled on the brain?

The analogy is historical and conceptual. Artificial layers are mathematical functions optimized on digital hardware; they are not faithful simulations of biological neurons.

Should every hidden layer use ReLU?

No. ReLU is a common baseline, but GELU, SiLU, tanh, gated units, normalization, initialization, and architecture can change the result. Test the choice on the task.

Should I apply softmax before cross-entropy?

Usually no when using a framework loss that expects logits, such as PyTorch CrossEntropyLoss. It applies the stable log-softmax operation internally. Follow the exact API contract.

Why can training loss fall while test performance worsens?

The model may overfit, the split may be contaminated, the metric may not represent the task, or preprocessing may differ. Inspect held-out slices, data lineage, calibration, and baseline comparisons.

Do nanoseconds or more layers imply a better model?

No. Hardware timing, parameter count, and depth are not quality measures by themselves. Compare task performance, resource cost, robustness, and uncertainty under a documented protocol.

Is backpropagation the same as biological learning?

No. Backpropagation is an algorithm for differentiating a computational graph. Biological learning involves mechanisms and constraints that are not represented by the ordinary training loop.

Primary Sources

• Goodfellow, Bengio, and Courville: Deep Learning
• He et al.: Delving Deep into Rectifiers
• Kingma and Ba: Adam
• Loshchilov and Hutter: Decoupled Weight Decay Regularization
• Vaswani et al.: Attention Is All You Need
• PyTorch documentation
• NIST AI Risk Management Framework

Conclusion

Neural networks are compositions of functions trained against an objective, not magic abstractions or literal digital brains. Track shapes, numerical stability, gradients, data lineage, generalization, and deployment behavior together. The broader deep learning guide connects these mechanics to architecture selection and production evaluation.