9.1 Affine Layers, Activation, and Logits: Depth Comes from Function Composition, Not from the Name "Neuron"

A linear classifier can accept all 784 pixels of a 28×28 image and output ten class logits. Its limitation is the affine score function: pairwise class boundaries are linear, and pixel locality or translation structure is not built into the model.

A neural network composes learnable affine transformations with nonlinear functions. The loss adjusts their parameters together through a computation graph; a neuron-shaped diagram is optional.

This lesson's objectives

• Clearly write the affine layer and batch computation using tensor shapes;
• Explain why a sequence of linear layers degrades to a single linear layer;
• Look at hidden representations with XOR, rather than memorize "one layer deeper" alone;
• Distinguish logits, probabilities, predictions, and loss;
• Understand the capacity, inductive bias, and applicable boundaries of MLPs.

1. A Layer is a batch affine transformation

For a single input $x\in {\mathbb{R}}^{d_{in}}$:

$$z=Wx+b,$$

where $W\in {\mathbb{R}}^{d_{\text{out}}\times d_{\text{in}}}$ and $b\in {\mathbb{R}}^{d_{\text{out}}}$.

For a batch $X\in {\mathbb{R}}^{B\times d_{in}}$, samples are typically stored row-wise:

$$Z=X{W}^T+b,Z\in {\mathbb{R}}^{B\times d_{out}}.$$

PyTorch's nn.Linear(in_features, out_features) is exactly $y=x{W}^T+b$; weight.shape is [out_features, in_features], not the reverse. Parameters:

$$d_{out}{d}_{in}+d_{out}.$$

This runnable Python 3 / PyTorch example checks the batch and parameter shapes; it completes without output when the assertions pass:

import torch
from torch import nn

layer = nn.Linear(in_features=784, out_features=128)
assert layer.weight.shape == (128, 784)
assert layer.bias.shape == (128,)

x = torch.randn(32, 784)
z = layer(x)
assert z.shape == (32, 128)

Broadcasting adds the same bias to each row of a batch. Understanding shape is more important than visualizing each scalar as a circle: convolution, attention, and parallel training all ultimately depend on tensor dimension contracts.

2. Stacking just linear layers won't increase expressive power

Two affine layers:

$$h=W_{1}x+b_1,y=W_{2}h+b_2,$$

Can be combined as:

$$y=(W_2{W}_1)x+(W_2{b}_1+b_2).$$

No matter how many layers you stack, as long as there's no nonlinearity in the middle, the entire system remains an affine map. Depth only increases parameterization; it doesn't make the decision boundary nonlinear.

A narrow hidden layer can additionally limit the rank of the combined matrix. Factorization changes optimization and capacity constraints even while the resulting map remains affine.

Activation function $\varphi$ breaks the merge:

$$h=\varphi (W_{1}x+b_1),y=W_{2}h+b_2.$$

3. XOR demonstrates a transformation representation

Binary XOR:

(0, 0) -> 0
(0, 1) -> 1
(1, 0) -> 1
(1, 1) -> 0

The original input space contains no straight line that can separate the two classes. Let $ s = x_1 + x_2 $, and construct two ReLU hidden features:

$$h_1=\mathrm{ReLU}(s),h_2=\mathrm{ReLU}(s-1).$$

Then at the four binary points:

$$f(x)=h_1-2h_2$$

This gives exactly the 0/1 output of XOR. The network isn't mysteriously "learning equality and inequality," but rather mapping inputs to a piecewise-linear representation that can be combined by the output layer. Whether training actually finds this solution still depends on the data, initialization, loss function, and optimization.

4. What Common Activation Functions Are Doing

ReLU

$$\mathrm{ReLU}(z)=max(0,z).$$

On the positive axis, the derivative is 1, making MLPs piecewise linear and easy to compute. On the negative axis, the gradient is 0, which might cause certain units to remain inactive for long periods; however, a derivative of 1 on the positive side does not guarantee that gradient vanishing won't occur in deep networks, since the gradient still needs to be continuously multiplied through the weight matrix/Jacobian.

Leaky ReLU / PReLU

Keeps a small slope on the negative axis, reducing the risk of permanent zero gradient, but does not automatically resolve issues with bad initialization, learning rate, or activation distribution.

Sigmoid / Tanh

The sigmoid maps values to (0,1), while tanh maps to (-1,1). In regions of large absolute values, the derivatives approach zero, making them prone to saturation during backpropagation in deep networks. They are still suitable for gates, probabilistic linking, or bounded states, and are not equivalent to being "obsolete."

GELU / SiLU

Smooth, non-pure hard gates, common in modern architectures. Which is better depends on the model, initialization, normalization, and computational power; "Transformer standard activation" is not a fixed fact across all implementations.

Nondifferentiable points (like ReLU at 0) are handled by the framework choosing a subgradient or convention; almost everywhere differentiability is sufficient for gradient-based optimization.

5. The MLP Outputs Logits

The final layer of a typical classification model usually outputs $K=10$ raw scores:

$$z=(z_1,\dots ,z_K).$$

These logits are not probabilities. Softmax:

