Skip to content

Add NegativeBinomial support to XGBoostLSS - #1222

Open
marrov wants to merge 3 commits into
sktime:mainfrom
marrov:enh/xgboostlss-negative-binomial
Open

marrov wants to merge 3 commits into
sktime:mainfrom
marrov:enh/xgboostlss-negative-binomial

Conversation

@marrov

@marrov marrov commented Oct 22, 2026

Copy link
Copy Markdown
Member

Reference Issues/PRs

None.

What does this implement/fix? Explain your changes.

Adds support for dist="NegativeBinomial" to the XGBoostLSS wrapper. Both skpro and xgboostlss already provide this distribution, but the wrapper currently cannot construct it or convert its predicted parameters.

Unlike the other supported distributions, xgboostlss's NegativeBinomial requires separate response functions for its parameters. This change:

  • Passes the wrapper's response_fn to response_fn_total_count and uses "sigmoid" for response_fn_probs.
  • Converts predictions to skpro's parameterization using mu = total_count * probs / (1 - probs) and alpha = total_count.
  • Clips probabilities to avoid non-finite moments at numerical saturation, including for float32 predictions.
  • Documents the distribution and its supported response functions and loss.
  • Adds a negative-binomial configuration to get_test_params for coverage by the generic estimator checks.

The fitted total_count is preserved as alpha, without imposing a dispersion floor that would change the predicted variance.

Does your contribution introduce a new dependency? If yes, which one?

No.

What should a reviewer concentrate their feedback on?

  • The parameter conversion between the PyTorch/xgboostlss and skpro negative-binomial distributions.
  • The handling of the separate response functions.
  • The probability-clipping safeguard, particularly for float32 predictions.

Did you add any tests for the change?

Yes. Two focused regression tests cover:

  • Mean and variance agreement with PyTorch's NegativeBinomial, including small total_count values, and finite means for saturated float32 probabilities.
  • Fitting on count data, routing a non-default response function to total_count, and returning a skpro NegativeBinomial with integer-valued quantiles.

The new get_test_params configuration also exercises the distribution through skpro's existing generic estimator checks, avoiding duplicate tests for the standard prediction interface.

Local validation:

  • All 8 XGBoostLSS wrapper tests passed with the VM-tag skip temporarily disabled for the test run.
  • All 269 generic estimator checks passed across 10 configurations.
  • All applicable pre-commit checks passed.

Any other comments?

Local fitting tests were run with OMP_NUM_THREADS=1 on a macOS laptop. Without this setting, negative-binomial fitting stalled on my macOS environment, including when using xgboostlss directly without skpro. The environment loads multiple OpenMP runtimes; the cause has not been confirmed.

No threading workaround is included in this PR. Feedback on whether this also occurs on the macOS CI runners would be welcome.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant