MLX-Boosting

GPU-accelerated gradient boosting algorithms for Apple Silicon, built on Apple MLX.

Features

• XGBoost-style implementation with second-order gradients
• Gradient Boosted Decision Trees (GBDT) for regression and classification
• Decision Trees as standalone estimators
• Optimized for Apple Silicon (M1/M2/M3/M4) using MLX
• Numba JIT compilation for fast tree building
• scikit-learn compatible API

Installation

pip install mlx-boosting

Requirements:

• macOS with Apple Silicon (M1/M2/M3/M4)
• Python 3.10+

Quick Start

XGBoost Regressor

import mlx.core as mx
from mlx_boosting import XGBoostRegressor

# Create sample data
X = mx.random.normal((1000, 10))
y = mx.random.normal((1000,))

# Train model
model = XGBoostRegressor(
    n_estimators=100,
    max_depth=6,
    learning_rate=0.1,
)
model.fit(X, y)

# Predict
predictions = model.predict(X)

XGBoost Classifier

import mlx.core as mx
from mlx_boosting import XGBoostClassifier

# Binary classification
X = mx.random.normal((1000, 10))
y = mx.array((mx.random.uniform((1000,)) > 0.5).astype(mx.int32))

model = XGBoostClassifier(n_estimators=100, max_depth=6)
model.fit(X, y)

# Predict probabilities
probs = model.predict_proba(X)

# Predict classes
classes = model.predict(X)

Gradient Boosting

from mlx_boosting import GradientBoostingRegressor, GradientBoostingClassifier

# Regression
reg = GradientBoostingRegressor(n_estimators=100, max_depth=4)
reg.fit(X, y)

# Classification
clf = GradientBoostingClassifier(n_estimators=100, max_depth=4)
clf.fit(X, y_class)

Decision Trees

from mlx_boosting import DecisionTreeRegressor, DecisionTreeClassifier

# Standalone decision tree
tree = DecisionTreeRegressor(max_depth=6)
tree.fit(X, y)
predictions = tree.predict(X)

Parameters

XGBoostRegressor / XGBoostClassifier

GradientBoostingRegressor / GradientBoostingClassifier

Performance

MLX-Boosting is optimized for Apple Silicon and achieves excellent performance on high-volume datasets:

MLX-Boosting achieves up to 3x faster training on high-volume data compared to sklearn's GradientBoosting, running natively on Apple Silicon.

Working with NumPy

MLX-Boosting works seamlessly with NumPy arrays:

import numpy as np
import mlx.core as mx
from mlx_boosting import XGBoostRegressor

# NumPy data
X_np = np.random.randn(1000, 10).astype(np.float32)
y_np = np.random.randn(1000).astype(np.float32)

# Convert to MLX
X = mx.array(X_np)
y = mx.array(y_np)

# Train
model = XGBoostRegressor(n_estimators=100)
model.fit(X, y)

# Predictions back to NumPy
preds = np.array(model.predict(X))

API Reference

Classes

• XGBoostRegressor - XGBoost-style regression
• XGBoostClassifier - XGBoost-style classification (binary and multiclass)
• GradientBoostingRegressor - GBDT regression
• GradientBoostingClassifier - GBDT classification
• DecisionTreeRegressor - Decision tree regression
• DecisionTreeClassifier - Decision tree classification

Common Methods

• fit(X, y) - Train the model
• predict(X) - Make predictions
• predict_proba(X) - Predict probabilities (classifiers only)

License

MIT License - see LICENSE for details.

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

Acknowledgments

• Apple MLX - The foundation for GPU acceleration
• XGBoost - Inspiration for the algorithm implementation
• scikit-learn - API design patterns