PidokuInfra
04Neural Networks from Scratch
On this page

Neural Networks from Scratch

Foundations Beginner 1h 30m Difficulty 2/5 Topic 04 of 12

Prerequisites 01, 02, 03

1. What is it?#

A neural network is a stack of alternating linear transformations and nonlinear functions:

x → [linear] → [nonlinear] → [linear] → [nonlinear] → ... → y

That’s the entire structure. Everything else — attention, convolutions, normalization — is a variation on “linear thing, then nonlinear thing, repeated.”

2. Why does it exist?#

The linear parts do the learning; the nonlinear parts make depth meaningful.

Without nonlinearity, stacking layers is pointless:

W₂(W₁x) = (W₂W₁)x = Wx

Two linear layers collapse into one. A hundred linear layers collapse into one. Depth buys you nothing.

Insert a nonlinearity and the collapse is impossible:

W₂ · relu(W₁x)   cannot be written as a single matrix multiply

Now depth composes: each layer can build features from the previous layer’s features. That is the whole reason “deep” learning is deep.

3. Simple analogy#

An assembly line of specialists.

Station 1 takes raw material and produces simple parts (edges, in a vision model; character patterns, in a language model). Station 2 combines simple parts into subassemblies. Station 3 combines those into components. By the last station you have something complex built from many layers of simple combinations.

The nonlinearity is the “decide whether this part is worth passing on” step. Without it, every station is just re-scaling the raw material, and you might as well have one station.

4. Tiny example#

A complete 2-layer network solving XOR — the classic demonstration that nonlinearity matters:

Go
package main

import "fmt"

// XOR is not linearly separable: no single layer can compute it.
// These hand-set weights solve it with one hidden layer.
var (
	W1 = [2][2]float32{{1, 1}, {1, 1}} // (2, 2)
	b1 = [2]float32{0, -1}
	W2 = [2]float32{1, -2} // (2,)
	b2 = float32(0)
)

func forward(x [2]float32) float32 {
	var h [2]float32
	for i := range h {
		h[i] = max(0, W1[i][0]*x[0]+W1[i][1]*x[1]+b1[i]) // ReLU
	}
	return W2[0]*h[0] + W2[1]*h[1] + b2
}

func main() {
	X := [][2]float32{{0, 0}, {0, 1}, {1, 0}, {1, 1}}
	y := []float32{0, 1, 1, 0}
	for i, x := range X {
		fmt.Println(x, "→", forward(x), "target", y[i])
	}
	// [0 0] → 0   [0 1] → 1   [1 0] → 1   [1 1] → 0
}

Now delete the ReLU (h = x @ W1.T + b1) and try to find any W1, b1, W2, b2 that solves XOR. You cannot — it’s a theorem. That’s the nonlinearity earning its place.

5. Technical explanation#

The building blocks of a transformer#

You only need five:

1. Linear (fully connected)

y = x @ W.T + b        # W: (out, in), b: (out,)

Parameters: out × in + out. This is where almost all parameters live.

2. Activation — elementwise, no parameters (file 05).

3. Normalization — stabilizes the scale of activations.

Go
// LayerNorm: normalize over the feature dimension
func layerNorm(x, gamma, beta []float32) []float32 {
	var mean, variance float32
	for _, v := range x {
		mean += v
	}
	mean /= float32(len(x))
	for _, v := range x {
		variance += (v - mean) * (v - mean)
	}
	variance /= float32(len(x))
	inv := 1 / float32(math.Sqrt(float64(variance)+1e-6))
	y := make([]float32, len(x))
	for i, v := range x {
		y[i] = (v-mean)*inv*gamma[i] + beta[i] // gamma, beta learned, shape (d,)
	}
	return y
}

// RMSNorm (used by Llama, most modern LLMs): cheaper, no mean subtraction
func rmsNorm(x, gamma []float32) []float32 {
	var sumSq float32
	for _, v := range x {
		sumSq += v * v
	}
	inv := 1 / float32(math.Sqrt(float64(sumSq/float32(len(x)))+1e-6))
	y := make([]float32, len(x))
	for i, v := range x {
		y[i] = v * inv * gamma[i] // gamma learned, no beta
	}
	return y
}

RMSNorm drops the mean subtraction and the bias — fewer ops, fewer parameters, empirically no worse. Two fewer passes over the data per norm, and there are 2 norms per layer × 80 layers. Memory-bound op, so those savings are real.

4. Residual connection

x = x + sublayer(x)

Free in parameters, enormous in effect: gradients flow through the identity path during training, and at inference it means the “residual stream” carries information straight through the network with each layer adding a correction. Understanding LLMs as “a residual stream that layers read from and write to” is a genuinely useful mental model.

5. Embedding — a lookup table (file 06).

A complete MLP block, written out#

Go
// SwiGLU, as used in Llama/Mistral/Qwen. x is one token's vector, shape (d,).
func mlpBlock(x []float32, Wgate, Wup, Wdown Mat) []float32 {
	gate := Wgate.Apply(x) // (d_ff,)
	up := Wup.Apply(x)     // (d_ff,)
	for i, g := range gate {
		silu := g / (1 + float32(math.Exp(float64(-g)))) // SiLU/Swish: g * sigmoid(g)
		gate[i] = silu * up[i]
	}
	return Wdown.Apply(gate) // (d,)
}

Three matmuls, one elementwise multiply, one sigmoid. That is 80% of a transformer’s parameters and roughly 60-65% of its FLOPs.

