Optimization

Overview

Optimization is the study of how to find the best choice from many possible ones, where "best" is defined by a number you want to make as small or as large as possible. In machine learning, that number is the loss, the choice is the set of model parameters, and the workhorse algorithm is gradient descent and its relatives. Training a model is nothing more, and nothing less, than running an optimizer on a loss.

For tabular models the story has two layers. Classical models such as linear regression and gradient-boosted trees are fitted by well-understood optimization procedures. Neural tabular models, including foundation models pre-trained on huge collections of tables, are fitted by stochastic gradient methods over billions of steps. Choosing the optimizer, the learning rate and the schedule decides whether training succeeds. Related pages: Linear algebra for the arithmetic, Statistics for why the loss looks the way it does, and /model/architecture for the model being trained.

What it is

An optimization problem has three parts.

• Variables are the numbers you may change, such as the weights of a network.
• An objective is the score to minimize, such as prediction error plus a penalty.
• Constraints are rules the variables must obey, such as weights staying non-negative. Many machine learning problems have none.

The solution is the setting of the variables that gives the lowest objective value.

Loss, cost and objective

These words are often used loosely. A loss measures the error on one example. A cost or objective averages the loss over a dataset and may add a penalty. The optimizer sees the objective and tries to push it down.

Common losses for tables:

A landscape analogy

Picture the objective as a hilly landscape, where location is the parameter setting and height is the loss. Optimization is walking downhill in fog. You cannot see the valley, but you can feel the slope under your feet. Gradient descent says: step in the direction that goes down most steeply, then check again.

Convex and non-convex problems

In a convex problem the landscape is one smooth bowl, so every downhill walk reaches the same lowest point. Linear regression and logistic regression with a standard penalty are convex. In a non-convex problem, such as a deep network, the landscape has many valleys, ridges and flat plains. Training still works well in practice, partly because large networks have many good valleys and partly because the noise in stochastic updates helps escape poor spots.

Where it came from

Optimization has deep roots in calculus: setting the derivative to zero finds the flat points of a function. The idea of repeatedly stepping downhill dates back to the nineteenth century. In the twentieth century, linear programming, the study of the best choice under linear constraints, became central to operations and planning, and brought in the simplex method and later interior-point methods. Convex optimization developed into a complete theory with strong guarantees.

For machine learning the key step was stochastic approximation: instead of computing the exact slope from all the data, use a small random piece of it. This idea, from the mid-twentieth century, is now the foundation of neural network training. The backpropagation algorithm, popularized in the 1980s, made it practical to compute slopes for networks with many layers.

Starting in the 2010s, adaptive methods arrived: AdaGrad, RMSProp and then Adam, which was introduced in a 2014 paper. Weight decay variants, warmup schedules and cosine decay followed, and large-scale training added gradient clipping, mixed precision and sharded optimizer states. Alongside, methods that need no gradients, such as Bayesian optimization, grew for tuning expensive black-box functions.

How it works

Step 1: Define the objective

Write down what you want to minimize, for example mean squared error over the training rows plus a small penalty on weight size. Everything else follows from this choice, so it should reflect what you actually care about.

Step 2: Compute the gradient

The gradient is the vector of partial derivatives of the objective with respect to every parameter. It points uphill, in the direction of the steepest increase. Its negative points downhill. In neural networks the gradient is computed by backpropagation, which applies the chain rule layer by layer. Frameworks such as PyTorch and JAX do this automatically.

Analogy: standing on a hillside with a ruler, you measure how much the ground drops if you shift one step east and one step north. Those two numbers form a tiny gradient.

Step 3: Take a step

The basic update is

parameters = parameters - learning_rate * gradient

The learning rate is the step size. Too small and training crawls. Too large and you overshoot the valley, bounce around or blow up.

Step 4: Use mini-batches

Computing the gradient on the full dataset is slow. Stochastic gradient descent uses a random mini-batch of rows for each step. The gradient is noisy but, on average, points the right way. Noise is not all bad: it acts as a mild regularizer and helps leave sharp, poor minima.

Analogy: asking a handful of customers for feedback each day, instead of the whole customer base. Each day's answer is rough, but the trend over weeks is reliable.

Step 5: Add momentum

Plain steps zigzag in narrow valleys. Momentum keeps a running average of past gradients and moves in that smoothed direction, like a ball rolling downhill that builds up speed along the valley floor and ignores small wiggles.

