traceot · research, unvalidated

Learn ADOPT: regression when each answer is a whole distribution

A plain-language guide to optimal-transport regression, the additive transport idea behind ADOPT, and what happened when we re-implemented it from the paper.

1. What optimal-transport regression is

Ordinary regression predicts one number per row: an average. Sometimes the thing you care about is a whole distribution: how spread out values are, how heavy the right tail is, not just where the middle sits. Distributional regression predicts the whole shape.

To compare and average distributions you need a distance between them. Optimal transport (OT) gives a natural one: the cheapest way to move the probability mass of one distribution so that it becomes the other. With squared-distance cost this is the Wasserstein-2 distance, written W2.

On a single number line OT is especially simple. The cheapest plan moves the 10th percentile of the first distribution to the 10th percentile of the second, the median to the median, and so on. So the transport map is "percentile to same percentile", and it is always increasing.

This gives a handy representation: describe each distribution by its quantile function (the value at each percentile). Then the W2 distance is just the ordinary distance between quantile functions, and the average distribution (the Fréchet mean) is the percentile-by-percentile average.

OT regression predicts, for each set of inputs x, a whole distribution, and judges errors with W2. Instead of "x raises the average by 2", the answer reads "x pushes the distribution along this transport map".

2. How additive transport works

Additive models (like generalised additive models) are popular because each input gets its own curve you can plot. In ordinary regression the effects simply add: prediction = average + effect of x1 + effect of x2 + ….

Distributions cannot be added. ADOPT's idea is to replace "add an effect" with "apply a transport". Each input x_j gets its own transport map T_j(x_j). Start from the average distribution and apply the maps one after another:

prediction(x) = T_p(x_p) ∘ … ∘ T_2(x_2) ∘ T_1(x_1) applied to the mean distribution

In the paper this chaining is written with a "⊕" sign, which means "do this transport, then that one". In ordinary flat space, chaining shifts is the same as adding them, so ADOPT reduces to the usual additive model there.

Each map is glass-box: you can show how one input value moves, widens or reshapes the outcome distribution while holding the others at their typical level.

3. ADOPT's components, one by one

Fréchet mean

The "average" distribution: the one with the smallest total squared W2 distance to all observed responses. For 1-D distributions it is the percentile-wise mean of the quantile functions.

Transport maps T_j(x_j)

For each input there is a curve of distributions g_j(x_j). The map T_j(x_j) is the OT map from the average of that curve to its value at x_j. This is the distribution version of "subtract the mean of each effect" that makes additive models identifiable.

Partial transport residuals

To estimate one map, undo all the others. For row i, take the map from the mean to the observed distribution, then apply the inverse maps of every other input. What remains is "the part of this row explained by input j, plus noise". This mirrors partial residuals in classic backfitting. Order matters because transports need not commute, and the paper fixes the order.

Local Fréchet regression (the smoother)

The residuals for input j are smoothed against x_j with kernel-weighted local linear regression, done in W2 geometry. For 1-D distributions this is a weighted average of quantile functions, followed by a step that makes the result increasing again (the weights can be negative).

Transport centring

After smoothing, each map is re-anchored at the Fréchet mean of its own curve, so that "typical" inputs mean "no change".

Transport backfitting loop

Cycle through the inputs, re-estimating each map with the others held fixed, until the fitted distributions stop changing. The bandwidth of the smoother is chosen by 5-fold cross-validation.

Cost