$$p_k=\frac{e^{z_k}}{\sum _j{e}^{z_j}}$$

Softmax normalizes them into probabilities over mutually exclusive categories. During training, do not manually softmax and then pass the result to PyTorch CrossEntropyLoss; the loss directly receives logits and combines log-softmax with negative log-likelihood in a numerically stable way.

The next example reuses the torch and nn imports above. It generates synthetic inputs and labels to check the loss interface, without training a useful classifier:

class MLP(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=128, classes=10):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, classes),
        )

    def forward(self, x):
        return self.network(x)  # logits, shape [batch, classes]

model = MLP()
x = torch.randn(32, 784)
target = torch.randint(0, 10, (32,), dtype=torch.long)

logits = model(x)
loss = nn.CrossEntropyLoss()(logits, target)
probability = logits.softmax(dim=1)  # For explanation/Evaluation, do not return to the same loss
prediction = logits.argmax(dim=1)

CrossEntropyLoss also supports soft/probability targets modes, with different shape and dtype rules. First, clearly define the label semantics, don't rely on broadcasting to "just work."

6. Don't Mix Up Binary, Multiclass, and Multilabel Classification

Single-label binary

It outputs one logit [B] or [B,1] and uses BCEWithLogitsLoss with float targets 0 or 1. The loss combines sigmoid with BCE for better numerical stability.

Single-label multiclass

Output [B,K] logits, paired with CrossEntropyLoss and long class indices of shape [B]. Classes are mutually exclusive.

Multilabel

Output [B,K] separate logits, paired with BCEWithLogitsLoss and [B,K] multi-hot/soft targets. Apply a separate sigmoid per class, without using mutually exclusive softmax.

Separate sigmoid outputs do not establish statistical independence between labels: they can share a learned representation. For BCE, targets must match the logit shape exactly, so [B] and [B,1] must not be mixed.

pos_weight, class weights alter the training objective and do not guarantee that the original overall probability remains calibrated; the calibration/threshold principles from Chapter 8 still apply.

7. Width, Depth, and Piecewise-Linear Regions

ReLU MLPs create complex regions through multilayer piecewise linear mappings. Increasing width adds more features per layer, depth reuses representations from earlier layers, and some functions can be expressed with fewer units.

The Universal Approximation Theorem only guarantees the existence of an approximation under certain conditions and sufficient capacity; it does not guarantee:

• Limited data can identify the objective function;
• SGD can find that parameter;
• Network/samples/computational requirements are acceptable;
• Still correct outside the distribution;
• Predictions are calibrated or have a causal interpretation.

Expressiveness, optimizability, and generalization are three distinct issues.

8. MLP Doesn't Automatically Understand Structure

After flattening a 28×28 image into 784 dimensions, a standard MLP doesn't know that adjacent pixels are more related than distant ones and lacks translation equivariance. It can learn some patterns from the data, but may have poor parameter and sample efficiency.

An architecture is an inductive bias:

• CNN favors local connections and weight sharing;
• RNN/state-space models favor sequential recursion;
• attention enables content-relevant cross-position interactions;
• Graph networks lean toward node-edge relationships.

You can't ignore data structure just because MLPs theoretically can fit it, the bias will be expanded on in the next chapter.

9. Input, Output, and Unit Boundaries

Neural networks still need data contracts:

• The scale of numerical inputs affects initialization and optimization;
• Categories need embedding/encoding, and ID numeric sizes typically have no ordinal meaning;
• Missing values should be explicitly represented;
• The target transform and inverse transform must be consistent;
• Padding/mask should not be mistaken for real tokens;
• Training and deployment normalization statistics must be frozen.

The network won't automatically fix leaks, mislabeled data, or time travel errors.

Common Misconceptions

• Linear models cannot handle 784-pixel inputs: They can; the limitations are the decision boundaries and inductive bias, not the input dimension itself.
• Adding one more layer can teach you XOR: Finding a solution doesn't mean optimization has definitely found it.
• ReLU addresses gradient vanishing: It alleviates some saturation issues, but the Jacobian chain may still vanish or explode.
• Softmax must be applied inside the model before using CrossEntropyLoss: during training, pass raw logits.
• Neural networks automatically discover local structures: Architecture and data augmentation determine many priors.

Practice

1. Write out the weight/bias/output shape and number of parameters for Linear(12, 5).
2. Manually calculate the four XOR points using two ReLU hidden units.
3. Write out the logits/target shapes and loss for binary, multiclass, and multilabel classification.
4. Remove the ReLU activation from the middle of an MLP, algebraically merge the two layers, and verify that the output remains equivalent.
5. What prior assumptions are needed to compare a flattened MLP with a model that has local weight sharing?

Summary

Neural networks are compositions of differentiable functions. An affine layer transforms coordinates, activation functions introduce nonlinearity, and the final layer produces logits matching the task loss. Hidden representations expand expressive power, but do not guarantee trainability, generalization, or understanding of data structure.

Next, we traverse the computation graph backward: backpropagation isn't vague "spreading responsibility," it's efficient reuse of the chain rule via reverse-mode automatic differentiation.