Boosting that knows what it doesn’t know, read closely

Ask XGBoost or LightGBM to predict tomorrow's return, or a patient's blood pressure, and you get one number. That number is an answer to the question “what's your best guess?” — but it silently throws away the answer to a question that often matters more: how sure are you? A guess of 0.5% that could plausibly be anything from −3% to +4% is a different object than a guess of 0.5% pinned to within a tenth of a percent. In finance you size a position by that spread; in medicine you decide whether to run another test by it. The spread is not a footnote to the prediction. It is half the prediction.

Probabilistic prediction means the model outputs a full distribution over the outcome — here a Normal with a mean μ(x) and a standard deviation σ(x) that both depend on the input — instead of a bare point. Neural nets have a few ways to do this (Bayesian layers, deep ensembles, Monte-Carlo dropout). NGBoost's pitch is to get it from the model that actually wins on tabular data: boosted trees.

The idea is disarmingly simple. A Normal distribution has two parameters. So run two boosters at once — one that builds up the mean μ(x), one that builds up the (log) spread s(x) = log σ(x) — and at every step nudge both in the direction that makes the observed data more likely. “More likely” has a precise meaning: minimise the negative log-likelihood (NLL), the penalty a Normal pays for putting little probability where the real value landed. It is a proper scoring rule: it is minimised, in expectation, only by the true distribution, so it rewards being honest about the spread — too narrow and an outlier costs you dearly; too wide and you pay a standing tax on every point.

We parameterise the Normal as (μ, s) with s = log σ — the paper's choice, and a good one: it keeps σ positive for free and disentangles “where” from “how uncertain.” Here is the loss and its gradient, exactly as the site runs it (the TypeScript is what draws every chart below; Python and C++ are line-for-line translations):

import numpy as np

# The Normal is parameterised as (mu, s) with s = log(sigma), so the two
# "channels" the booster steers are the mean and the log of the spread.
def nll(y, mu, s):                    # negative log-likelihood = the log score
    inv = np.exp(-2 * s)              # 1 / sigma^2
    return s + 0.5 * np.log(2 * np.pi) + 0.5 * (y - mu) ** 2 * inv

def grad(y, mu, s):                   # ordinary gradient of the NLL
    inv = np.exp(-2 * s)
    return -(y - mu) * inv, 1 - (y - mu) ** 2 * inv

def fisher(s):                        # Fisher information of N(mu, e^s): diag(1/sigma^2, 2)
    return np.exp(-2 * s), 2.0

def natural_grad(y, mu, s):           # I^{-1} @ grad
    g_mu, g_s = grad(y, mu, s)
    i_mu, i_s = fisher(s)
    return g_mu / i_mu, g_s / i_s     # mu-channel collapses to -(y - mu): no sigma at all

To see why this is worth the trouble, we need heteroscedastic data — a fancy word for “the noise level changes across the input.” Our toy: x in [0, 1], a true mean m(x) = sin(2πx), and a noise scale that starts tiny on the left and grows about 11× toward the right, σ(x) = 0.05 + 0.5x. A point predictor can only ever draw the wiggle. NGBoost draws the wiggle and the fog around it:

Now the part a skim misses. You would think the “natural gradient” is a performance flourish — a faster optimiser, nice-to-have. It is not. It is the thing that makes the two-parameter boost trainable at all, and the reason is visible right in the gradient above. Look at the μ-channel of the ordinary gradient:

∂NLL / ∂μ = −(y − μ) / σ²

That 1/σ² is poison for boosting. Where the data is quiet (σ small), it multiplies the mean-error by a huge number; where the data is noisy (σ large), it shrinks it to nothing. So a single regression tree fit to this gradient spends all its capacity chasing the low-noise region and effectively ignores the high-noise one — and a single global learning rate cannot be right for both at once. The two parameters are measured in incompatible units, and the boost lurches.

The natural gradient is the fix, and it is not a hack — it is the geometrically correct notion of “steepest descent” when your parameters live on a curved statistical manifold rather than flat Euclidean space. You get it by pre-multiplying the ordinary gradient by the inverse Fisher information I(θ) — the matrix that measures how much the distribution actually moves when you jiggle each parameter. For a Normal in the (μ, s) chart, the Fisher information is beautifully simple, I = diag(1/σ², 2), and inverting it does something almost magical to the mean channel:

σ² · (−(y − μ) / σ²) = −(y − μ)

The σ cancels exactly. The natural-gradient mean update is just the plain residual −(y − μ) — the same target ordinary regression boosting already uses, with no scale distortion anywhere. Every point, quiet or noisy, contributes on equal footing. That single cancellation is the whole paper in one line, and it is why NGBoost's loop is a one-word switch away from the broken version:

def ngboost(xs, ys, rounds, lr, natural=True):
    mu = np.full_like(ys, ys.mean())              # init: the marginal MLE
    s  = np.full_like(ys, 0.5 * np.log(ys.var()))
    trees_mu, trees_s, rhos = [], [], []
    for _ in range(rounds):
        step = natural_grad if natural else grad  # <-- the entire experiment is this switch
        g_mu, g_s = np.array([step(y, m, v) for y, m, v in zip(ys, mu, s)]).T
        t_mu, t_s = fit_tree(xs, g_mu), fit_tree(xs, g_s)   # a regression tree per parameter
        f_mu, f_s = t_mu.predict(xs), t_s.predict(xs)
        rho = line_search(ys, mu, s, f_mu, f_s)   # scale that minimises the total NLL
        mu -= lr * rho * f_mu                      # ...then shrink by the learning rate
        s  -= lr * rho * f_s
        trees_mu.append(t_mu); trees_s.append(t_s); rhos.append(rho)
    return trees_mu, trees_s, rhos

We can run exactly that switch. Same data, same trees, same learning rate, same line search — only natural flipped. Here is the training NLL (lower is better) for both:

The mean curve is the easy half — any regressor draws a wiggle. The claim worth checking is the spread: does the second booster recover σ(x), the thing a point model can't represent? We plot the learned σ̂(x) against the truth:

On real benchmarks the paper reports this pays off on the metric that scores the whole distribution — held-out NLL — where it is competitive with far heavier machinery like deep ensembles and Monte-Carlo dropout, while being a plain boosted-tree model you can fit in a line. Its Table 1 (numbers are the paper's, lower is better):

Here is the caveat the error bars themselves won't tell you. NGBoost's interval is only as honest as the distribution family you assumed. We assumed a Normal. Real financial returns, insurance losses, and plenty of medical measurements have fat tails — extreme events far more often than a Normal predicts. Feed NGBoost data with the same variance but heavier tails, and it will happily learn a σ̂(x), print a crisp 95% interval, and be quietly wrong about exactly the events you most need it to get right.

We can measure it. We keep the mean and the noise scale σ(x) identical, and only swap Gaussian noise for Laplace noise — same variance, heavier tails. Both models fit fine. The central 95% interval barely notices: 96.4% coverage on Gaussian, 95.1% on Laplace, near enough the nominal 95% that a dashboard would call it calibrated. But walk out to the tail — the fraction of points that land beyond ±2.5σ̂, where a Normal expects a mere 1.24% — and the mismatch explodes:

• Fit a Student-t head. NGBoost is distribution-agnostic; the whole demo above would run with a heavy-tailed t in place of the Normal, with its own Fisher information. Does letting the model learn the tail thickness close the 4× gap in Fig 4? That is the honest fix, and it is a few lines.
• Watch the ordinary gradient actually diverge. Our line search tames the ordinary-gradient run into merely-worse rather than unstable. Remove the line search, or crank the learning rate, and the 1/σ² blow-up should turn Fig 2's orange curve from lagging into oscillating — the failure the natural gradient prevents, made loud.
• Push the spread into 2-D. The mean booster here is the same axis-aligned tree the model pages train. Growing both the μ- and σ-trees on a 2-D dataset would draw an uncertainty surface — high where the data is sparse or noisy — which is the version of this that a decision-maker can actually read off a map.

The mean channel in Fig 1 is doing the same job as the boosters on the XGBoost and LightGBM pages — go watch one fit a boundary, then picture a second copy running beside it, fitting not the answer but the model's own doubt. That second copy is the whole idea; the natural gradient is what keeps it standing.