The Forward Pass: Shapes, XOR and Parameter Counts

0 · From one affine decision to a learned representation

A single affine score cannot separate XOR. Adding another affine map does not help: (X @ W1 + b1) @ W2 + b2 equals X @ (W1 @ W2) + b1 @ W2 + b2. The effective linear part also satisfies rank(W1@W2) ≤ min(rank(W1),rank(W2)), so a narrow bottleneck can restrict the affine maps representable by the stack. A hidden nonlinearity changes the function class. This lesson makes that change visible in nine numbers, then constructs a classifier for all four XOR inputs.

Rumelhart, Hinton and Williams's 1986 Nature paper, pages 533–536, describes learning internal representations by back-propagating errors. It is not an eight-page proof that one paper ended an AI winter; neither a citation count nor an invented floppy-disk anecdote explains the mechanism. Forward evaluation, not historical priority, is our subject today. The original two-linear-layer transformer FFN has related algebra, but not every sequence architecture keeps that exact block.

1 · What you'll build today

Draw and evaluate a 2-2-1 MLP, batch it without changing its parameters, construct a 2-3-1 XOR classifier, and count parameters from explicit dimensions. You will distinguish logits passed to fused losses from probabilities used to interpret predictions. Finally, compare width, depth and modern gated/recurrent blocks without turning one architecture into a universal rule.

2 · Checkpoint · you should already know

You need matrix multiplication, trailing-axis broadcasting, a single neuron's affine score, and ReLU/sigmoid. Reconstruct these from geometry, broadcasting, the neuron and activations. Optional intuition: 3Blue1Brown, StatQuest, Nielsen chapter 1, and Olah's geometry. These are pre-reading, not evidence for an unmeasured training result.

3 · The 2-2-1 architecture — draw it once, forever

X (N,2) -- W1 (2,2), b1 (2,) --> Z1 (N,2)
Z1 -- ReLU --> H (N,2)
H -- W2 (2,1), b2 (1,) --> z (N,1) -- sigmoid --> p (N,1)

Four first-layer weights, two hidden biases, two output weights and one output bias give nine parameters, independent of batch size. A batch is many inputs evaluated with the same parameters, not many separately parameterized networks.

3.1 The forward equations (row-batch convention)

We store weights as (in,out). Column-vector notation instead writes h = ReLU(W_col @ x_col + b_col), with W_col = W_row.T. PyTorch nn.Linear accepts row batches but stores (out,in) and computes X @ weight.T + bias. A row-batch API does not imply our storage convention. For a three-dimensional input (B,T,D), matmul contracts the last feature axis while preserving both leading axes.

A hidden unit changes branch at w·x+b=0. Its weight controls the hyperplane's orientation and the bias translates it away from the origin. Removing every bias forces the ReLU kink hyperplanes through the origin and restricts the resulting homogeneous network; it is not automatically fatal in architectures with other offsets or suitable tasks. Parameter counting must follow the actual bias configuration.

4 · One forward pass — by hand, on numbers you can hold in your head

Take x=[1,2], W1=[[.1,.2],[.3,.4]], b1=[.1,.1], W2=[[.5],[.6]], b2=[.2].

4.1 Layer 1

The two dot products are .1+2*.3=.7 and .2+2*.4=1. Add the bias to obtain Z1=[.8,1.1]. Both are positive, so H=[.8,1.1].

4.2 Layer 2

The output logit is .8*.5+1.1*.6+.2=1.26. Thus sigmoid(1.26)=0.7790261078, or 0.7790 to four decimal places. The former 0.7788 and printed 0.77894181 were not the sigmoid of 1.26. This probability is a model output, not a calibration guarantee.

4.3 Verify in NumPy

import numpy as np
x = np.array([1.0, 2.0])
W1 = np.array([[0.1, 0.2], [0.3, 0.4]])
b1 = np.array([0.1, 0.1])
W2 = np.array([[0.5], [0.6]])
b2 = np.array([0.2])
h_pre = x @ W1 + b1
h = np.maximum(0, h_pre)
y_pre = h @ W2 + b2
y = 1 / (1 + np.exp(-y_pre))
assert np.allclose(h, [.8, 1.1])
assert np.allclose(y_pre, [1.26])
assert np.allclose(y, [.7790261077798122])
print('hand trace', h, y_pre, y)

The sigmoid here is safe for these small logits; the general loss should use stable logits, as in DL010 and DL011.

5 · Batched forward pass — the shape you'll actually use

The same parameters act on each row. Bias (2,) broadcasts across rows, whereas a column-shaped (2,1) bias fails for a (4,2) result. For a square (2,2) result it can succeed with the wrong semantics, so shape validation must reflect meaning, not merely legal broadcasting.