Each cycle smooths n residual distributions at n points for each of p inputs: about n² · p · q operations for q percentiles (the paper's Supplement C says the same). That is fine for hundreds or thousands of distributions, and too slow for millions of rows. See section 5.

4. Our reproduction of the paper's simulations

There is no public code, so we re-implemented ADOPT from the paper: All computation runs in a Rust core. A separate plain-numpy version, used only by the tests, checks every Rust result (agreement within 1e-9, in practice about 1e-14). We then re-ran the paper's two simulation designs with fixed seeds and B = 200 replicates per setting. Every guess we had to make is listed in REPRODUCTION_NOTES.md.

The metric is MISE: the squared W2 (or log-Cholesky) error of the predicted distribution, averaged over the unit cube of inputs and over replicates. Lower is better. "Paper" columns are copied verbatim from the published tables. "Ours" are our runs.

Mean (standard error). *L2-additive is an extra comparator of ours (additive model on quantile functions treated as vectors), not part of the paper. We did not re-implement the ADR method of Han et al. (2020), so its column is "not run".

What matched and what did not

• Distributional responses, Case I (Beta family): the ranking ADOPT < GF < FM is reproduced, and ADOPT at n = 300 matches the paper almost exactly (0.179 vs 0.181 ×10⁻³). At n = 100 our ADOPT is about 25% higher than the paper's. Our GF is 15–25% higher and our FM about 5% higher. The FM gap is small but larger than its standard error, which points to a small difference in the true distributions or in how the MISE integral was approximated.
• Case II (Normal family): GF and FM land within about 1–11% of the paper. Our ADOPT error is 3–6 times lower than the paper's, and our L2-additive comparator is just as low. The most likely cause is bandwidth selection or another tuning detail the paper does not give. Here the true maps are pure shifts, which ADOPT can fit very well.
• SPD matrices (log-Cholesky): the ranking ADOPT < GF < FM is reproduced, but the levels are not. Only FM for Case I, m = 10 matches. Our ADOPT and GF errors are 10–50 times smaller than the paper's. A diagnostic run (not used in the table) shows that noise about ten times larger than the stated σ = 0.01 brings ADOPT into the paper's range for m = 10. It does not fix GF or FM, though. Under our reading of the design, FM grows in proportion to m, which the paper's FM column does not show. So at least part of the SPD design must differ from how we read it.

Bottom line: the method and its qualitative claims reproduce (ADOPT is best everywhere, and its error falls as n grows). The exact numbers in Table 1 are close for Case I and for the baselines. Our ADOPT in Case II and all of Table 2 do not match. We report this as is and did not tune anything towards the paper's values.

What we found

An earlier arXiv copy (2512.08118) was withdrawn on 15 December 2025. The notice says: "Not all co-authors had agreed in advance to the public release." That is a consent issue, not a stated technical correction. The paper appears in AISTATS 2026.

We checked our code for leaked test points, a wrong noise formula, a wrong error integral, and mixed-up seeds. The evaluation grid does not contain training rows. Responses are noisy in the way the paper writes, and the score compares predictions with the noise-free truth. One bug was real: while choosing the bandwidth, the range of each input was computed on the whole sample, including the fold held out of that check. The fit never saw the held-out responses. The tables above are after that fix. The gap with the paper is still there, so the bug was not the cause.

The method itself has limits the paper's algorithm does not remove. Details and numbers are in ADOPT_AUDIT.md.

• Order. Transport maps do not commute, so the order of the inputs changes the fit. Averaging every order removes that. On problems where the maps do commute, the average matches the original fit.
• Shared inputs. Two identical inputs make the algorithm put the whole effect on whichever one it visits first. Averaging the orders splits the effect evenly, and the prediction stays the same.
• The loop can walk uphill. On a stiff example the in-sample error rose at every cycle. A check that rejects an uphill step stops at the best cycle so far.
• Crossed quantiles. Adding curves in the tangent space can decrease a quantile function. Projecting back onto increasing functions removes the crossings. Composing the paper's increasing maps does not cross.
• Cost. Each cycle compares every row with every other row. Binning the inputs keeps the error on a smooth additive problem and grows much more slowly.
• Edges. Local-linear smoothing already handles the edge of the sample far better than a plain kernel average. Reflecting points through the endpoints is optional and, on a curved target, made the edge worse. It stays off.

5. From ADOPT to TRACE

ADOPT needs one whole distribution per row (for example, one per group or per period), and its smoother compares every row with every other row. TRACE is our own proposed method, not published and not validated. It keeps ADOPT's key idea and changes the setting:

• Same idea, unit level. On a number line the quantile function is the OT map from a uniform variable. So an additive model of the conditional quantile function, Q(u | x) = β₀(u) + Σ f_j(x_j, u), is an additive model of transport displacements. Each shape function says how x_j moves, widens or fattens the outcome distribution, read straight from individual rows rather than pre-aggregated distributions.
• EBM-style fitting. Shape functions are lookup tables over histogram bins, learned by slow cyclic boosting on the pinball (quantile) loss: linear passes over the data, no row-by-row comparisons.
• Built for large tables. The Rust core already has the scalable pieces: mergeable streaming quantile sketches, 8/16-bit histogram binning in columnar layout, multi-threaded gradient histograms, and pinball gradients that never store an n × K matrix.
• Explanations of change and of oddness. Planned: splitting a shift between periods or groups into composition, structure and unexplained parts, transport-based oddness scores, and leakage and placebo checks that run inside the fitting loop.

ADOPT stays in the package as a faithful baseline for row-level distributional data, and as a segment-level comparator for TRACE.

6. Source and citation

Song, W. and Müller, H.-G. (2026). "ADOPT: Additive Optimal Transport Regression." Proceedings of the 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:1558–1566. proceedings.mlr.press/v300/song26c.html. An arXiv preprint (2512.08118) was withdrawn on 15 December 2025 with the comment "Not all co-authors had agreed in advance to the public release". The conference paper is the version we reproduced.

Other background: Petersen and Müller (2019), Fréchet regression; Lin (2019), log-Cholesky metric; McCann (1997), displacement interpolation.