Handout #8 — Neural Network Training and Tokenization

A correct training loop can still learn slowly, lose its gradient signal, or fit the training set while getting worse on new data. This reading connects activation scale, gradient flow, and generalization through worked examples and diagnostic experiments. We then turn to the interface between language and the network: which pieces of text become its inputs and prediction targets.

Use the worked examples to predict what a training run will do, then test the explanation. Try each exercise before opening its answer.

1Activation functions and their derivatives

Slides 183–188

At zero, nonsmooth functions need a chosen derivative convention. Saturation is gradual: sigmoid derivatives at finite inputs are mathematically positive, although finite precision can round them to zero. A “dead ReLU” stays inactive across the relevant training inputs, not merely on one example. Leaky ReLU avoids an exactly zero negative slope but cannot guarantee learning.

Monotonicity is not required for a valid activation. GELU and SiLU/Swish provide smooth, nonmonotone alternatives. For the feedforward experiments in these slides, ReLU is a useful baseline; compare alternatives with the same data split and training budget. PReLU learns its negative-side slope rather than fixing it. Softplus is a smooth approximation to ReLU; SwiGLU is a gated construction using two projections, rather than simply a scalar activation replacement.

For an elementwise activation $h=f(z)$, backward computes $g_z=g_h\odot f\prime(z)$. At $z=5$, sigmoid contributes about 0.00665; at $z=3$, tanh contributes about 0.00987. These are small, not mathematically zero. Zero-centered outputs can improve optimization geometry, but tanh's mean is not guaranteed to be zero for an arbitrary input distribution.

2Vanishing gradients, exploding gradients, and residuals

Slides 189–193

Handout #5 derived $g_i=A_i^\top g_{i+1}$. Repeating that rule multiplies many Jacobians. For a scalar illustration, twenty factors of $0.5$ give about $9.54\times10^{-7}$; twenty factors of 2 give $1{,}048{,}576$. In multiple dimensions, amplification depends on directions and singular values, not just whether individual matrix entries exceed 1.

This is an upper bound, not a prediction that every direction grows at that rate. Monitor activation and gradient norms at several layers. Nonfinite gradients, nearly zero gradients, and large but finite gradients require different diagnoses. Gradient clipping limits the gradient norm before the optimizer uses it; it does not directly bound every optimizer’s parameter update, and it cannot repair a NaN already produced in the forward pass.

A residual block computes $h_{i+1}=h_i+F(h_i)$, so its backward rule is

The identity path supplies a direct gradient contribution alongside the learned branch. Shapes must agree for addition, or a projection is needed. Residuals improve gradient flow in many networks; they do not guarantee that gradients cannot vanish or explode.

Consider a block with two linear layers and a final activation (the example on slide 193): out = linear2(relu(linear1(x))); out = relu(out + x). The skip connection is the addition of x, not the assignment identity = x by itself. Because that example applies ReLU after the sum, its full Jacobian is $D_{\mathrm{out}}(I+J_F)$, where $D_{\mathrm{out}}$ is the final ReLU mask. The identity contribution can still be blocked by that mask. In a plain $x+F(x)$ block, $F(x)=0$ passes $x$ through; the branch learns a correction to the identity.

This connects three complementary controls: residuals add a direct path, initialization sets starting signal scales, and normalization adjusts hidden scales during computation. None substitutes for checking the actual forward and backward values.

3Initialization: break symmetry without losing scale

Slides 194–195

If hidden units start with identical incoming weights and biases and identical outgoing connections, they receive identical gradients and remain interchangeable. Random weights break this symmetry. Zero biases can be fine when the weights already distinguish the units; the warning is about initializing the whole network symmetrically.

The noise scale matters too: large weights can saturate activations or amplify gradients, while very small weights can weaken signals. To see why fan-in matters, let $z=\sum_{j=1}^n w_jx_j$, with independent, zero-mean weights of variance $s^2$, independent of the inputs. Under the usual independence assumptions, $\mathbb E[z^2]=n s^2\mathbb E[x_j^2]$. Unit-variance weights can therefore enlarge signals as width grows; the scale should depend on the number of summed inputs.

Xavier/Glorot's gain-1 choice uses $\operatorname{Var}(w)=2/(n_{\rm in}+n_{\rm out})$ to balance forward and backward scales. Kaiming's ReLU fan-in choice uses $2/n_{\rm in}$: for a symmetric preactivation, ReLU keeps half of its second moment. This is a second-moment argument, not a claim that ReLU outputs have mean zero or exactly preserve centered variance. For leaky ReLU slope $\alpha$, the analogous choice is $2/((1+\alpha^2)n_{\rm in})$.

import torch
from torch import nn

layer = nn.Linear(100, 64)  # weight shape: (fan_out, fan_in)
nn.init.kaiming_normal_(layer.weight, mode="fan_in", nonlinearity="relu")
nn.init.zeros_(layer.bias)
assert layer.weight.shape == (64, 100)
# Suitable for a layer followed by ReLU; choose separately for the output layer.

These are starting-scale heuristics under simplified assumptions, not guarantees throughout training. See PyTorch's initialization reference for gains and weight-layout conventions.

4Normalization: choose what to center

Slides 196–198

For a numerical input feature, standardization uses $(x-\mu)/\sqrt{v+\epsilon}$, where $\mu$ and $v$ are estimated from the training split and reused on validation and test inputs. Do not standardize discrete token IDs: their integer values are labels, not magnitudes.

Centering can change which update directions are available: for a single affine unit and one example, $\nabla_w\ell=\delta x$. If every coordinate of $x$ is positive, the weight-gradient coordinates share a sign. Centering can improve this geometry; gradients summed across examples need not share a sign, and centering is not a universal convergence guarantee.

Normalization inside a network is an operation with its own axes. For $X$ shaped $(B,T,d)$, LayerNorm(d) computes mean and variance over the final $d$ coordinates separately for each token. Batch normalization uses statistics across the batch and relevant spatial axes and normally tracks running statistics for evaluation. These choices are not interchangeable. For $(B,T,d)$, LayerNorm has $BT$ separate groups of $d$ features and typically $2d$ learned scale/shift parameters. For a $(B,d)$ batch, BatchNorm has $d$ groups of $B$ examples; for sequence input in PyTorch's $(B,d,T)$ layout, BatchNorm1d(d) aggregates over $B$ and $T$. Padding then needs care because those statistics can include padded positions. LayerNorm uses the current token's statistics in both training and evaluation; BatchNorm normally switches to running estimates for evaluation.

4.1  What information does each normalization remove?

Take two examples with two features each, $X=\begin{bmatrix}1&3\\101&103\end{bmatrix}$. Ignore $\epsilon$ for this calculation and set $\gamma=1,\beta=0$. LayerNorm groups the two features within each row; BatchNorm groups the two examples within each column. Use population variance (divide by the number of values in the group).

LayerNorm removes the 100-unit offset between these two rows: it represents which features are high or low relative to the same example. BatchNorm represents which examples are high or low relative to the current batch for that feature. Their outputs answer different questions even though both operations center and rescale numbers.

The following experiment checks LayerNorm's axes and formula on a three-dimensional tensor.

import torch
from torch import nn

torch.manual_seed(1)
X = torch.randn(2, 3, 4, dtype=torch.float64)
layer = nn.LayerNorm(4, elementwise_affine=False, eps=1e-5).double()
manual = (X - X.mean(-1, keepdim=True)) / torch.sqrt(
    X.var(-1, unbiased=False, keepdim=True) + 1e-5)
torch.testing.assert_close(layer(X), manual)
torch.testing.assert_close(layer(X).mean(-1), torch.zeros(2, 3, dtype=X.dtype), atol=1e-12, rtol=0)
print("One mean and variance per token; output shape:", layer(X).shape)

5Validation, early stopping, and dropout

Slides 199–205

Training loss measures fit to the data used for updates. Validation loss helps choose stopping time and hyperparameters. Save the parameters with the best validation score and restore that checkpoint; a patience window avoids reacting to a single noisy epoch. Keep the test set for the final evaluation. Patience delays the decision to stop so that one noisy evaluation does not end the run; it does not change which checkpoint had the best score. Save a copy of those weights (for example with copy.deepcopy(model.state_dict())), not just a reference to a model that subsequent updates will change.

Compare losses on the same target definition and reduction. Dropout makes training predictions noisier, so training-mode loss can even exceed evaluation loss. If both losses stay high, check optimization and capacity; if only held-out performance worsens, investigate overfitting and data shift. Repeatedly choosing hyperparameters on validation also adapts to that split; reserve the test set for the final report.

One way to understand dropout is as training many subnetworks with shared weights. A layer with $n$ independently masked activations has $2^n$ possible masks, but we sample masks rather than train and store that many independent networks. A fresh mask perturbs which features are available; the model cannot always rely on the same activation being present.

Dropout randomly removes activation contributions during training. Let $p$ be the probability of dropping and $m\sim\operatorname{Bernoulli}(1-p)$. PyTorch uses inverted dropout:

Since $\mathbb E[m]=1-p$, $\mathbb E[\widetilde h\mid h]=h$. For an unscaled training mask, the alternative is multiplying activations by the keep probability at evaluation. Do not apply both conventions. Expectation matching is exact for the activation and a following affine operation; a nonlinear whole network is not exactly the average of all masked networks.

import torch
from torch import nn

torch.manual_seed(2)
drop = nn.Dropout(p=0.25)
x = torch.ones(10000)
drop.train()
y = drop(x)
assert torch.all((y == 0) | torch.isclose(y, torch.full_like(y, 4 / 3)))
assert abs(y.mean().item() - 1) < 0.03
# Disabling graph recording does not turn dropout off.
with torch.no_grad():
    assert (drop(x) == 0).any()
drop.eval()
torch.testing.assert_close(drop(x), x)
print("Training mean:", y.mean().item(), "; eval is identity")

6First debug on a tiny dataset

Slides 206–208

Before a long training run, ask whether a sufficiently expressive model can fit a handful of consistent examples. Disable dropout and other regularization for this check, use a fixed seed, and inspect the loss. Success checks that some parts work together; it does not prove labels, data splits, or evaluation are correct. Failure can mean a bug, unsuitable learning rate, too little capacity, or contradictory labels.

import torch
from torch import nn

torch.manual_seed(3)
X = torch.tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y = torch.tensor([0, 1, 1, 0])  # XOR: requires nonlinear features
model = nn.Sequential(nn.Linear(2, 16), nn.Tanh(), nn.Linear(16, 2))
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
model.train()
for step in range(300):
    optimizer.zero_grad(set_to_none=True)
    loss = nn.functional.cross_entropy(model(X), y)
    assert torch.isfinite(loss)
    loss.backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all()
               for p in model.parameters())
    optimizer.step()
model.eval()
with torch.no_grad():
    logits = model(X)
    accuracy = (logits.argmax(-1) == y).float().mean().item()
    final_loss = nn.functional.cross_entropy(logits, y).item()
assert accuracy == 1.0 and final_loss < 0.02
print("Tiny-data accuracy:", accuracy, "; loss:", final_loss)

If loss is nonfinite, find the first nonfinite intermediate. If gradients are missing, check whether the parameter participates in the loss and whether an operation detached it. If gradients are zero, inspect activation saturation and masks. If gradients are finite but parameters do not change, check the optimizer's parameter list, learning rate, and call to step(). A low training loss with worsening validation loss instead calls for checking generalization.

6.1  From the first successful fit to a useful experiment

Slides 207–208

The loss as a function of a neural network's parameters is generally nonconvex. Backprop computes gradients efficiently; it does not guarantee a global optimum, convergence for an arbitrary learning rate, or generalization. Different seeds can expose sensitivity to initialization and data order. Compare several runs where feasible, using a fixed validation protocol and compute budget.

Ensembling averages predictions from trained models (or uses voting for class decisions), which may improve held-out performance but costs more computation. It does not make an individual run escape a poor solution. Likewise, a larger model can sometimes be easier to optimize, but scale does not generally turn the objective into a convex function.

1. Check shapes, paired labels, finite values, and gradient flow on one batch.
2. Fit a tiny consistent dataset with regularization disabled.
3. Restore regularization, train on the training split, and monitor held-out loss.
4. Change a justified hyperparameter, record the result, and restore the best validation checkpoint.

The same reasoning carries beyond feedforward networks: track what is computed, what is saved, what is differentiated, and what evidence shows that the model generalizes.

6.2  Diagnostic capstone: three curves, more than one explanation

7Reading real code: PyTorch dropout

Companion to slide 205 · torch/nn/modules/dropout.py · Dropout.forward

This method from the installed PyTorch 2.12.0 source delegates to the functional operation. The excerpt below preserves its executable lines and omits only the docstring:

    def forward(self, input: Tensor) -> Tensor:
        return F.dropout(input, self.p, self.training, self.inplace)

Trace self.p to the drop probability, self.training to the mode toggled by train() and eval(), and self.inplace to whether storage may be modified in place. No check of gradient recording appears here: that is why no_grad() and evaluation mode have distinct jobs. See the Dropout documentation and source link. The normalization axes used above are documented in LayerNorm.

8From text to tokens: choosing the units the model sees

Tokenization slides 1–12 · Atomic Units of Language