Step 6: Adapt the step per parameter

Adaptive methods give each parameter its own effective step size, based on the history of its gradients. Adam combines momentum with a running estimate of each parameter's gradient size. Parameters with consistently large gradients get smaller steps, and rare ones get larger steps. The result is robust training with less tuning, which is why it is the default for many neural networks.

Step 7: Schedule the learning rate

The best step size changes during training. A common recipe is:

1. Warmup. Start tiny and ramp up over the first few hundred or thousand steps, so early noisy gradients do not wreck the weights.
2. Hold or decay. Keep the rate high for a while, then lower it.
3. Cosine or linear decay. Shrink smoothly toward zero, so the final steps settle into a valley.

Step 8: Control the size of updates

Gradient clipping limits the size of a step when the gradient explodes. Weight decay pulls weights gently toward zero and acts as a regularizer. Both are small additions that prevent many failures.

Step 9: Stop at the right time

Stop when validation error stops improving, a practice called early stopping. Training error will keep falling, but past some point the model is memorizing noise.

Why scaling of inputs matters

Imagine a valley that is very long in one direction and very narrow in another. Gradient descent bounces between the steep walls and creeps along the floor. This happens when columns have very different scales, for example age in tens and income in hundreds of thousands. Standardizing columns makes the valley rounder, and the optimizer reaches the bottom in far fewer steps.

Regularization as part of the objective

Adding a penalty to the objective changes the answer the optimizer finds.

• L2 penalty adds the sum of squared weights. It shrinks all weights smoothly.
• L1 penalty adds the sum of absolute weights. It pushes many weights to exactly zero, giving sparse models.
• Elastic net mixes the two.
• Early stopping acts like an implicit penalty by cutting the path short.
• Dropout randomly silences units, which acts like averaging many smaller networks.

Constrained problems in one minute

Some problems restrict the variables, for example probabilities must be non-negative and sum to one. Common remedies are to reparameterize (use a softmax so the outputs automatically obey the rule), to project back onto the allowed region after each step, or to add a penalty for violating the rule. Reparameterizing is the simplest and is used all the time in neural networks.

Key concepts and terms

Comparing common optimizers

A worked example

We will minimize a simple function by hand, then train a tiny regression model with mini-batches.

Hand calculation

Minimize f(w) = (w - 3)^2. The derivative is 2 (w - 3). Start at w = 0 with a learning rate of 0.25.

Each step halves the remaining distance to 3. With a rate of 1.0 the update would jump to 6 and then back to 0 forever. With a rate above 1.0 it would diverge. This tiny example shows the whole learning-rate story.

Gradient descent for linear regression

import numpy as np

rng = np.random.default_rng(0)
n, d = 500, 3
X = rng.normal(size=(n, d))
w_true = np.array([2.0, -1.0, 0.5])
y = X @ w_true + 0.3 * rng.normal(size=n)

w = np.zeros(d)
lr = 0.1
batch = 32

for epoch in range(30):
    idx = rng.permutation(n)
    for start in range(0, n, batch):
        b = idx[start:start + batch]
        err = X[b] @ w - y[b]
        grad = 2 * X[b].T @ err / len(b)
        w -= lr * grad
    if epoch % 10 == 0:
        loss = np.mean((X @ w - y) ** 2)
        print(epoch, round(float(loss), 4))

print(w.round(3))

The loss drops quickly and the weights approach the true ones. Because the mini-batches are noisy, the last digits wobble, which is why a decaying learning rate helps at the end.

A neural network with a modern recipe

import torch
from torch import nn

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 1))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
steps = 1000
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=3e-3, total_steps=steps)

X = torch.randn(2048, 20)
y = (X[:, :3].sum(dim=1, keepdim=True) + 0.1 * torch.randn(2048, 1))

for step in range(steps):
    idx = torch.randint(0, X.shape[0], (128,))
    pred = model(X[idx])
    loss = nn.functional.mse_loss(pred, y[idx])
    opt.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    sched.step()

Note the ingredients: an adaptive optimizer with decoupled weight decay, a schedule with warmup and decay, mini-batches and gradient clipping. This pattern, scaled up, is how large models are trained.

One step of Adam, written out

import numpy as np