X = np.array([[1., 2.], [.5, .5], [-1., .3], [2., -1.5]])
H_pre = X @ W1 + b1
H = np.maximum(0, H_pre)
Y_pre = H @ W2 + b2
Y = 1 / (1 + np.exp(-Y_pre))
assert H.shape == (4, 2) and Y.shape == (4, 1)
assert np.allclose(Y_pre.ravel(), [1.26, .59, .257, .2])
assert np.allclose(Y, np.vstack([1/(1+np.exp(-(
    np.maximum(0, row @ W1+b1) @ W2+b2))) for row in X]))
try:
    X @ W1 + b1[:, None]
except ValueError:
    pass
else:
    raise AssertionError('column bias should reject batch size four')

5.1 The three inputs traced through

For [.5,.5], Z1=[.3,.4], so z=.15+.24+.2=.59 and p≈.643365. For [-1,.3], Z1=[.09,.02], so z=.045+.012+.2=.257 and p≈.563899. For [2,-1.5], Z1=[-.15,-.1], H=[0,0] and z=.2: only the output bias contributes. This last row has zero hidden-layer parameter derivatives through ReLU, but its output-bias gradient need not be zero.

For MNIST, substitute dimensions 784→128→10 and a multiclass endpoint. For a transformer, use (B,T,D)→(B,T,F)→(B,T,D) with a hidden activation and no final sigmoid. Dimensions 4096→16384 are an illustrative FFN, not GPT-3 175B's dimensions.

6 · Solving XOR by hand — the payoff of stacking

Let s=x1+x2. The piecewise-linear tent t(s)=ReLU(s)-2*ReLU(s-1)+ReLU(s-2) is 0,1,0 at s=0,1,2. Use output logit 4*t(s)-2. This yields logits −2,+2,+2,−2 on the XOR truth table: a strict positive margin, unlike the old construction that gave both [0,0] and [0,1] logit zero.

xor = np.array([[0.,0.],[0.,1.],[1.,0.],[1.,1.]])
xor_W1 = np.ones((2,3))
xor_b1 = np.array([0.,-1.,-2.])
xor_W2 = np.array([[4.],[-8.],[4.]])
xor_b2 = np.array([-2.])
xor_H = np.maximum(0, xor @ xor_W1 + xor_b1)
xor_z = xor_H @ xor_W2 + xor_b2
assert np.array_equal(xor_z.ravel(), [-2,2,2,-2])
assert np.array_equal(xor_z.ravel() > 0, [False,True,True,False])
# Original faulty construction: different labels have identical logits.
bad = np.maximum(0, xor @ np.ones((2,3))+[.5,-.5,-1.5]) @ [1.,-2.,1.]-.5
assert bad[0] == bad[1] == bad[2] == 0
# Two hidden units also suffice on these four points.
two = np.maximum(0, xor @ np.ones((2,2))+[0.,-1.]) @ [4.,-8.] - 2
assert np.array_equal(two, xor_z.ravel())

At [0,1], hidden coordinates are [1,0,0]; at [1,1] they are [2,1,0]; at [0,0] all are zero. The output hyperplane 4h1−8h2+4h3−2=0 separates these features. A projection onto arbitrary two hidden coordinates or a PCA projection need not preserve separability; the explicitly chosen first two do here.

Why are two ReLU units sufficient and one insufficient? With one unit, a*ReLU(w·x+b)+c followed by a threshold produces a half-space, the whole space, or the empty set on non-boundary points. XOR is not any such set. The two-unit formula above proves sufficiency on the four binary points, not a universal minimal-width theorem for arbitrary domains. The strip interpretation is a thresholded tent, not a literal paper fold that stacks both red corners: reflection across x1=x2 leaves those corners fixed.

7 · Why ReLU between layers, but not at the very end

A hidden activation prevents affine collapse. Output choices follow the statistical model:

• Binary class: train with raw logits and a stable BCE-with-logits loss; use sigmoid for probabilities.
• Mutually exclusive classes: train logits with cross entropy; use softmax to inspect a distribution.
• Unbounded regression: raw output, often squared error.
• Bounded regression: sigmoid/tanh only if the range and saturation are appropriate.
• Positive rates: softplus or exponential with an appropriate count/rate likelihood; positivity alone does not specify a loss.

For mean binary BCE, dz=(p-y)/N; for multiclass index CE, (P-one_hot)/N. For half squared error the derivative is prediction minus target; ordinary squared error has an extra factor two. ReLU is not a probability map: it allows values above one and has zero slope for negative logits. Replacing it with a fused-logit loss fixes the model/loss mismatch, but promises no universal convergence time.