A neural network operates on numbers, but our observations are strings. Tokenization chooses how to express a string as a sequence of discrete symbols. That choice determines what counts as one prediction, how many predictions represent a sentence, and which distinctions the model can observe. It is part of the modeling setup, not just a formatting step.

8.1  Text pieces, integer IDs, and vectors are different things

A tokenizer maps text to vocabulary entries and their integer IDs. An embedding lookup then maps each ID to a learned vector. For vocabulary size $V$ and embedding dimension $d$, an embedding matrix $E\in\mathbb R^{V\times d}$ gives token $i$ the vector $E[i]$. A batch of IDs has shape $(B,T)$; its embeddings have shape $(B,T,d)$. These are the tensor axes from Handout #7.

The ID is a row address, not a numerical measurement of meaning. If “cat” has ID 20 and “dog” has ID 21, their proximity does not make them similar; the learned vectors can express similarity. Renumbering a vocabulary while consistently permuting the corresponding embedding and output rows changes the labels, not the model's predictions over token strings. Mixing IDs from one tokenizer with another model's unpermuted vocabulary instead gives the rows the wrong meanings.

At the output, a language model produces $V$ scores, one per vocabulary entry. A common affine output layer has $Vd$ weights plus any bias. Thus tokenization affects both the sequence axis $T$ and vocabulary axis $V$. We can discuss that interface without yet choosing a neural language-model architecture. Beginning/end markers such as BOS and EOS are additional symbols when the model's convention uses them; they are not literal words in the text.

8.2  Where does the cost go?

Consider the exact string cats run., including one space and a period. The following segmentations are illustrative choices, not predictions from a particular trained tokenizer:

Words keep sequences short but need many entries. New names, inflections, compounds, and typos continually create unfamiliar forms. Mapping them to one UNK symbol lets the model continue, but erases their identities. Whitespace also does not reliably define word boundaries across languages. Even in English, a simple split treats run and run. as different entries unless another preprocessing rule separates punctuation.

Characters reuse a smaller inventory but lengthen sequences. In this example, one word-to-word relationship spans several character positions. More positions mean more activation storage and prediction steps for many architectures. A character vocabulary learned only from training data can still miss a new character; “fewer unknowns” is not “none under every encoding.” The word “character” also needs a convention for combining marks and other multi-part written symbols.

Subwords trade between these costs. Frequent strings can get whole entries, while uncommon strings reuse shorter pieces. Those pieces sometimes align with meaningful parts such as a plural suffix, but frequency-driven pieces need not be morphemes or pronounceable words. Encoding an unfamiliar word requires that the tokenizer's underlying pieces can cover it; the label “subword” alone does not guarantee complete coverage.

The slides also ask about bits: a tiny symbol inventory can represent text under a specified encoding, but it expands the sequence further. For this all-ASCII example, UTF-8 uses 9 bytes, or 72 bits before any boundary markers. A small vocabulary is not by itself a cheap language model.

8.3  Read the tokenizer output, not the word count

For The unbelievably fluffy cat purred., slide 11 records the following outputs with added special tokens disabled. These are the lecture's observed examples, not a promise that every tokenizer revision will produce the same result.

The vertical bars above mark boundaries we added for readability. In these displayed token spellings, Ġ represents a leading space; it is not a letter that appeared in the original sentence. A tokenizer's display strings are an inspection format. Use its decoder to recover text rather than assuming that concatenating displayed spellings always reproduces it.

Notice that the same printed word can have different IDs and boundaries across tokenizers. The split Ġp | urred is not a linguistic analysis of “purred”; it is a vocabulary choice. Neither eight tokens nor ten tokens alone tells you which resulting language model is better. You have measured segmentation length, not prediction quality.

8.4  Spaces, unfamiliar words, and numbers are part of the input

Slide 12 uses The cat saw a quokka 12345., with two spaces before the number. An unfamiliar word may split into pieces; the number may split too. A space can be attached to a following piece or represented separately. Do not assume “one word = one token,” “one number = one token,” or “spaces are free.”

Try predicting boundaries for four inputs: cat, cat, cat, and cat.. Then inspect them with the same tokenizer and special-token setting used in a class example. Record the IDs, display pieces, token count, and decoded text. Keep leading/repeated spaces visible when printing, for example with a string representation that includes quotes. A changed count or ID sequence is an observation to explain, not a failure of the tokenizer.

To make a comparison meaningful, record the tokenizer/model identifier and revision when available, exact input text, and whether special tokens were added. Inspect one-token decodings carefully: some tokens represent only part of an encoded character. Decoding the full sequence is the appropriate round-trip check; tokenizers with normalization may intentionally change the text.