Multi-Layer Perceptron — Forward Pass

Why this session exists

A multi-layer perceptron alternates affine maps and nonlinear activations. Understanding its matrix form lets you locate parameters, debug shapes, reason about numerical stability and recognize position-wise feed-forward blocks inside larger architectures. Calling a framework layer is useful; knowing what it computes is what lets you check it.

You will implement a two-hidden-layer NumPy classifier, trace every shape, count parameters exactly, distinguish universal approximation from learnability, and compare width/depth under a validation protocol. The implementation accepts MNIST's 784 inputs, but the runnable default uses sklearn's bundled 8×8 digits so it requires no dataset download. No 96% MNIST result is promised or inferred from that different dataset.

Prerequisites

Use S096 activations, matrix multiplication, broadcasting and S095 validation discipline. Each row is an example throughout; PyTorch's stored linear weights use the transposed convention, discussed in S100.

(a) Intuition · 5 min

A committee of committees first constructs features, then combines those features into a decision. A hidden layer similarly learns intermediate coordinates rather than receiving a human label for each unit. The analogy does not imply that successive MLP layers reliably become “edges, textures, eyes, cats,” or that each neuron has a unique human-interpretable role.

Each affine map can rotate, scale, shear or project the representation; the activation bends it. Composing these changes can make classes easier to separate at the output. Width controls available intermediate coordinates, while depth controls composition. This is a useful picture, not a guarantee that every trained layer improves separability.

The forward pass in one paragraph

For a batch $X\in\mathbb R^{B\times D}$,

The weights have shapes $(D,H_1)$, $(H_1,H_2)$ and $(H_2,K)$. Biases have widths $H_1,H_2,K$; each broadcasts across rows. All intermediate activations retain batch size $B$. A bias shaped (1,H) also broadcasts correctly; a single example can be represented as (1,D). Not every shape error is caused by a one-dimensional bias: check axes, transposes and flattening explicitly.

For exclusive-class classification, $Z_3$ contains logits. Softmax converts them to probabilities when needed; stable cross entropy should operate on logits. For unbounded regression the output may remain affine; constrained targets may need another output parameterization.

The universal approximation theorem — and its practical trap

A classical result for continuous sigmoidal activations states that finite sums $\sum_{j=1}^m a_j\sigma(w_j^Tx+b_j)$ can uniformly approximate continuous functions on the unit cube to arbitrary accuracy. This is an existence statement with sufficient width, not a promise about the weights found by SGD, finite-data generalization or a fixed network size. ReLU is unbounded, so do not justify its approximation results by calling it a bounded sigmoid.

Depth can represent some compositional function families efficiently; it does not follow that shallow networks are obsolete or always memorize. Nor does “arbitrary continuous function” include an exact discontinuous classifier everywhere under uniform error. Approximation assumptions, target domain and error criterion must be stated before interpreting universality.

An MLP as a stack of matrix multiplications

The biases are part of the affine maps, not absent operations between the arrows. On a tiny fixture $X=(1,2)$, $W_1=[[1,-1],[2,1]]$, $b_1=(0,1)$, the first pre-activation is $(5,2)$ and ReLU leaves it unchanged. With output weights $[[1,0],[-1,2]]$ and zero bias, logits are $(3,4)$ and the second-class probability is $e^4/(e^3+e^4)\approx0.7311$. You can trace larger matrices using exactly this row-wise calculation.

Parameter counting — where the weights live

A dense layer costs $d_{in}d_{out}+d_{out}$ trainable scalars. Therefore [784,128,64,10] contains

The first layer dominates this particular model, not every MLP. [100,64,10] has 7,114 parameters. [784,32,10] has 25,450 and [784,512,256,10] has 535,818. These counts do not include activations, gradients or optimizer state. Doubling a single hidden width affects its adjacent matrices linearly if neighbors stay fixed; doubling both sides of an interior matrix quadruples that matrix's weights.

Width vs depth — which one buys you what

