Stochastic Gradient Descent › Implementation: SGD for NLP
Language AI HandbookAll chapters Stochastic Gradient Descent › Implementation: SGD for NLPOne section · Sections in this chapter11 / 16 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement SGD from scratch and apply it to a practical NLP task: training a logistic regression classifier on sentiment analysis, then extending to a small multilayer perceptron.
Setting Up the Problem
We will train on the 20 Newsgroups dataset, classifying posts between two sports categories using TF-IDF features. This lets us see SGD in action on a real text classification task without requiring any external downloads.
import numpy as np
from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
# Use 20 Newsgroups with 2 categories for binary sentiment analogy
# (A practical NLP classification task available without download)
categories = ["rec.sport.hockey", "rec.sport.baseball"]
data = fetch_20newsgroups(
subset="all",
categories=categories,
remove=("headers", "footers", "quotes"),
random_state=42,
)
# Build TF-IDF features
vectorizer = TfidfVectorizer(max_features=5000, sublinear_tf=True)
X = vectorizer.fit_transform(data.data).toarray().astype(np.float32)
y = np.array(data.target, dtype=np.float32)
# Train/validation/test split
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
X_train, X_val, y_train, y_val = train_test_split(
X_train, y_train, test_size=0.15, random_state=42
)Training samples: 1354 Validation samples: 240 Test samples: 399 Feature dimensions: 5000
The training set has several hundred thousand parameters to learn across the TF-IDF feature space. SGD processes mini-batches of this data to iteratively improve the model.
Logistic Regression with SGD
Before building a multi-layer network, let us implement logistic regression with pure SGD. This isolates the optimization algorithm from architecture complexity.
def sigmoid(z):
"""Numerically stable sigmoid."""
return np.where(z >= 0, 1 / (1 + np.exp(-z)), np.exp(z) / (1 + np.exp(z)))
def binary_cross_entropy(y_pred, y_true):
"""Binary cross-entropy loss (averaged over the batch)."""
eps = 1e-9
return -np.mean(
y_true * np.log(y_pred + eps) + (1 - y_true) * np.log(1 - y_pred + eps)
)
def compute_accuracy(X, y, w, b):
"""Compute binary classification accuracy."""
probs = sigmoid(X @ w + b)
predictions = (probs > 0.5).astype(float)
return np.mean(predictions == y)
# Initialize weights to small random values (important for breaking symmetry)
np.random.seed(42)
input_dim = X_train.shape[1]
w = np.random.randn(input_dim).astype(np.float32) * 0.01
b = np.float32(0.0)def sgd_logistic(
X_train,
y_train,
X_val,
y_val,
learning_rate=0.1,
batch_size=64,
n_epochs=30,
seed=42,
):
"""Train logistic regression using mini-batch SGD."""
np.random.seed(seed)
n = X_train.shape[0]
d = X_train.shape[1]
w = np.random.randn(d).astype(np.float32) * 0.01
b = np.float32(0.0)
train_losses, val_losses, train_accs, val_accs = [], [], [], []
for epoch in range(n_epochs):
# Shuffle training data at the start of each epoch
perm = np.random.permutation(n)
X_shuffled = X_train[perm]
y_shuffled = y_train[perm]
epoch_loss = 0.0
n_batches = 0
for start in range(0, n, batch_size):
X_batch = X_shuffled[start : start + batch_size]
y_batch = y_shuffled[start : start + batch_size]
# Forward pass
logits = X_batch @ w + b
probs = sigmoid(logits)
# Compute gradients (chain rule through cross-entropy + sigmoid)
error = probs - y_batch # shape: (B,)
grad_w = X_batch.T @ error / len(y_batch)
grad_b = np.mean(error)
# SGD update
w -= learning_rate * grad_w
b -= learning_rate * grad_b
epoch_loss += binary_cross_entropy(probs, y_batch)
n_batches += 1
# Record metrics
val_probs = sigmoid(X_val @ w + b)
train_losses.append(epoch_loss / n_batches)
val_losses.append(binary_cross_entropy(val_probs, y_val))
train_accs.append(compute_accuracy(X_train, y_train, w, b))
val_accs.append(compute_accuracy(X_val, y_val, w, b))
return w, b, train_losses, val_losses, train_accs, val_accs
# Train with default settings
w_sgd, b_sgd, train_losses, val_losses, train_accs, val_accs = sgd_logistic(
X_train,
y_train,
X_val,
y_val,
learning_rate=0.1,
batch_size=64,
n_epochs=30,
)Final train accuracy: 0.916 Final val accuracy: 0.825 Final val loss: 0.6552
The model learns quickly and reaches high accuracy. Notice how the training and validation losses move together, indicating the model is not badly overfitting with this batch size and learning rate.
Comparing Batch Sizes
One of the most illuminating experiments is to compare how different batch sizes affect the training trajectory.
batch_sizes = [1, 32, 256, len(X_train)]
batch_results = {}
for bs in batch_sizes:
actual_bs = min(bs, len(X_train))
_, _, tl, vl, ta, va = sgd_logistic(
X_train,
y_train,
X_val,
y_val,
learning_rate=0.1,
batch_size=actual_bs,
n_epochs=30,
seed=42,
)
batch_results[actual_bs] = {
"train_loss": tl,
"val_loss": vl,
"train_acc": ta,
"val_acc": va,
}
Under this fixed 30-epoch budget, batch size 1 makes enough updates to reduce loss quickly and reaches strong validation accuracy first. Batch size 32 improves more gradually and eventually approaches the same accuracy range. The larger batches make far fewer updates per epoch, so batch size 256 and full-batch training remain close to chance accuracy in this run. The smooth epoch-level curve for batch size 1 reflects averaging across all of its noisy per-example updates.
Comparing Learning Rate Schedules
Next, let us compare different learning rate schedules on the same classification task.
def sgd_with_schedule(
X_train,
y_train,
X_val,
y_val,
schedule="constant",
eta0=0.1,
n_epochs=40,
batch_size=64,
decay_factor=0.5,
step_size=10, # for step decay
T_warmup=100,
seed=42, # for warmup
):
"""SGD with pluggable learning rate schedules."""
np.random.seed(seed)
n, d = X_train.shape
w = np.random.randn(d).astype(np.float32) * 0.01
b = np.float32(0.0)
n_batches_per_epoch = int(np.ceil(n / batch_size))
total_steps = n_epochs * n_batches_per_epoch
val_losses, val_accs, lr_history = [], [], []
step = 0
for epoch in range(n_epochs):
perm = np.random.permutation(n)
X_s = X_train[perm]
y_s = y_train[perm]
for start in range(0, n, batch_size):
# Compute learning rate for this step
if schedule == "constant":
lr = eta0
elif schedule == "step":
drops = epoch // step_size
lr = eta0 * (decay_factor**drops)
elif schedule == "cosine":
lr = eta0 * 0.5 * (1 + np.cos(np.pi * step / total_steps))
elif schedule == "warmup_cosine":
if step < T_warmup:
lr = eta0 * step / max(1, T_warmup)
else:
adjusted = step - T_warmup
adjusted_total = total_steps - T_warmup
lr = (
eta0
* 0.5
* (1 + np.cos(np.pi * adjusted / adjusted_total))
)
else:
lr = eta0
lr_history.append(lr)
X_batch = X_s[start : start + batch_size]
y_batch = y_s[start : start + batch_size]
logits = X_batch @ w + b
probs = sigmoid(logits)
error = probs - y_batch
w -= lr * (X_batch.T @ error / len(y_batch))
b -= lr * np.mean(error)
step += 1
val_probs = sigmoid(X_val @ w + b)
val_losses.append(binary_cross_entropy(val_probs, y_val))
val_accs.append(compute_accuracy(X_val, y_val, w, b))
return val_losses, val_accs, lr_history
schedules = ["constant", "step", "cosine", "warmup_cosine"]
schedule_results = {}
for sched in schedules:
vl, va, lrh = sgd_with_schedule(
X_train,
y_train,
X_val,
y_val,
schedule=sched,
eta0=0.1,
n_epochs=40,
batch_size=64,
seed=42,
)
schedule_results[sched] = {"val_loss": vl, "val_acc": va, "lr_history": lrh}
The learning-rate schedules produce visibly different update profiles. In this run, the constant rate continues reducing validation loss and finishes lowest. Cosine and warmup-cosine converge to nearly the same intermediate loss, while step decay finishes slightly higher. This result is specific to the shared training budget and illustrates why a schedule must be evaluated on the task rather than assumed to outperform a constant rate.
Visualizing the Loss Landscape
To develop intuition for SGD's path, we can visualize a low-dimensional slice of the loss surface around the trained solution.
# Create a 2D slice of the loss surface around the trained weights
# We perturb along two random directions and evaluate the loss
def plot_loss_landscape(w_center, b_center, X, y, n_grid=40, scale=2.0, seed=0):
"""Compute a 2D loss landscape slice around (w_center, b_center)."""
rng = np.random.RandomState(seed)
# Two random perturbation directions in weight space
d1 = rng.randn(len(w_center)).astype(np.float32)
d1 /= np.linalg.norm(d1)
d2 = rng.randn(len(w_center)).astype(np.float32)
d2 -= d2.dot(d1) * d1 # orthogonalize
d2 /= np.linalg.norm(d2)
alphas = np.linspace(-scale, scale, n_grid)
betas = np.linspace(-scale, scale, n_grid)
Z = np.zeros((n_grid, n_grid))
for i, a in enumerate(alphas):
for j, bv in enumerate(betas):
w_perturbed = w_center + a * d1 + bv * d2
probs = sigmoid(X @ w_perturbed + b_center)
Z[i, j] = binary_cross_entropy(probs, y)
return alphas, betas, Z
# Use a random subset for speed (cap at available validation samples)
rng = np.random.RandomState(42)
n_subset = min(300, len(X_val))
subset_idx = rng.choice(len(X_val), size=n_subset, replace=False)
X_sub = X_val[subset_idx]
y_sub = y_val[subset_idx]
alphas, betas, Z = plot_loss_landscape(
w_sgd, b_sgd, X_sub, y_sub, n_grid=40, scale=1.5
)
The broad color bands indicate low curvature along these sampled directions. Their steady slope also shows that the held-out validation subset would continue improving away from the training solution along this particular slice, so the center should not be interpreted as the validation-loss minimum.