Why three matrices instead of the classic two? The gating (act * up) empirically works better per-parameter. To keep the parameter count comparable to a 2-matrix FFN with d_ff = 4d, d_ff is set to about 8d/3 (e.g. 11008 for d=4096, or 14336 in Llama 3 8B).

Counting parameters#

Linear(in, out):     in × out (+ out if bias)
LayerNorm(d):        2d
RMSNorm(d):          d
Embedding(V, d):     V × d

Per transformer layer (Llama-style, GQA):

q_proj: d × d
k_proj: d × d_kv          where d_kv = n_kv_heads × head_dim
v_proj: d × d_kv
o_proj: d × d
gate:   d × d_ff
up:     d × d_ff
down:   d_ff × d
2 RMSNorms: 2d

Exercise the formula on a model you know. It’s the fastest way to internalize architecture.

6. Under the hood#

What one layer costs, at inference, for one token:

d = 4096, d_ff = 14336, d_kv = 1024 (GQA 8 heads)

FLOPs (2 × params touched):
  q_proj    2 · 4096 · 4096  =  33.6 M
  k_proj    2 · 4096 · 1024  =   8.4 M
  v_proj    2 · 4096 · 1024  =   8.4 M
  o_proj    2 · 4096 · 4096  =  33.6 M
  gate      2 · 4096 · 14336 = 117.4 M
  up        2 · 4096 · 14336 = 117.4 M
  down      2 · 14336 · 4096 = 117.4 M
  ────────────────────────────────────
  total                       436 M FLOPs per token per layer
  × 32 layers                = 14 GFLOPs      ✓ matches 2P for an 8B model
  attention over KV          + 4·L·S·d (grows with context)

Norms, activations, residuals: ~0.1% of FLOPs, but they read and write the full
activation tensor each time → disproportionate share of MEMORY traffic.

That last line is why kernel fusion matters (Section IV.09): the norms and elementwise ops are FLOP-trivial and bandwidth-expensive, so combining them with their neighbors is nearly free performance.

7. Performance implications#

  • Matmuls dominate FLOPs (>95%) but norms/activations/residuals take 10-25% of time because they’re memory-bound.
  • RMSNorm over LayerNorm saves a pass over the data — measurable at scale.
  • The residual add is a separate kernel unless fused. Fusing norm+residual is a standard optimization.
  • SwiGLU’s three matmuls can be fused into one by concatenating gate and up weights — a common trick that turns two GEMMs into one larger, more efficient GEMM.

8. Production implications#

  • You will read model code constantly. Being able to look at modeling_llama.py and immediately see “7 linears, 2 norms, 2 residuals per layer” is the baseline skill.
  • Architecture variations matter for serving. GQA changes KV size; SwiGLU changes FFN shape; RMSNorm changes kernel fusion opportunities; RoPE (file 09) changes what you cache.
  • Weight fusion at load time (concatenating gate+up, or q+k+v) is a standard startup optimization that improves GEMM efficiency.

9. Common mistakes#

Thinking the nonlinearity is a minor detail. Without it there is no network.

Confusing LayerNorm and RMSNorm. Different parameter counts, different kernels, and using the wrong one silently degrades quality.

Forgetting that d_ff is not always 4d. SwiGLU models use ~8d/3 to keep parameters comparable. Getting this wrong ruins your parameter-count estimate.

Ignoring the memory cost of “cheap” ops. A residual add is 1 FLOP per element and 12 bytes of traffic per element in FP32.

Assuming bias terms exist. Most modern LLMs omit biases in linear layers.

10. Hands-on exercise#

A. XOR by hand. Verify the network in section 4 on paper for all four inputs. Then prove (algebraically) that no linear model can solve XOR.

B. Build a transformer layer in Go. Implement rmsnorm, silu, swiglu_mlp, and a residual connection. No attention yet — just the MLP half of a layer. Verify shapes.

C. Parameter counting. For Llama-3-8B (d=4096, d_ff=14336, L=32, n_heads=32, n_kv_heads=8, head_dim=128, V=128256), compute the parameter count using the formulas in section 5. Compare to 8.03B. Explain any discrepancy.

D. Where does time go? Profile a real transformer layer. What fraction of time is in matmuls vs norms vs elementwise? Compare to the FLOP fractions. Explain the difference.

E. Fusion win. Implement rmsnorm as (i) three separate PyTorch ops and (ii) one fused op (torch.compile it, or use an existing fused kernel). Measure the difference.

11. Interview questions#

  1. Why do neural networks need nonlinearities? Prove it.
  2. What is the difference between LayerNorm and RMSNorm, and why do modern LLMs prefer RMSNorm?
  3. Write out a SwiGLU FFN and explain why d_ff ≈ 8d/3.
  4. What fraction of a transformer’s parameters are in the FFN? Why does that matter for quantization and MoE?
  5. Why do norms and residual adds take a disproportionate share of time relative to FLOPs?
  6. What is the residual stream and why is it a useful mental model?

12. Further reading#

  • [FUNDAMENTAL] Karpathy, “Neural Networks: Zero to Hero” (video series) and micrograd
  • [ESTABLISHED] Zhang & Sennrich, “Root Mean Square Layer Normalization” (2019)
  • [ESTABLISHED] Shazeer, “GLU Variants Improve Transformer” (2020) — where SwiGLU comes from
  • Next: 05 — Activation functions
esc

↑↓ navigate↵ openesc close