JAXBoost

JAX autodiff for XGBoost/LightGBM objectives.

Write a loss function, get gradients and Hessians automatically. No manual derivation needed.

Works with XGBoost and LightGBM.

Install

pip install jaxboost

Quick Start

XGBoost

import xgboost as xgb
import jax.numpy as jnp
from jaxboost import auto_objective, focal_loss, huber, quantile

# Prepare your data
dtrain = xgb.DMatrix(X_train, label=y_train)
params = {"max_depth": 4, "eta": 0.1}

# Built-in objectives - just use them
model = xgb.train(params, dtrain, num_boost_round=100, obj=focal_loss.xgb_objective)
model = xgb.train(params, dtrain, num_boost_round=100, obj=huber.xgb_objective)
model = xgb.train(params, dtrain, num_boost_round=100, obj=quantile(0.9).xgb_objective)

# Custom objective - write the loss, autodiff handles the rest
@auto_objective
def asymmetric_mse(y_pred, y_true, alpha=0.7):
    error = y_true - y_pred
    return jnp.where(error > 0, alpha * error**2, (1 - alpha) * error**2)

model = xgb.train(params, dtrain, num_boost_round=100, obj=asymmetric_mse.xgb_objective)

LightGBM

import lightgbm as lgb
from jaxboost import huber

train_data = lgb.Dataset(X_train, label=y_train)
params = {"max_depth": 4, "learning_rate": 0.1}

model = lgb.train(params, train_data, num_boost_round=100, fobj=huber.lgb_objective)

Available Objectives

Regression

Binary Classification

Multi-class Classification

Survival Analysis

Ordinal Regression

Multi-task Learning

Uncertainty Estimation

Ordinal Regression

XGBoost/LightGBM have no native ordinal objective. JAXBoost implements proper Cumulative Link Models:

from jaxboost import ordinal_logit, qwk_ordinal

# Wine quality: 6 ordered classes (3-8 mapped to 0-5)
ordinal = ordinal_logit(n_classes=6)
ordinal.init_thresholds_from_data(y_train)

# Works with XGBoost
model = xgb.train(params, dtrain, obj=ordinal.xgb_objective)

# Or LightGBM
model = lgb.train(params, train_data, fobj=ordinal.lgb_objective)

# Get class probabilities
probs = ordinal.predict_proba(model.predict(dtest))
classes = ordinal.predict(model.predict(dtest))

Evaluation Metrics

When using custom objectives, use matching evaluation metrics:

from jaxboost import ordinal_logit
from jaxboost.metric import qwk_metric, mae_metric

ordinal = ordinal_logit(n_classes=6)
ordinal.init_thresholds_from_data(y_train)

# Train with custom metric monitoring
model = xgb.train(
    {'disable_default_eval_metric': 1, 'max_depth': 4},  # Disable default metrics!
    dtrain,
    obj=ordinal.xgb_objective,
    custom_metric=ordinal.qwk_metric.xgb_metric,  # Built-in QWK metric
    evals=[(dtest, 'test')]
)

Available Metrics

Custom Objectives

The @auto_objective decorator turns any loss function into an XGBoost/LightGBM objective:

import xgboost as xgb
import lightgbm as lgb
import jax.numpy as jnp
from jaxboost import auto_objective

@auto_objective
def my_custom_loss(y_pred, y_true, **kwargs):
    # Write your loss here - JAX computes grad/hess automatically
    return (y_pred - y_true) ** 2

# Use with XGBoost
dtrain = xgb.DMatrix(X_train, label=y_train)
params = {"max_depth": 4, "eta": 0.1}
model = xgb.train(params, dtrain, num_boost_round=100, obj=my_custom_loss.xgb_objective)

# Use with LightGBM
train_data = lgb.Dataset(X_train, label=y_train)
params = {"max_depth": 4, "learning_rate": 0.1}
model = lgb.train(params, train_data, num_boost_round=100, fobj=my_custom_loss.lgb_objective)

# Pass parameters
model = xgb.train(
    params, dtrain, num_boost_round=100,
    obj=my_custom_loss.get_xgb_objective(alpha=0.5)
)

Multi-class Example

import xgboost as xgb
import jax
import jax.numpy as jnp
from jaxboost import multiclass_objective

@multiclass_objective(num_classes=3)
def custom_multiclass(logits, label):
    # logits: (num_classes,), label: scalar
    probs = jax.nn.softmax(logits)
    return -jnp.log(probs[label] + 1e-7)

dtrain = xgb.DMatrix(X_train, label=y_train)
model = xgb.train(
    {"num_class": 3, "max_depth": 4, "eta": 0.1},
    dtrain,
    num_boost_round=100,
    obj=custom_multiclass.xgb_objective
)

Sklearn Interface

Use custom objectives with XGBClassifier and XGBRegressor:

from xgboost import XGBClassifier
from jaxboost import focal_loss

clf = XGBClassifier(
    objective=focal_loss.sklearn_objective,
    n_estimators=100,
    max_depth=4
)
clf.fit(X_train, y_train)
y_pred = clf.predict(X_test)

Why jaxboost?

Benchmark Results

JAXBoost shines when XGBoost/LightGBM have no native solution:

Bounded Regression (Proportions in [0, 1])

Predicting proportions where standard MSE can predict outside valid range.

9.5% improvement + guaranteed valid outputs.

@auto_objective
def soft_crossentropy(y_pred, y_true):
    mu = sigmoid(y_pred)
    return -(y_true * jnp.log(mu) + (1 - y_true) * jnp.log(1 - mu))

Ordinal Regression (Wine Quality)

Predicting ordered categories (ratings 3-8) with Quadratic Weighted Kappa.

JAXBoost ordinal objectives provide proper probability distributions over classes.

When to Use JAXBoost

📊 Full benchmark details →

Requirements

• Python >= 3.10
• JAX >= 0.4.20

Documentation

Full documentation available at: https://jxucoder.github.io/jaxboost/

License

Apache 2.0