def adam_step(w, g, m, v, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
    m = b1 * m + (1 - b1) * g           # running mean of gradients
    v = b2 * v + (1 - b2) * g * g       # running mean of squared gradients
    m_hat = m / (1 - b1 ** t)           # correct the early bias
    v_hat = v / (1 - b2 ** t)
    w = w - lr * m_hat / (np.sqrt(v_hat) + eps)
    return w, m, v

The division by the square root of v_hat is the adaptive part. A parameter whose gradients are always large gets divided by a large number, so its steps stay modest. A parameter with tiny gradients gets a boost.

A learning rate schedule in code

import math

def lr_at(step, total, base=1e-3, warmup=100):
    if step < warmup:
        return base * step / warmup
    progress = (step - warmup) / max(1, total - warmup)
    return 0.5 * base * (1 + math.cos(math.pi * progress))

for s in [0, 50, 100, 500, 900, 1000]:
    print(s, round(lr_at(s, 1000), 6))

The rate climbs linearly during warmup, peaks, then glides down along a half cosine to nearly zero.

Reading a loss curve

Reading a loss curve

Strengths and limitations

Strengths

• General. One framework fits regression, classification, ranking, generation and many other tasks.
• Scalable. Stochastic methods train on datasets far larger than memory.
• Automatic. Combined with automatic differentiation, almost any differentiable model can be optimized without hand-derived formulas.
• Well studied. Convex problems have strong guarantees, and many practical tricks are documented.
• Flexible objectives. You can combine accuracy, penalties and constraints in one score.

Limitations

• Local minima and plateaus. Non-convex landscapes offer no guarantee of the global best.
• Hyperparameter sensitivity. Learning rate, batch size and schedule can make or break a run.
• Cost. Training a large model needs many steps, memory and energy.
• Mismatch between loss and goal. What you minimize may differ from what you care about, such as accuracy versus cross-entropy.
• Instability. Exploding gradients, vanishing gradients and low-precision overflow can derail training.
• Poor fit for non-differentiable steps. Trees and discrete choices need other methods.

Bias in the optimizer

Even when two models have the same capacity, different optimizers can land in different valleys with different generalization. This is called implicit bias. It is one reason that the same network trained with SGD and with Adam can behave differently on new data, and why reporting the optimizer and schedule is part of honest benchmarking. See /model/benchmarks.

How it compares with alternatives

Gradient boosting deserves a note. It is optimization in function space: each new tree is a step that reduces the loss, fitted to the negative gradient of the loss with respect to the current predictions. So boosted trees for tables and neural networks share the same optimization logic with different building blocks.

Role in tabular foundation models

A tabular foundation model is trained once, at great cost, on a vast supply of tables, and then used many times with no further training. Optimization is where that one-time cost is spent, and it also shapes what the model can do at use time. See /model/research and /model/architecture.

Pre-training as optimization over many tables

During pre-training, each step draws a batch of tables, often generated synthetically, and asks the network to predict held-out targets from the labeled rows in the same table. The loss is the average prediction error across tables. Minimizing it forces the network to learn a general strategy for reading a table and predicting from it. This is a form of meta-learning: the optimizer improves not a model for one dataset, but a procedure for learning from any dataset. See Meta-learning and Few-shot learning.

In-context learning replaces per-task optimization

At use time, there is no gradient step. The labeled rows go in as context and the network answers in a forward pass. You can think of the forward pass as an optimization that has been amortized: the heavy search for good parameters was done during pre-training, so each new task needs only one pass. This is one reason such models are fast on small and medium tables. See TabPFN and Transformers.

Row and column structure shapes the loss

Targets are masked at random, so the loss is computed only on the rows whose labels were hidden. Because tables vary in number of rows, columns and classes, the batches are ragged, and the loss needs careful normalization. The optimizer must also cope with big differences in gradient size between easy and hard tables, which is one reason adaptive optimizers and gradient clipping are standard.

Mixed types and missing values

Different column types create different loss terms. A numeric target may use a distributional loss such as binned cross-entropy over value ranges, which behaves better than squared error when targets span many scales. A categorical target uses cross-entropy. Missing cells are randomly hidden during pre-training and the model learns to predict or ignore them. Balancing these terms is a weighting problem within the objective.

No column names

Because pre-training tables are often generated with random, meaningless column order and labels, the optimizer drives the network to rely on values and relationships rather than names. This makes the learned strategy robust to cryptic headers. When names are available at use time, they act as a small extra signal rather than a requirement.

Imputation and anomaly detection

Imputation is trained by hiding cells and minimizing the error of reconstructing them. This is the same masked-prediction objective used in Self-supervised learning. Anomaly detection can reuse the same trained predictor: rows whose cells are poorly predicted by the rest get high scores. See Imputation networks and Anomaly detection.

Continual learning

Updating a model with new data risks overwriting what it already knows. Optimization offers several controls: a small learning rate for updates, replay of old examples mixed into new batches, penalties that discourage moving important weights, and parameter-efficient updates such as LoRA that change only a small low-rank part. Each is a way of shaping the optimization problem so it respects the past.

Sector and semantic classification

When a model labels a dataset or a column by sector or meaning, it is trained with a classification loss on embeddings. Class imbalance, label noise and many similar categories make the loss landscape tricky, so label smoothing, class weights and careful learning rate schedules are common.

Systems side of optimization

Training at scale is an engineering problem as much as a mathematical one. Mixed precision halves memory, gradient checkpointing trades compute for memory, and data parallelism splits batches over devices. See Mixed precision, Gradient checkpointing and CUDA.

Optimization around the model

Hyperparameters of the foundation model, such as depth, width, learning rate and the mix of synthetic generators, are themselves chosen by an outer search. Sample-efficient tools such as Bayesian optimization help here. For end users, the same tools can tune the few settings that remain at inference time, such as the number of ensemble members.

Compute budgeting

Practical guidance and best practices

A sensible default recipe

• Start with AdamW, a learning rate around 1e-3 for small networks and lower for large ones, and weight decay around 1e-2.
• Use a warmup of a few percent of total steps followed by cosine or linear decay.
• Clip the global gradient norm to about 1.
• Choose the largest batch size that fits, then scale the learning rate with it.
• Standardize numeric inputs and check for extreme outliers.
• Track both training and validation loss every few hundred steps.

Tuning order

1. Get a tiny model to overfit a tiny batch. If it cannot, there is a bug.
2. Find a good learning rate with a short sweep over a log-spaced range.
3. Set the schedule length to match your compute budget.
4. Tune regularization: weight decay, dropout, early stopping.
5. Only then explore architecture.

Hyperparameter search

For tables specifically

• Use trees as a baseline before training any neural network. See XGBoost.
• Scale and, where useful, transform skewed columns before a neural model.
• Use early stopping on a validation fold that matches how the model will be used.
• When data is small, consider a pre-trained tabular model instead of training from scratch.

Stability checklist

• Check for NaN or infinite values in inputs.
• Use a stable form of softmax and cross-entropy.
• Keep an eye on the largest activation values in lower precision.
• Save checkpoints often so a crash costs minutes, not days.
• Log the learning rate, gradient norm and loss together.

Debugging recipe when training does not work

1. Print the loss on the very first step. For classification with k classes, it should be close to the log of k.
2. Overfit one small batch. If the loss does not go near zero, look for a bug in the model, the loss or the labels.
3. Disable all regularization and augmentation and retry.
4. Check that inputs are scaled and free of NaN.
5. Compare against a tiny linear baseline.
6. Turn features back on one at a time.
7. Inspect gradient norms per layer for explosions or vanishing.

Common mistakes

1. Tuning on the test set. Use a separate validation set for choices, and open the test set once.
2. Learning rate far off. It is the single most common cause of failure. Sweep it on a log scale.
3. No warmup for large models. Early updates can destroy useful structure.
4. Forgetting to shuffle. Ordered data produces biased gradients.
5. Putting weight decay on everything. Biases and normalization scales are usually exempt.
6. Confusing loss with metric. A lower loss does not always mean a better metric. Track the metric you care about.
7. Not resetting gradients. Accumulating gradients by accident is a classic bug.
8. Ignoring scale of inputs. Columns with huge ranges slow or break training.
9. Stopping too early or too late. Use validation curves, not hope.
10. Comparing optimizers with one seed. Differences between runs can exceed differences between methods.
11. Overlooking leakage. An optimizer is good at exploiting any shortcut in the data, including the ones you did not mean to include.
12. Using big models on tiny tables. The optimizer will memorize. A pre-trained model or a regularized simple model is safer.

Frequently asked questions

What is optimization in machine learning?

It is the process of adjusting a model's parameters to minimize a loss that measures error on the training data. The optimizer repeatedly computes how the loss changes with each parameter and moves the parameters in a direction that lowers it. Training a neural network is running an optimizer over many such steps.

What is gradient descent and how does it work?

Gradient descent computes the slope of the loss with respect to every parameter and then moves each parameter a small distance in the opposite direction of the slope. Repeating this walks the parameters downhill toward lower loss. The size of each step is controlled by the learning rate, which must be neither too small nor too large.

Why use mini-batches instead of the full dataset?

Mini-batches make each step cheap, so you can take many more steps in the same time and train on data that does not fit in memory. The gradient from a batch is noisy but unbiased on average. The noise can even help generalization by keeping the optimizer from settling into sharp, fragile minima.

What is the difference between SGD and Adam?

SGD uses one global step size and, optionally, momentum. Adam keeps running estimates of each parameter's gradient mean and size and scales the step accordingly, so parameters with small or rare gradients can still move. Adam usually works well with little tuning, while SGD with a good schedule can generalize slightly better on some tasks.

How do I choose a learning rate?

Run a short sweep over a log-spaced range and look for the largest rate at which loss still falls smoothly. Then use warmup and a decaying schedule. If loss spikes or turns to NaN, lower the rate or add gradient clipping. The best rate also depends on batch size and optimizer.

What are local minima, and do they matter?

A local minimum is a valley that is not the lowest in the landscape. In very large neural networks, most local minima turn out to be nearly as good as the best ones, and the bigger risks are flat plateaus and saddle points. In small or hand-built models, a poor local minimum can matter more, so restarts or better initialization help.

How does optimization relate to tabular foundation models?

Their pre-training is a large-scale optimization over many tables, with the loss measuring prediction error on hidden targets. This teaches the network a general procedure for learning from a table. At use time the optimization has already been done, so new tasks are handled by conditioning on context rows in a single pass instead of running a fresh training loop.

Is gradient boosting a kind of optimization?

Yes. Gradient boosting performs gradient descent in the space of functions. Each new tree is fitted to the negative gradient of the loss with respect to the current predictions, then added with a small step size, the learning rate. That is why boosting libraries share concepts such as shrinkage, regularization and early stopping with neural network training.

When should I use second-order methods?

Use them on small to medium, smooth, mostly deterministic problems, such as fitting a logistic regression or a small statistical model, where curvature information speeds up convergence a lot. They are rarely used for large networks trained with noisy mini-batches, because the cost and memory of curvature outweigh the gain.

How do I know training has converged?

Watch the validation loss and the metric you care about. If they stop improving over a patience window while the learning rate has decayed, you have most likely converged. A flat training loss alone is not enough, since it can be flat because of a learning rate that is too small.

Glossary

• AdamW: Adam with weight decay applied separately from the gradient update.
• Backpropagation: applying the chain rule backward through a network to compute gradients.
• Batch size: the number of rows used for one update.
• Convex function: a function with a single bowl-shaped landscape.
• Early stopping: halting training when validation error stops improving.
• Epoch: one full pass through the training data.
• Gradient: the vector of slopes of the loss with respect to each parameter.
• Gradient clipping: limiting the size of a gradient step.
• Hyperparameter: a setting chosen before or outside training.
• Learning rate: the size of each update step.
• Local minimum: a point lower than its neighbors but not the lowest overall.
• Loss function: a measure of error on an example.
• Momentum: a running average of past gradients used to smooth updates.
• Objective: the quantity being minimized, usually loss plus penalty.
• Regularization: any constraint or penalty that reduces overfitting.
• Saddle point: a flat point that slopes up in some directions and down in others.
• Schedule: a rule for changing the learning rate over time.
• Warmup: a gradual increase of the learning rate at the start.
• Weight decay: a pull that shrinks weights toward zero.

Summary

Optimization turns a goal into a number and finds the parameters that make it small. Gradient descent, stochastic mini-batches, momentum, adaptive scaling, careful schedules, clipping and early stopping form a reliable toolkit. Most training problems come down to the learning rate, the scale of the inputs and the match between the loss and the real goal.

For tabular foundation models, optimization is the large, one-time investment that makes later use cheap. Pre-training minimizes prediction error on hidden targets across a huge variety of tables, which teaches a general procedure for reading rows and columns, handling mixed types and missing values and ignoring column names. At use time, in-context learning replaces the per-task training loop, while optimization continues to matter for continual updates, hyperparameter search and the engineering of large runs. Related pages: Variational inference, Bayesian deep learning, Statistics and /model/benchmarks.

Related

It belongs to Foundations in my stack, under AI and machine learning. Others in the same group: