Neural Networks and Back Propagation

Conditional probability distribution

Our mental model about neural networks is a system where we feed an input $x$ and the model will predict an output $y$. But the real world systems are not that easy. Data we train on contains noise, ambiguity, and natural uncertainty. To build stronger models, **we train it to approximate a conditional probability distribution, $p(y|x)$ rather than a single value.

Conditional probability distribution: instead of predicting a single value, predicting a range of possibilities.

If we look at the true underlying data distribution, $p_{\text{data}}(x, y)$, we can break it down using the product rule of probability: $$p_{\text{data}}(x, y) = p_{\text{data}}(x)p_{\text{data}}(y|x)$$

To generate realistic data, we first sample an input $x$, and then we sample the corresponding output $y$ from a conditional distribution (more specifically a bell curve). Therefore, approximating the true generating function is fundamentally equivalent to estimating $p(y|x)$ i.e, given $x$ what is $y$.

Instead of predicting $y$ directly, our neural network outputs the parameters ($\theta(x)$) of a chosen probability distribution. The architecture of our network doesn't change drastically; only our interpretation of the output layer does.

Binary Outputs (The Bernoulli Distribution)

If we are dealing with a binary classification problem ($y \in \{0, 1\}$), we can model the output using a Bernoulli distribution. Here, our network needs to output a single parameter for each dimension: the probability $\mu$ that the output is 1.

Continuous Outputs (The Gaussian Distribution)

For regression tasks where $y$ is a real number, we naturally lean toward a Gaussian (Normal) distribution. Instead of predicting a single value, our network outputs two types of parameters:

The mean ($\mu$): Our most likely prediction.
The standard deviation ($\sigma$): Our uncertainty or risk.

Multi-Class Outputs (The Categorical Distribution)

When we scale up to multi-class classification (where an item belongs to one of $K$ classes), we use a Categorical distribution. We need the network to output a probability for each class, ensuring they all range between 0 and 1 and sum up to 1.

To achieve this, we use the Softmax function to transform raw, unbounded network outputs (logits) into a valid probability distribution:

$$\mu_k = \frac{\exp(w_k^T \phi(x) + b_k)}{\sum_{k'=1}^K \exp(w_{k'}^T \phi(x) + b_{k'})}$$

Why Predicting Distributions Matters

Predicting distribution is important, because single-point predictions often throw away valuable context. On a single point system, the model cannot predict how uncertain it response is, whereas a distribution naturally captures a model's uncertainty and confidence. For instance, in a continuous task using a Gaussian output, a massive variance indicates that the network is highly uncertain about its guess. Similarly, in a classification task, a flat probability split signals genuine ambiguity.

Consider the classic "Duck or Rabbit" optical illusion: if the model is forced to provide a single answer, a network must arbitrarily choose one and discard the alternative. By predicting a distribution instead, the network can output $p(\text{duck}|x) = 0.5$ and $p(\text{rabbit}|x) = 0.5$, accurately reflecting the natural ambiguity of the input.

Connecting Distributions to Loss Functions

We need to train the model to learn distribution, i.e, we want our predicted distribution $p(y|x)$ to be as close to the true data distribution $p_{\text{data}}(y|x)$. The standard mathematical tool to measure the distance between two distributions is Kullback-Leibler (KL) divergence.

Because we don't have access to the true perfect data distribution, we approximate this over our training set. Minimizing the KL divergence simplifies directly into minimizing the Negative Log-Likelihood (NLL) empirical cost function:

$$\tilde{C}(\theta) = \frac{1}{N} \sum_{n=1}^N - \log \tilde{p}(y_n|x_n)$$

This relationship bridges probability theory and practical deep learning. Depending on the distribution we choose, this objective function naturally manifests as the loss functions we use every day:

Cross-Entropy Loss: When we apply the NLL to a Bernoulli or Categorical distribution, we get the standard cross-entropy loss used in classification.
Mean Squared Error (MSE): If we apply the NLL to a Gaussian distribution and assume a constant variance ($\sigma = 1$), the cost function simplifies beautifully into the Euclidean distance between the true value $y$ and the predicted mean $\mu$.

Feature Extraction

When we build machine learning models, we rarely feed raw data directly into a linear predictor. Instead, we transform the raw input $x$ into a more meaningful representation, a feature vector $\phi(x)$.

For example, the feature vector $\phi(x)$ of house listing is;

$$\phi(x) = \begin{bmatrix} \text{Square Footage} \\ \text{Number of Bedrooms} \\ \text{Distance to Downtown (miles)} \\ \text{Age of House (years)} \end{bmatrix} = \begin{bmatrix} 1850 \\ 3 \\ 4.2 \\ 12 \end{bmatrix}$$

Historically, this step was known as manual feature extraction, or feature engineering. We relied on domain experts to create features by hand, such as extracting edges from images or counting specific keywords in text. This approach is not satisfactory because it doesn't work across different areas - a skilled expert in computer vision may not be able to help with speech recognition.

This brings us to a fundamental question: Can we automate feature extraction?

The breakthrough realization is that the feature extraction process $\phi(x)$ is simply a function. If we chain a prediction function $f$ with a feature extraction function $\phi$, the result is still a function: $$\hat{y} = f(\phi(x))$$

Thus we can actually turn the feature creator into a neural network itself by giving it its own adjustable weights and biases. By merging these new weights into our model's overall parameters $\theta$, the model gains the incredible ability to learn both the best way to represent the data and how to make predictions from it all at once. This creates a huge problem: a neural network can be shaped into almost any mathematical form, so there are countless ways to transform the input $x$.

Without constraints or structure, the model;

will get lost
overfit to random noise
become impossible to train.

To avoid this, we need to determine how to structure and restrict the neural network, such as organizing it into sequential layers, so it can learn meaningful patterns.

Instead of trying to find one massive, complex function, we build a stack of simple transformations. A single layer of this transformation can be as simple as an affine transformation (matrix multiplication plus a bias) followed by an element-wise nonlinearity: $$\phi_0(x) = g(W_0x + b_0)$$

Here, $W_0$ is the weight matrix, $b_0$ is the bias vector, and $g$ is a non-linear activation function like $\tanh$ or ReLU.

This structure reveals a fascinating mathematical property known as the Universal Approximation Theorem: if we project our low-dimensional input $x$ into a much higher-dimensional space by making the size of our hidden layer $\phi_0(x)$ significantly larger than $x$, the network can approximate almost any continuous function. This property still applies even if the hidden layer's weights $W_0$ and biases $b_0$ are completely random and remain unchanged from the start. Since these random hidden parameters do not change, the network no longer acts like a complex, non-linear deep learning puzzle. The problem simplifies to linear regression operating on the fixed, random features $\phi_0(x)$.

This means we don't need to use computationally expensive gradient descent, such as backpropagation, to train the hidden layers. Since the output layer is purely linear, you can skip iterative training and calculate the optimal output weights directly in one step using the Normal Equation: $$\theta = (X^T X)^{-1} X^T y$$

This architecture led to a huge and contentious debate in the machine learning community, known as Extreme Learning Machines (ELMs). Supporters praised ELMs for their fast training speeds and simple math, but critics said the idea was just a new name for older methods like Random Vector Functional Link networks and kernel methods. This made ELMs one of the most divisive topics in modern AI research.

Why is this highly debated? While randomly projecting data into a massive hidden layer enables incredibly fast training, it is wildly parameter-inefficient and fails to capture the data's underlying structure, requiring an exponentially larger network to match the accuracy of one that explicitly learns its features.

What We Actually Want

While random feature mapping is an interesting shortcut, it's not our goal. Random features don't take into account the specific characteristics of the dataset. What we really want is to learn representations - the ability to adjust every weight and parameter in the entire pipeline at the same time. When we optimize the feature extractor and the predictor together using gradient descent, the network learns to ignore noise and focus on the exact information needed for the task. This approach, where the entire system is optimized from start to finish, is the basis of modern Deep Learning.

Multilayer Perceptron (MLP)

If we stack a sequence of the simple feature-extraction layers we just discussed, we get a Multilayer Perceptron (MLP). The idea of stacking layers has been around since the 1960s, but it took a long time to develop a practical way to train all the layers at the same time. To understand how these networks are trained, we need to look at the math behind how gradients move backward through a network. This is the basis of the Backpropagation Algorithm.

Case Study: Single Hidden Unit (Scalar Input)

Let us start with a simple neural network: a single scalar input $x$, one hidden unit $\phi(x)$, and a single binary output prediction $\mu \in [0, 1]$ representing the probability that the true label $y = 1$.

The Forward Pass

Hidden Layer: We apply an affine transformation followed by a sigmoid activation function $\sigma$: $$\phi(x) = \sigma(ux + c)$$
Output Layer: We calculate the final prediction $\mu$ using a second sigmoid layer: $$\mu = \sigma(w\phi(x) + b)$$ Where the sigmoid function is defined as $\sigma(z) = \frac{1}{1 + \exp(-z)}$.

The Loss Function

Because we are dealing with a binary classification problem ($y \in \{0, 1\}$), we use Binary Cross-Entropy Loss ($C_x$): $$C_x = -y \log \mu - (1 - y) \log (1 - \mu)$$

If $y = 1$, the right term disappears, leaving $- \log \mu$ (we want $\mu$ close to 1).
If $y = 0$, the left term disappears, leaving $- \log (1 - \mu)$ (we want $\mu$ close to 0).

The Backward Pass (Applying the Chain Rule)

To train the network, we want to know how a tiny change in our weights ($w, u$) or biases ($b, c$) affects our overall loss $C_x$. We do this by working backward from the output to the input.

Step 1: The Output Layer Error

First, we see how the cost changes with respect to the pre-activation value at the output layer. By applying the chain rule: $$\frac{\partial C_x}{\partial \mu} \cdot \frac{\partial \mu}{\partial \text{pre-activation}} = \frac{\partial C_x}{\partial \mu} \cdot \mu(1-\mu)$$

When we compute this explicitly using the derivative of the sigmoid ($\mu' = \mu(1-\mu)$), a cancellation happens: $$\frac{\partial C_x}{\partial \mu} = \frac{\mu - y}{\mu(1-\mu)}$$ $$\text{Output Error} = \left(\frac{\mu - y}{\mu(1-\mu)}\right) \cdot \mu(1-\mu) = \mu - y$$

The Intuition: The initial error signal injected into the network is simply the difference between our prediction and the true label ($\mu - y$). If our prediction is perfect, the error is zero.

Step 2: Gradients for Output Parameters ($w$ and $b$)

Now we multiply this output error by the local derivatives of the output layer: $$\frac{\partial C_x}{\partial w} = (\mu - y)\phi(x)$$ $$\frac{\partial C_x}{\partial b} = (\mu - y)$$

Step 3: Gradients for Hidden Parameters ($u$ and $c$)

To update the hidden layer, we backpropagate the error further back. Notice how we reuse the output error $(\mu - y)$ instead of recomputing it: $$\frac{\partial C_x}{\partial u} = (\mu - y) \cdot w \cdot \phi(x)(1 - \phi(x)) \cdot x$$ $$\frac{\partial C_x}{\partial c} = (\mu - y) \cdot w \cdot \phi(x)(1 - \phi(x))$$

Generalizing: Multiple Hidden Units (Vector Inputs)

Real networks don't use single scalars; they handle vectors and matrices. Let's scale up the equations. Our input $\mathbf{x}$ is now a vector, our hidden layer features $\mathbf{\phi(x)}$ form a vector of size $l$, and the hidden weights become a matrix $U$. $$\mathbf{\phi(x)} = \sigma(U\mathbf{x} + \mathbf{c})$$ $$\mu = \sigma(\mathbf{w}^T \mathbf{\phi(x)} + b)$$

Vectorized Backpropagation

The output error remains the scalar value $(\mu - y)$. However, when we backpropagate this error to the hidden layers, we must maintain vector-matrix compatibility.

Output Weight Vector Gradient: $$\frac{\partial C_x}{\partial \mathbf{w}} = (\mu - y)\mathbf{\phi(x)}$$
The Error at the Hidden Layer: To push the error past the output weights and through the element-wise hidden activation function, we use the element-wise multiplication ($\odot$) with the derivative of the hidden activations ($\mathbf{\phi}'$): $$\frac{\partial C_x}{\partial \mathbf{\phi}} = (\mu - y)\mathbf{w} \odot \mathbf{\phi}'(\mathbf{x})$$
Hidden Weight Matrix Gradient ($U$): Finally, by utilizing the vector outer product ($\mathbf{a}\mathbf{b}^T$), we calculate the gradient for the entire hidden weight matrix simultaneously: $$\frac{\partial C_x}{\partial U} = \left[ (\mu - y) (\mathbf{w} \odot \mathbf{\phi}'(\mathbf{x})) \right] \mathbf{x}^T$$

[!NOTE] Computational Efficiency By working backward from the output layer to the input layer, we compute intermediate error terms, such as $\mu - y$, exactly once and store them. The lower layers then use these stored values. Since we share these intermediate results, we can calculate the gradient for millions of parameters efficiently. This process takes no more time than running the data forward through the network.

How Many Layers and How Wide Should They Be?

When we decide on the number of layers and the number of hidden units per layer in a Multilayer Perceptron, we have to balance the network's ability to learn complex patterns against the risk of overfitting and the amount of computation required.

Width Heuristic: A common approach is to start with a hidden layer size between the input and output layer sizes. For efficiency, people often choose sizes that are powers of 2, such as 64, 128, or 256. Another option is to make the first hidden layer larger than the input layer, which can help by transforming the data into a higher-dimensional space where it may be easier to separate.
Depth Heuristic: For standard tabular data, a network with 1 to 3 hidden layers is usually sufficient. As the network gets deeper, it builds a hierarchy of features. The first layer learns simple edges or boundaries, the second layer combines these into shapes, and the deeper layers understand more complex concepts.
The Golden Rule: We always prefer to start small and incrementally add width and depth until we notice our training loss decreases but our validation loss begins to worsen (the definition of overfitting).

Automating Backpropagation

The cleanest way to conceptualize backpropagation is simply as a clever, automated application of the chain rule of calculus over a Directed Acyclic Graph (DAG). We call this DAG a computational graph.

The Anatomy of a Computational Graph

In a computational graph, we break down our complex neural network equations into an interconnected web of basic operations. The graph consists of two main elements:

Variable Nodes: Represent data ($x, y$) or learnable parameters (weights $W$, biases $b$).
Function Nodes: Represent modular mathematical operations—such as Matrix Multiplication (MatMul), Matrix Addition (MatSum), Element-wise Sigmoid ($\sigma$), or the Cost/Loss function ($C_y$).

The Contract of a Function Node

Every single function node in the graph acts as an independent black box that complies with a two-part contract:

The Forward Pass ($\to$): It takes its inputs, computes the specified mathematical operation, and passes the result to the next node.
The Backward Pass ($\leftarrow$): It receives a downstream gradient from the node ahead of it. It multiplies that incoming gradient by its own local derivative (evaluated at the cached forward pass values), and passes the resulting gradient backward to its input variables.

For example, when the error flows back to a MatSum node ($A + B$), its local derivative with respect to both inputs is 1. It mechanically multiplies the incoming gradient by 1 and splits it equally back to both paths.

Because each node only needs to know its own local derivative, the backend framework (like PyTorch or TensorFlow) can automatically calculate gradients for any arbitrary network architecture we stitch together. This technique is known as Reverse-Mode Automatic Differentiation.

Interview Angle (Next-Level Graph Intuition): If you push this concept further, a node doesn't have to return a numerical derivative during the backward pass. Instead, it can return a computational subgraph that represents its derivative. Because the gradient itself becomes a graph, we can run a forward pass on that gradient graph to effortlessly compute higher-order derivatives (like Hessians for second-order optimization).

Differentiability: What Happens When Functions Break?

Since backpropagation relies entirely on the chain rule, every function node in our graph must be differentiable. But what happens with the non-linearities we use every day?

Smooth Functions (Sigmoid and Tanh)

Sigmoid ($\sigma$) and Hyperbolic Tangent ($\tanh$) are perfectly smooth and differentiable everywhere. Their derivatives are simple functions of their outputs:

$\sigma'(x) = \sigma(x)(1 - \sigma(x))$

While mathematically convenient, these functions suffer from the vanishing gradient problem: as the input gets very large or small, the curve flattens out, the local derivative drops to nearly 0, and the error signal dies.

Piece-wise Linear Functions (ReLU and Maxout)

To fix vanishing gradients, modern deep learning relies heavily on the Rectified Linear Unit (ReLU): $$\text{rect}(x) = \max(0, x)$$

Mathematically, ReLU is not strictly differentiable because it has a sharp corner at exactly $x = 0$. How do we get around this? In practice, the probability of a continuous input landing exactly at $0.000000$ is effectively zero. We can safely bypass this by using a subgradient: $$\text{rect}'(x) = \begin{cases} 1 & \text{if } x > 0 \\ 0 & \text{if } x \le 0 \end{cases}$$

The "Dying ReLU" Problem: Notice that for half of the input space ($x \le 0$), the derivative is precisely zero. If a neuron gets stuck in this region, it stops propagating gradients entirely—a phenomenon known as the dead neuron problem.
The Maxout Alternative: To bypass zero-derivative zones, we can use a Maxout unit. Maxout takes the maximum over $k$ arbitrary inputs ($\max(x_1, \dots, x_k)$). Its derivative is $1$ for whichever input happens to be the maximum, and $0$ for the rest. This ensures that a gradient is always actively backpropagating down at least one path.

The Ultimate Wall: Stochastic Variables

What if we purposefully introduce a non-differentiable step? Consider a binary stochastic node: it takes a probability, flips a coin, and outputs a hard $0$ or $1$.

Because a step-function or random sample contains no smooth slope, its mathematical derivative is zero everywhere it is defined, making standard backpropagation impossible. When we absolutely require random sampling inside our networks (like in Variational Autoencoders or Reinforcement Learning), we cannot use standard reverse-mode auto-diff directly. Instead, we must rely on advanced workarounds like the reparameterization trick or score-function estimators (REINFORCE) to bypass the non-differentiable roadblock.

What Other Transformation Layers Are There?

While the standard MLP relies entirely on fully connected (dense) layers followed by an activation function, modern deep learning introduces specialized architectures designed to exploit specific structures in data:

Convolutional Layers (CNNs): Instead of connecting every input pixel to every neuron, convolutional layers apply small, localized filters. We use these for spatial data, like images or audio spectrograms, because they drastically reduce the parameter count and achieve translation invariance.
Recurrent/Attention Layers (RNNs, LSTMs, Transformers): When our data has a sequential or temporal nature (like text, time-series, or DNA sequences), we use layers that can track history and maintain a memory state across time steps.
Normalization Layers (Batch Norm, Layer Norm): These act as internal stabilizing transformations. They scale the activations of a layer to have a mean of 0 and a variance of 1, which prevents gradients from exploding or vanishing and allows us to train much deeper networks.

Selecting and Scheduling the Learning Rate ($\eta$)

The learning rate is arguably the most critical hyperparameter we tune. If it is too large, our gradient updates will overshoot the minimum and cause training to diverge. If it is too small, training will take an impractical amount of time or get trapped in poor local minima.

What kind of learning rate should we start with? A baseline default like 1e-3 ($0.001$) or 3e-4 is standard, particularly when using adaptive optimizers like Adam. Adam automatically scales the learning rate for each individual parameter based on its running history of gradients, removing a lot of the manual guesswork.
How should we schedule it over training? As our model gets closer to the optimal solution, large steps will cause it to bounce back and forth around the minimum. To fix this, we implement a Learning Rate Scheduler to decay $\eta$ over time:
- Step Decay: We reduce the learning rate by a fixed factor (e.g., multiplying by $0.1$) every few epochs.
- Cosine Annealing: We smoothly lower the learning rate following a cosine curve, gently slowing down our updates as training progresses.
- Reduce on Plateau: We monitor the validation loss, and if it fails to improve for a set number of epochs (a plateau), we automatically drop the learning rate to give the network a finer resolution to find a breakthrough.