Train Multiclass Classifier with Early Stopping

Implement train_with_early_stopping: train a linear classifier on the training set, use validation loss to choose a checkpoint, and return the best saved weights. Training and validation have different jobs: only training examples contribute to gradients.

Inputs and output

• x_train: (N_train, d), training features.
• y_train: (N_train,), integer-valued labels in [0, C), delivered as floats.
• x_val: (N_val, d), validation features.
• y_val: (N_val,), validation labels with the same class range.
• w0: (d, C), the supplied initial weights.
• lr: a finite, nonnegative learning rate.
• max_epochs: a nonnegative integer, the maximum number of full-batch updates.
• patience: a positive integer, the allowed number of consecutive updates without a new best validation loss.

Both datasets are nonempty, d is positive, and C is at least two. Shapes are consistent. Floating-point tensors are finite and share a dtype and device. Cast labels to integers when indexing. There is no separate bias; a supplied column of ones can act as one.

Return the selected weight matrix of shape (d, C), preserving dtype and device. Leave all inputs unchanged. If max_epochs=0, return w0's values. Otherwise, select among the weights after updates; do not evaluate w0 as an initial checkpoint.

One training update

Start from w=w0. At each epoch, compute class scores and probabilities from the current weights, then take a full-batch gradient step:

logits = x_train @ w
probs = softmax(logits, axis=-1)
grad = x_train.T @ (probs - one_hot(y_train, C)) / N_train
w = w - lr * grad

This is the gradient of mean training cross-entropy. Average over training examples once, not over classes or features. Use the updated weights on the next epoch, without shuffling, mini-batches, regularization or random initialization.

Validation loss

Immediately after each update, evaluate mean cross-entropy on the validation set using those updated weights. For each row of validation logits, compute stable log probabilities directly:

scores = x_val @ w
shifted = scores - row_max(scores)
log_probs = shifted - log(row_sum(exp(shifted)))
val_loss = -sum(log_probs[i, y_val[i]] for i in range(N_val)) / N_val

Keep the reduced axis when subtracting row maxima and row sums. Use the same subtract-max approach for training softmax. Do not add an epsilon to class probabilities before taking a log: that changes the loss and can make distinct large losses appear equal.

Saving and stopping

Initialize best_loss to infinity and no_improve to zero. After each validation:

• If val_loss is strictly less than best_loss, save these weights, record the new loss, and reset no_improve to zero.
• Otherwise, increment no_improve. Stop when it reaches patience.

Compare against the best loss so far, not just the preceding epoch's loss. A tie is not an improvement, and there is no minimum-improvement tolerance. The first updated checkpoint becomes the initial best. Save an independent copy where arrays are mutable, so later updates cannot overwrite it.

Stop after at most max_epochs updates. Whether patience or the epoch limit ends training, return the saved best weights rather than the latest weights. With lr=0, the weights stay unchanged: the first validation sets the best loss, and the next patience validations trigger stopping if the epoch limit allows them.

Example of the stopping rule

Suppose losses after successive updates are 0.8, 0.6, 0.6, 0.7 and patience=2. Update two is best. Update three ties the best, so the counter becomes one. Update four is worse, so the counter becomes two. Stop and return the weights from update two. These losses illustrate the rule independently of the training data.