Week 12: Modern Deep Learning & the Road to Transformers
Stat 577 — Statistical Learning Theory
Spring 2026
The Problem with SGD
Two fundamental limitations of vanilla SGD:
1. Oscillation in high-curvature directions. When the loss surface is much steeper in one direction than another, SGD oscillates back and forth across the steep direction while making slow progress along the shallow direction.
2. One learning rate for all parameters. Some parameters need large steps (sparse gradients, infrequent features), others need small steps (dense gradients, frequent features). Vanilla SGD uses the same \(\eta\) for everything.
Momentum
Fix for oscillation. Accumulate a running average of past gradients:
\[\mathbf{m}^{(t)} = \gamma \mathbf{m}^{(t-1)} + \eta \nabla R(\boldsymbol{\theta}^{(t)})\]
\[\boldsymbol{\theta}^{(t+1)} = \boldsymbol{\theta}^{(t)} - \mathbf{m}^{(t)}\]
Typical \(\gamma = 0.9\).
Intuition: A ball rolling downhill accumulates speed in consistent directions and dampens oscillations. Gradients that point the same way reinforce each other; gradients that flip sign cancel out.
BatchNorm in Practice
Where to place it? Before activation (original paper) or after? Both work. No definitive answer.
During inference: No minibatch at test time. Use running averages of \(\mu\) and \(\sigma^2\) computed during training.
In PyTorch: nn.BatchNorm1d(num_features) — one line.
Code: BatchNorm comparison on make_moons
import torch
import torch.nn as nn
from sklearn.datasets import make_moons
# Generate data
X, y = make_moons(n_samples=1000, noise=0.25, random_state=42)
X_t = torch.tensor(X, dtype=torch.float32)
y_t = torch.tensor(y, dtype=torch.float32).unsqueeze(1)
# Model WITHOUT BatchNorm
class PlainMLP(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 64), nn.ReLU(),
nn.Linear(64, 64), nn.ReLU(),
nn.Linear(64, 64), nn.ReLU(),
nn.Linear(64, 1), nn.Sigmoid()
)
def forward(self, x):
return self.net(x)
# Model WITH BatchNorm
class BNormMLP(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 64), nn.BatchNorm1d(64), nn.ReLU(),
nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(),
nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(),
nn.Linear(64, 1), nn.Sigmoid()
)
def forward(self, x):
return self.net(x)
# Train both and record loss curves
def train(model, X_t, y_t, epochs=200, lr=0.01):
opt = torch.optim.Adam(model.parameters(), lr=lr)
loss_fn = nn.BCELoss()
losses = []
for _ in range(epochs):
pred = model(X_t)
loss = loss_fn(pred, y_t)
opt.zero_grad(); loss.backward(); opt.step()
losses.append(loss.item())
return losses
torch.manual_seed(42)
plain = PlainMLP(); losses_plain = train(plain, X_t, y_t)
torch.manual_seed(42)
bnorm = BNormMLP(); losses_bnorm = train(bnorm, X_t, y_t)
# Plot
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(losses_plain, color="#D62828", label="Without BatchNorm", linewidth=2)
ax.plot(losses_bnorm, color="#06A77D", label="With BatchNorm", linewidth=2)
ax.set_xlabel("Epoch"); ax.set_ylabel("BCE Loss")
ax.set_title("Effect of Batch Normalization on Training")
ax.legend(); ax.set_ylim(0, 0.8)
plt.tight_layout()
plt.savefig("../figures/week12/03_batchnorm_effect.png", dpi=200, bbox_inches="tight")
plt.show()
Dropout Mechanism
During training: Randomly set each neuron’s output to zero with probability \(p\).
- Typical \(p = 0.5\) for hidden layers, \(p = 0.2\) for the input layer
During testing: Use all neurons. But their outputs are now \(p\) times too large (since during training only a fraction \(1 - p\) were active on average).
Two equivalent fixes:
- Scale outputs by \((1 - p)\) at test time
- Inverted dropout (standard): scale by \(\frac{1}{1-p}\) during training, do nothing at test time
Convolution = Local + Shared Weights
A small filter (e.g., \(3 \times 3\)) slides across the image. At each position, it computes a dot product between the filter weights and the local image patch.
Same weights everywhere — weight sharing — so a pattern detector learned in one location works in all locations.
Dramatically reduces parameters: a \(3 \times 3\) filter has 9 parameters regardless of image size.
The Bottleneck
In the standard encoder-decoder architecture, the encoder processes the entire input sequence and compresses it into a single fixed-length vector \(\mathbf{h}_T\):
\[\underbrace{x_1, x_2, \ldots, x_T}_{\text{input sequence}} \;\xrightarrow{\text{encoder}}\; \mathbf{h}_T \;\xrightarrow{\text{decoder}}\; \underbrace{y_1, y_2, \ldots, y_{T'}}_{\text{output sequence}}\]
The decoder must reconstruct the entire output from this one vector.
Imagine summarizing a 500-word paragraph into a single number. That is what the encoder-decoder bottleneck asks the network to do.
For short sequences, this works. For long sequences, too much information is lost. This is the problem attention solves.
Attention as Soft Retrieval
Three steps:
1. Score: How relevant is each encoder state \(\mathbf{h}_j\) to the current decoder state \(\mathbf{s}_t\)?
\[e_{tj} = \text{score}(\mathbf{s}_t, \mathbf{h}_j)\]
2. Normalize: Softmax converts scores to weights that sum to 1:
\[\alpha_{tj} = \text{softmax}_j(e_{tj}) = \frac{\exp(e_{tj})}{\sum_{j'} \exp(e_{tj'})}\]
3. Aggregate: Context vector is a weighted average of all encoder states:
\[\mathbf{c}_t = \sum_j \alpha_{tj} \mathbf{h}_j\]
A soft lookup — instead of selecting one encoder state (hard attention), we blend all of them with learned, context-dependent weights.
What Self-Attention Captures
Consider the sentence: “The animal didn’t cross the street because it was too tired.”
When processing “it,” self-attention attends strongly to “animal” — resolving the pronoun.
This requires looking far back in the sequence — something RNNs struggle with due to vanishing gradients.
Self-attention gives every position a direct connection to every other position. No bottleneck, no sequential processing.
Attention Demo: From Scratch
Code: Scaled dot-product attention in ~10 lines
import numpy as np
def scaled_dot_product_attention(Q, K, V):
"""Scaled dot-product attention (Vaswani et al., 2017)."""
d_k = Q.shape[-1]
scores = Q @ K.T / np.sqrt(d_k) # (n_queries, n_keys)
scores = scores - scores.max(axis=-1, keepdims=True) # numerical stability
weights = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True) # softmax
return weights @ V, weights # (n_queries, d_v), (n_queries, n_keys)
# Small example: 4 tokens, d_k = d_v = 3
np.random.seed(42)
X = np.random.randn(4, 3) # 4 tokens, embedding dim 3
W_Q = np.random.randn(3, 3) * 0.5
W_K = np.random.randn(3, 3) * 0.5
W_V = np.random.randn(3, 3) * 0.5
Q, K, V = X @ W_Q, X @ W_K, X @ W_V
output, attn_weights = scaled_dot_product_attention(Q, K, V)
# Visualize attention weights
fig, ax = plt.subplots(figsize=(6, 5))
tokens = ["The", "cat", "sat", "down"]
im = ax.imshow(attn_weights, cmap="Blues", vmin=0, vmax=1)
ax.set_xticks(range(4)); ax.set_xticklabels(tokens)
ax.set_yticks(range(4)); ax.set_yticklabels(tokens)
ax.set_xlabel("Key (attending to)"); ax.set_ylabel("Query (attending from)")
ax.set_title("Self-Attention Weights")
for i in range(4):
for j in range(4):
ax.text(j, i, f"{attn_weights[i,j]:.2f}", ha="center", va="center",
color="white" if attn_weights[i,j] > 0.5 else "black", fontsize=12)
plt.colorbar(im, ax=ax, label="Attention Weight")
plt.tight_layout()
plt.savefig("../figures/week12/09_attention_demo.png", dpi=200, bbox_inches="tight")
plt.show()
Note: These are random, untrained weights — the attention pattern is meaningless. In a trained model, the weights would reflect learned relationships between tokens.