A wider layer offers more simultaneous features, while extra layers offer more composition. A narrow bottleneck deliberately discards information in an autoencoder or compressed representation; it is not inherently a mistake. A shallow model can be a strong baseline. Residuals and normalization can improve trainability of deep designs, but no universal “ten layers” boundary requires them.

At equal parameter count, the models may still differ in activation memory, sequential dependencies and arithmetic. Evaluate actual task quality and resource constraints; parameter count alone is neither a latency measurement nor a sample-complexity theorem. Compare a modest MLP with a well-tuned booster on heterogeneous tabular data. The cited tabular benchmark reports tree strength in its studied regime, not a rule that neural networks require millions of rows to win.

Initialization: break symmetry and control scale

If hidden units have identical incoming parameters and equivalent outgoing connections, deterministic updates can preserve their symmetry. Random initialization is a practical symmetry breaker; it is not the only possible one. Zeroing only the first ReLU layer with zero bias creates an additional issue: the chosen derivative at zero blocks its data gradient. That is different from claiming any zeroed matrix always limits accuracy to 10%.

Under independent zero-mean weights, pre-activation variance is approximately $n\sigma_w^2 E[a^2]$. For symmetric zero-mean $z$, ReLU satisfies $E[\operatorname{ReLU}(z)^2]=E[z^2]/2$. This is a second-moment relation, not exactly halving variance: ReLU's mean is positive. He variance $2/n$ compensates approximately for this moment reduction. Xavier variance $2/(n_{in}+n_{out})$ balances fan-in/fan-out considerations, with activation-dependent gain. These initial estimates do not guarantee signal scales remain stable throughout training.

The MLP inside a transformer

The original Transformer defines a position-wise FFN

It shares these weights across token positions. Attention mixes information between positions; the FFN transforms each position's features. The original base configuration uses $d=512$ and inner width 2048. Ignoring biases, a $d\to4d\to d$ FFN has $8d^2$ weights; four $d\times d$ attention projections contribute $4d^2$. The FFN is thus two-thirds of those matrices, not necessarily two-thirds of every full model after embeddings, gating, experts and other choices.

Both residual edges are shown. This is the original post-norm pattern, not a template for every transformer. A gated FFN has three projections, with about $3dh$ weights; choosing $h\approx8d/3$ matches the ordinary $8d^2$ budget before rounding. A 1×1 convolution similarly mixes channels at each spatial position with shared weights, but not every ResNet block has one. Recommendation towers may use MLP feature interactions, while dot-product embedding baselines remain important.

(c) Hands-on · 25 min

The lab keeps the complete forward/backward/SGD mechanism in NumPy. It uses stable logit cross entropy rather than log(probability + epsilon) with a mismatched derivative, records validation loss during training and restores the chosen checkpoint before test evaluation. The arbitrary layer list supports two hidden layers and other widths.

import numpy as np
from scipy.special import logsumexp
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
 
class MLP:
    def __init__(self, sizes, seed=0):
        rng=np.random.default_rng(seed)
        self.W=[rng.normal(0,np.sqrt(2/a),(a,b)) for a,b in zip(sizes[:-1],sizes[1:])]
        self.b=[np.zeros(b) for b in sizes[1:]]
    def forward(self,X):
        a=X; cache={'a':[X], 'z':[]}
        for i,(W,b) in enumerate(zip(self.W,self.b)):
            z=a@W+b; cache['z'].append(z)
            a=np.maximum(z,0) if i<len(self.W)-1 else z
            cache['a'].append(a)
        return a,cache
    def loss_and_grads(self,X,y):
        logits,c=self.forward(X); n=len(y)
        logp=logits-logsumexp(logits,axis=1,keepdims=True)
        loss=-logp[np.arange(n),y].mean()
        dz=np.exp(logp); dz[np.arange(n),y]-=1; dz/=n
        dW=[None]*len(self.W); db=[None]*len(self.b)
        for i in reversed(range(len(self.W))):
            dW[i]=c['a'][i].T@dz; db[i]=dz.sum(0)
            assert dW[i].shape==self.W[i].shape and db[i].shape==self.b[i].shape
            if i:
                dz=(dz@self.W[i].T)*(c['z'][i-1]>0)
        return float(loss),dW,db
    def step(self,dW,db,lr):
        for i in range(len(self.W)):
            self.W[i]-=lr*dW[i]; self.b[i]-=lr*db[i]
    def state(self):
        return ([w.copy() for w in self.W],[b.copy() for b in self.b])
    def restore(self,state):
        self.W=[w.copy() for w in state[0]]; self.b=[b.copy() for b in state[1]]
 
def evaluate(m,X,y):
    z,_=m.forward(X)
    return float(np.mean(logsumexp(z,axis=1)-z[np.arange(len(y)),y])), float(np.mean(z.argmax(1)==y))
 
X,y=load_digits(return_X_y=True); X=X.astype(float)/16.  # fixed known pixel range
ids=np.arange(len(y)); dev,test=train_test_split(ids,test_size=.2,stratify=y,random_state=1)
train,val=train_test_split(dev,test_size=.25,stratify=y[dev],random_state=2)
assert not (set(train)&set(val) or set(dev)&set(test))
m=MLP([64,32,16,10]); rng=np.random.default_rng(3)
initial=evaluate(m,X[train],y[train])[0]; best=float('inf'); chosen=None
for epoch in range(1,26):
    rows=rng.permutation(train)
    for start in range(0,len(rows),64):
        b=rows[start:start+64]; loss,dW,db=m.loss_and_grads(X[b],y[b])
        assert np.isfinite(loss) and all(np.isfinite(g).all() for g in dW+db)
        m.step(dW,db,.1)
    vl,va=evaluate(m,X[val],y[val])
    if vl<best:
        best=vl; chosen=m.state(); chosen_epoch=epoch
    if epoch in [1,5,25]:print('epoch',epoch,'validation loss/accuracy',vl,va)
m.restore(chosen)
assert evaluate(m,X[train],y[train])[0]<initial
print('chosen epoch',chosen_epoch,'train',evaluate(m,X[train],y[train]),'test',evaluate(m,X[test],y[test]))
for sizes,total in [([784,32,10],25450),([784,128,64,10],109386),([784,512,256,10],535818),([100,64,10],7114)]:
    net=MLP(sizes); assert sum(w.size+b.size for w,b in zip(net.W,net.b))==total
    assert net.forward(np.zeros((2,sizes[0])))[0].shape==(2,sizes[-1])
# Stable loss remains finite even when explicit softmax probabilities underflow.
z=np.array([[1000.,-1000.]])
assert np.isclose(logsumexp(z,axis=1)[0]-z[0,1],2000.)

Anatomy of the script

forward stores the input to each affine layer plus pre-activations needed by ReLU's backward rule. The output remains logits. loss_and_grads computes the gradient for an unweighted mean over examples; it divides by the actual last-batch size, not a configured nominal batch. Backprop uses the old weights for all derivatives before step changes any parameter.

At the output, $\Delta_3=(P-Y)/B$. Then $dW_3=A_2^T\Delta_3$, $db_3=\sum_b\Delta_3$, $\Delta_2=(\Delta_3W_3^T)\odot1[Z_2>0]$, and similarly $\Delta_1=(\Delta_2W_2^T)\odot1[Z_1>0]$, $dW_1=X^T\Delta_1$. Shapes and the mean-loss normalization are part of the derivation, not afterthoughts.

For a separate optional MNIST experiment, supply trusted official training/test arrays, flatten each 28×28 image to 784 features, divide pixels by 255, split validation only from the official training partition, and instantiate [784,128,64,10]. Do not random-mix official test images into training. Keep the same training functions, select checkpoints on validation and use official test once after choices. This data run was not performed; the 784-dimensional forward and parameter contracts were tested locally.

Width practice with an answer

Exercise: Compare [784,32,10], [784,128,64,10] and [784,512,256,10] under a declared budget. What should the larger model achieve?

Answer: The exact counts are 25,450, 109,386 and 535,818. Equal epochs do not mean equal compute. Use identical splits, repeated seeds, validation selection and a recorded optimizer/step budget. Larger capacity can help, hurt or leave validation quality unchanged. No fixed 90/96/97% accuracies or 2× timing follows from these sizes. On the bundled digits dataset replace 784 by 64 and recompute the counts rather than labeling it MNIST.

(d) Production reality · 15 min

Tabular baselines: Compare a tuned MLP, a simple linear model and a booster under the same feature availability and validation boundaries. Published benchmarks are evidence for their datasets, not a categorical ban on neural networks for structured data.

Scaling arithmetic: The Chinchilla paper reports a compute-efficient regime in which model size and training tokens grow together, and a 70B model trained on four times Gopher's data. Its abstract states that compute-optimal model size and training-token count should scale equally: doubling one calls for doubling the other in the studied regime. That does not support the earlier reversed claim “optimal parameters ≈ 20 × training tokens.” It is an empirical compute-allocation result, not a universal deployment law or a prescription for every later model's width/depth. Inference cost and data reuse can motivate different choices.

Recommendation: A nonlinear MLP over user/item embeddings should earn its cost against a tuned dot product. The cited matrix-factorization comparison supports retaining that baseline, not a universal claim about every company's online A/B result. Offline loss and the operational target may differ.

Where this shows up in the rest of the plan

S098 derives and checks gradients; S099 changes the update without changing the forward map; S100 ports the matrices into registered modules; S101 adds explicit train/eval behavior. Convolutions share weights spatially and attention mixes token positions, so both add structure beyond an ordinary dense stack.

(e) Recall + stretch · 10 min

1. All shapes? $X:(B,D)$; $W_1:(D,H_1)$; $A_1:(B,H_1)$; $W_2:(H_1,H_2)$; $A_2:(B,H_2)$; $W_3:(H_2,K)$; logits $(B,K)$, with one bias per output feature.
2. UAT limitation? Existence under an activation/domain/error contract is neither an efficient-size, optimization nor generalization guarantee.
3. MNIST architecture count? 109,386, including all three biases.
4. Why bias? It translates affine surfaces instead of restricting every threshold to the origin.
5. Zero initialization? Symmetry can tie units together; zero ReLU paths can additionally block gradients. Inspect the whole architecture, not just one weight matrix.
6. CE logit gradient? $(P-Y)/B$ for normalized targets and unweighted batch mean; weighted/masked reductions need the correct denominator.
7. Why fourfold FFN width? It is one documented expansion design, not a theorem. Gated blocks have a different projection count.
8. Tabular MLP or booster? Compare controlled baselines using the actual types, missingness, embeddings and transfer needs; row count alone does not choose.
9. What does batching change? It replaces many row matrix products with one larger matrix product and can amortize overhead. No universal hardware speedup is asserted.
10. Stretch derivative: Walk $\Delta_3\to\Delta_2\to\Delta_1$ with ReLU masks and unchanged weights; then $dW_1=X^T\Delta_1$. Divide once for the mean, and check every bias too.

Explain-out-loud test

Draw every shape, locate the largest adjacent layer product, explain learned coordinates without semantic guarantees, and distinguish a position-wise FFN from token mixing.

What comes next

Backpropagation tests the exact gradients rather than treating decreasing loss as proof.

Sources and execution boundary

Sources: Cybenko's 1989 Approximation by superpositions of a sigmoidal function; PyTorch initialization definitions; original Transformer; tabular benchmark; Chinchilla; matrix-factorization comparison. Exact passages and artifact hashes accompany the review. Executed results are the bundled-digits CPU program and mathematical/shape fixtures; no MNIST, GPU or large-model benchmark was run.