8 · Parameter counting — and why it matters for GPT

For layer widths d0,...,dL, count sum(d[i-1]*d[i]+d[i]) when each layer has a bias. Thus 2-2-1 has 9; 784-128-10 has 101,770; 784-512-512-10 has 669,706, not 670,530; 784-256-10 has 203,530.

GPT-2 Small's conventional FFN has D=768,F=3072: 2DF+F+D=4,722,432 per block and 56,669,184 across 12 blocks. Count the whole architecture before asserting a fraction. The 124M release config has vocabulary50,257 and context1,024. Tied token/output embeddings plus positional embeddings contribute 768*(50257+1024); each block contributes 12*768**2+13*768 including attention, FFN and two norms; the final norm adds2*768. Total: 124,439,808, so FFNs are about 45.5%, not60%. Two projection matrices dominate many dense blocks, but vocabulary embeddings, tied weights, gated widths, MoE routing and attention design change the total.

For the GPT-3 175B exercise, D=12288,F=49152,L=96 gives 2DFL=115,964,116,992 FFN weights, about 116B; including these FFN biases gives 115,970,015,232. This is not the total model count. A three-matrix bias-free SwiGLU has 3DF weights; matching a two-matrix F=4D block gives ideal F=8D/3, often rounded to a hardware multiple. Count actual configured dimensions.

def count_params(widths):
    return sum(a*b+b for a,b in zip(widths, widths[1:]))
assert [count_params(w) for w in ([2,2,1],[784,128,10],[784,512,512,10],[784,256,10])] == [9,101770,669706,203530]
assert 12*(2*768*3072+3072+768) == 56669184
assert 768*(50257+1024)+12*(12*768**2+13*768)+2*768 == 124439808
assert 96*(2*12288*49152) == 115964116992

9 · Code all together — the forward pass to remember

This retained small-logit forward function exposes the cache required by DL010. An engine saves values needed by derivatives, not necessarily every intermediate; checkpointing may recompute them.

def mlp_forward(X, params):
    W1, b1, W2, b2 = params
    H_pre = X @ W1 + b1
    H     = np.maximum(0, H_pre)
    Y_pre = H @ W2 + b2
    Y     = 1 / (1 + np.exp(-Y_pre))
    # cache intermediates for backprop (Session 010)
    cache = (X, H_pre, H, Y_pre, Y)
    return Y, cache

For arbitrary depth, check adjacent widths and keep a separate output activation. The loop is over layers, not scalar data; vectorization does not remove architectural sequencing.

def forward_layers(X, Ws, bs, activation, output_activation):
    if len(Ws) != len(bs) or not Ws:
        raise ValueError('nonempty aligned layers required')
    for i, (W,b) in enumerate(zip(Ws,bs)):
        if X.shape[-1] != W.shape[0] or b.shape != (W.shape[1],):
            raise ValueError('layer shape mismatch')
        X = X @ W + b
        X = output_activation(X) if i == len(Ws)-1 else activation(X)
    return X
rng = np.random.default_rng(0)
Ws = [rng.normal(size=(2,4)),rng.normal(size=(4,3)),rng.normal(size=(3,1))]
bs = [rng.normal(size=n) for n in (4,3,1)]
relu = lambda x: np.maximum(0,x)
a = forward_layers(xor, Ws, bs, relu, lambda x:x)
b = relu(relu(xor@Ws[0]+bs[0])@Ws[1]+bs[1])@Ws[2]+bs[2]
assert np.allclose(a,b)
# Affine-collapse control, with no activation.
flat = (xor @ Ws[0] + bs[0]) @ Ws[1] + bs[1]
assert np.allclose(flat, xor@(Ws[0]@Ws[1])+bs[0]@Ws[1]+bs[1])
assert not np.allclose(flat, relu(xor@Ws[0]+bs[0])@Ws[1]+bs[1])

10 · MLPs, approximation and modern sequence blocks

Cybenko's 1989 theorem establishes approximation of continuous functions on a compact cube by finite sums of continuous sigmoidal units. It does not promise a given width, efficient training or generalization. Hornik's 1991 result gives approximation in finite-measure Lp criteria for finite p, using bounded nonconstant activations with sufficiently many hidden units; adding continuity supports uniform approximation of continuous maps on compact sets. These are existence results, not a recipe to keep adding layers until validation improves. Depth can represent some compositional functions efficiently; linear-region bounds are architecture-dependent upper bounds, not a guarantee that every added unit or layer realizes them.

Modern blocks illustrate why the exact architecture matters:

• Mamba (2023) combines selective state-space computation with a gated block and removes conventional separate attention/MLP blocks in its homogeneous design. It does not retain an unchanged transformer FFN.
• Jamba (2024) mixes Transformer and Mamba layers with mixture-of-experts components. Expert FFNs add conditional parameter use, not just another dense width.
• Griffin/RecurrentGemma (2024) combine gated linear recurrence with local attention; the recurrent mixing operation and the surrounding feed-forward processing have different jobs.
• RWKV-7 (2025) uses a generalized delta rule for sequence-state updates. Its channel-mixing computation does not justify calling the entire architecture an unchanged two-layer transformer.
• SwiGLU and MoE retain learned projections but change gating, parameter count or which experts execute. The row-batch and chain-rule tools transfer; architectural identity does not.

11 · Failure cases and width–depth decisions

Wrong bias: (4,2)+(2,1) raises; (2,2)+(2,1) can silently add a per-row value. Test a nonsquare batch and distinct bias entries. Dead hidden units: if a unit's preactivation stays negative for all training examples, its incoming gradient through ReLU is zero; the output bias may still learn. Zero hidden bias with appropriate random scaling is a starting point, not a guarantee against death. Transposed weights: square shapes can hide the mistake; compare an asymmetric matrix to a coordinate-level dot product. Repeated first batch: repeatedly using X[:bs] fits only that subset; the model does not literally observe or memorize row order unless order is a feature. Track sampled IDs and paired labels.

With a fixed budget, a shallow wide model shortens sequential depth but can still saturate or optimize poorly; deep narrow models reuse features but may create bottlenecks and unstable Jacobian products. Residuals provide an identity path when dimensions agree, or a learned projection otherwise; they do not abolish optimization difficulty. Compare matched parameter counts, data splits, optimizer settings, multiple seeds and latency. Flat validation results do not prove the data lacks hierarchical structure: the training procedure might be inadequate.

12 · Try it yourself — worked solutions

1. Train XOR: use DL010's complete random-initialized program; inspect learned signs, loss and margins, not a promised 300-step deadline. Hidden-unit permutations and rescaling can describe the same function. A dead-ReLU start may fail even when a hand solution exists.
2. Arbitrary depth: the forward_layers check above matches explicit three-layer composition. Removing hidden activations collapses matrices; changing the output activation changes the modeled task.
3. GPT-3 count: 115,964,116,992 FFN weights, computed in §8, not a blanket 60% rule.
4. Visualize the bend: evaluate h1=ReLU(s), h2=ReLU(s-1) on a grid; draw 4h1−8h2−2=0. The four feature points (0,0),(1,0),(1,0),(2,1) are separated. A general PCA view can hide that separation.

13 · Recap

A hidden nonlinearity changes representational power; matrix shapes and biases fix the arithmetic. Probability endpoints and training logits are distinct. XOR has a verified strict-margin construction, and the transformer parameter fraction must be calculated for a named architecture. The cache is a contract for the next reverse pass, not a reason to save everything forever.

Retention scaffold — eight answers

1. Shapes: X:(N,2), W1:(2,2), b1:(2,), Z1,H:(N,2), W2:(2,1), b2:(1,), z,p:(N,1).
2. Hand trace: [.8,1.1]→1.26→.7790261078; the original .7788 was wrong.
3. 784-256-10: 784*256+256+256*10+10=203530.
4. Output ReLU: not normalized or bounded by one; negative branch has zero slope. Use raw logits for fused BCE.
5. Cache: dW1=X.T@dZ1 requires X; the ReLU pullback needs the positive mask. Checkpointing can recompute rather than retain values.
6. Fraction: GPT-2 Small's FFNs are about 45.5% of its total; GPT-3 175B's FFN weights are about 66.3%. No universal percentage.
7. Learned symmetry: the task is symmetric but individual hidden units need not be. Compare predictions and margins across all four points, not weight equality.
8. Column form: transpose the stored matrices and input batch; Z_col=W_row.T@X.T+b[:,None]. Transposes express layout, not a competing mathematical answer.

Reconstruct §4 tomorrow; reproduce the XOR margin and the affine-collapse control a week later.

Sources and execution boundary

Primary references: PyTorch 2.9 Linear implementation, OpenAI GPT-2, GPT-3 paper, GLU variants, Mamba, Jamba, Griffin, RWKV-7, and the 1986 paper. Dated exact passages and hashes accompany the review. Local NumPy assertions are execution evidence; no large-model training or benchmark was reproduced.