FullSupportBarDistribution.sample()/cdf()/icdf() disagree with forward() in outer buckets

[ppguo]

FullSupportBarDistribution.sample()/cdf()/icdf() disagree with forward() in outer buckets

Summary

FullSupportBarDistribution.forward() evaluates the first and last components as
half-normal tails with support on all of R, but cdf(), icdf(), and sample()
inherit the finite, within-bin-uniform implementation from BarDistribution.

Consequently, values returned by sample() are not distributed according to the
density returned by -forward(). The same inconsistency affects tail quantiles,
median(), and methods such as ucb() that call the inherited icdf().

This is still present in PriorLabs/TabPFN commit
930ea728b40f7f321d965250838bca8507bd70b9 (2026-08-20):

• FullSupportBarDistribution.forward() adds the half-normal tail density.
• FullSupportBarDistribution.sample() calls self.icdf(...).
• icdf() is inherited from BarDistribution and linearly interpolates between
  the two finite borders of every selected bucket, including both outer buckets.

Minimal reproduction

import torch

try:
    from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution
except ImportError:
    from tabpfn.model.bar_distribution import FullSupportBarDistribution

torch.manual_seed(7)
dist = FullSupportBarDistribution(torch.tensor([-2.0, -1.0, 1.0, 2.0]))
probs = torch.tensor([0.25, 0.50, 0.25])
logits = probs.log().repeat(20_000, 1)

samples = dist.sample(logits)
buckets = dist.map_to_bucket_idx(samples).clamp(0, dist.num_bars - 1)
reported_logq = -dist(logits, samples)
sample_logq = probs[buckets].log() - dist.bucket_widths[buckets].log()
edge = (buckets == 0) | (buckets == dist.num_bars - 1)

print(samples.min().item(), samples.max().item())
print(((samples < -2.0) | (samples > 2.0)).sum().item())
print((reported_logq[edge] - sample_logq[edge]).min().item())
print((reported_logq[edge] - sample_logq[edge]).max().item())

Observed:

sample range: approximately [-2, 2]
outside [-2, 2]: 0 / 20,000
reported_logq - sample_logq in edge buckets: [-0.847, -0.620]

The documented full-support construction puts half of each outer component's
mass beyond its finite outer border. With probabilities [0.25, 0.50, 0.25],
the expected outside fraction is therefore

0.5 * (0.25 + 0.25) = 0.25,

or about 5,000 of 20,000 samples, rather than zero.

The complete reproducible audit also compares empirical and analytic CDFs. It
obtains KS distance 0.00585 to the inherited finite ICDF distribution versus
0.125 to the density implemented by forward().

Expected behavior

For borders b[0] < ... < b[K] and component probabilities pi, the two outer
components should be sampled from the same half-normal densities used by
forward():

left:  Y = b[1]   - HalfNormal(scale_left)
right: Y = b[K-1] + HalfNormal(scale_right)

where each scale keeps half of the component mass inside its nominal outer bin:

scale_left  = (b[1] - b[0])     / HalfNormal(1).icdf(0.5)
scale_right = (b[K] - b[K - 1]) / HalfNormal(1).icdf(0.5)

For a uniform variate u that selects the left component, with conditional
rank r = u / pi[0], the inverse CDF should be

b[1] - HalfNormal(scale_left).icdf(1 - r).

For the right component it should be

b[K-1] + HalfNormal(scale_right).icdf(r).

Interior buckets should retain their current within-bin-uniform interpolation.

Requested fix

1. Override cdf() and icdf() in FullSupportBarDistribution with the
   half-normal formulas already used by forward().
2. Make sample() use that corrected inverse CDF, or sample the selected outer
   components directly from their half-normal distributions.
3. Preserve batch shape, device, and dtype; avoid rebuilding the result through
   torch.tensor(list_of_tensors), which can move CUDA results back to CPU.
4. Keep the existing borders, component probabilities, and half-normal scale
   convention so checkpoint likelihoods remain unchanged.
5. Audit inherited methods that depend on cdf()/icdf(), including
   quantile(), median(), ucb(), and border-probability translation.
6. State the affected releases and include the correction in the release notes.

Proposed regression tests

1. Deterministic inverse-CDF checkpoints. For borders [-2,-1,1,2] and
   probabilities [0.25,0.50,0.25], verify
   icdf(0.125) == -2 and icdf(0.875) == 2 within numerical tolerance.
   The current inherited implementation returns -1.5 and 1.5.
2. CDF/ICDF round trip. Check cdf(logits, icdf(logits, u)) == u for ranks
   inside both tails and all interior buckets.
3. Sampling tail mass. A seeded large sample should place approximately 25%
   of samples outside [-2,2] in the reproduction above.
4. Density/sampling consistency. A goodness-of-fit test should accept the
   distribution defined by forward() and reject the finite inherited-ICDF law.
5. Device/dtype coverage. Test CPU and CUDA when available, including batched
   logits and float32/float64.

Why this matters

The mismatch invalidates any use that assumes samples and density values refer
to the same distribution, including importance sampling, Monte Carlo
expectations, density estimation, generative sampling, and calibrated tail
quantiles. For importance sampling it can additionally violate the proposal
support condition because the implemented sampler has finite support while the
reported proposal density is full-support.

[oscarkey]

hi @ppguo , thanks for reporting this! I agree that this is incorrect. I'll briefly check if they'll be any unintended consequences from changing this, and then fix it.

[ppguo]

Thanks for confirming! I'd be happy to work on the fix and open a PR for this. If you haven't started on it yet, feel free to assign the issue to me. Once you've checked for any potential unintended consequences, I can take it from there.

[oscarkey]

Oh that would be great, thank you! Yes go for it. I can run the fix through our benchmarks once it's ready.

[ppguo]

Hi @oscarkey, the fix is now ready for review in #1215. It aligns cdf(), icdf(), and sample() with the existing half-normal tail density, and includes targeted CPU/CUDA regression tests plus a float32 quantile-endpoint guard. If you have time, it would be great to run it through the benchmarks you mentioned. Thanks!