Leaves Project Evolution Plan (Draft v1.0)

> Status: Draft v1.0 · Date: 2026-06-13 · Reference projects: XGBoost, GoMLX · Current version: v0.8.0 (pure-Go GBRT prediction library)

*Note: this is a structured English summary of a long Chinese-language design document; section numbering follows the original.*

Key points

• Core question: How to evolve leaves from a pure prediction library (loading LightGBM, XGBoost, and scikit-learn models) into a full GBRT framework with training, GPU acceleration, and interpretability — using GoMLX as the compute base — without breaking ecosystem compatibility?
• Design principles: progressive enhancement (each phase ships a usable release), stable public Ensemble API with replaceable internals, GoMLX-first with native fallback, learn from XGBoost without cloning it, and test-driven evolution.

Current state (Section 1)

• leaves today loads LightGBM (text/JSON), XGBoost (binary), and scikit-learn (pickle) models; supports gbtree, dart, gblinear, random forest; dense and CSR inputs; transforms (Raw, Logistic, Softmax, Exponential, LeafIndex); NaN routing via default_left; LightGBM categorical one-hot rules; and multithreaded batch prediction via goroutines.
• Gaps vs. XGBoost (prioritized): model formats and dump (high), Tree SHAP (high), GPU inference (high), feature importance, evaluation metrics, native categorical features, multi-output trees, JSON serialization, training (long-term), survival analysis, distributed training.
• Performance: single-threaded inference is within ~20% of the C API; the bottleneck is scalar per-node tree traversal, which cannot use SIMD or GPUs.

GoMLX feasibility (Section 2)

• GoMLX uses symbolic graph building + JIT compilation (XLA to CPU/GPU/TPU) similar to JAX/TF2. Key primitives for tree inference: LessThan, Where, Gather, While, If, Const, Reshape, ReduceSum, IsNaN.
• Backends: XLA (fastest, CPU/GPU/TPU), SimpleGo (WASM/edge portable), go-coreml (Apple Metal).
• Limitations and mitigations: fixed shapes at build time (use Exec shape cache, 32 slots); dynamic shapes still in progress (one subgraph per tree); While-loop overhead (fully unroll loops for max_depth ≤ 8); no direct CSR gather (support Dense first).

Roadmap overview (Section 3)

| Phase | Version | Theme | |-------|---------|-------| | 0 | v0.9.0 | Architecture dual-track (compatibility) | | 1 | v1.0.0 | GoMLX inference engine (acceleration) | | 2 | v1.1.0 | Format & feature alignment | | 3 | v1.2.0 | Interpretability & advanced prediction | | 4 | v2.0.0 | Training capability | | 5 | v3.0.0 | High performance & distributed |

Phase 0 — Architecture dual-tracking (v0.9.0)

• Add GoMLX dependency (requires Go ≥ 1.25; upgrade Go in two steps from the current Go 1.12 baseline).
• Split the monolithic root package into model/, tree/, io/, ensemble/, transform/, mat/, internal/, util/; root keeps type aliases (e.g., type Ensemble = model.Ensemble) for zero-breakage compilation; aliases removed at v1.0.0.
• Define a pluggable tree.Engine interface with PredictDense, PredictCSR, PredictSingle, NOutputGroups, Close, plus a Factory for backend creation.
• Introduce a framework-agnostic intermediate representation (ForestIR with NumFeatures, NumOutputGroups, BaseScore, Trees), where TreeIR stores arrays of SplitFeature, SplitThreshold, DefaultLeft, LeftChild, RightChild, LeafValue, and LightGBM categorical fields (CatBoundaries, CatThresholds).

Phase 1 — GoMLX inference engine (v1.0.0)

• Approach A (recommended start): fully unroll the traversal loop for fixed-depth trees (typical XGBoost max_depth = 6, 63 internal nodes) into a static graph of Gather/LessThan/Where ops.
• Approach B: While loop for variable-depth trees (e.g., LightGBM's num_leaves-growth trees), looping while the node index is non-negative (leaves encoded as negative).
• Approach C: LightGBM categorical nodes via vectorized equality checks reduced with OR.
• Parallelize all trees into a [batch, numTrees] tensor and sum with ReduceSum.
• Missing values: insert IsNaN checks before comparisons, routing to defaultLeft.
• Transforms reuse GoMLX built-ins (Logistic, fused Softmax with numerical stability, Exp).
• Backend switching via WithBackend(...): BackendAuto (GoMLX/XLA with native fallback), BackendGoMLX, BackendNative, BackendSimpleGo (WASM).
• Performance expectations: native path remains best for small batches (1–16); GoMLX/XLA GPU wins at medium (100–1K) and large (10K+) batches; SimpleGo is the only WASM option.

Phase 2 — Format & feature alignment (v1.1.0)

• XGBoost JSON model loader (XGEnsembleFromJSON): parse gbtree_model_param, the trees[] array (split_indices, split_conditions, left_children, right_children, base_weights), and train_param (base_score, num_class), converting to ForestIR.
• Metrics package: a Metric interface (Name, Evaluate, EvaluatePerGroup, HigherIsBetter) implementing a core subset of XGBoost's 24+ metrics — RMSE, MAE, MAPE, RMSLE, LogLoss, AUC, Error, multi-class LogLoss/Error, NDCG@k, MAP.
• Native categorical features (XGBoost 1.6+): parse the ftype field; categorical nodes use partition-based membership tests rather than numeric comparisons; extend TreeIR with IsCategorical and Categories.
• Multi-output trees (XGBoost 3.x MultiOutputTree): vector leaves stored as LeafValues[numLeaves * outputDim]; GoMLX Reshape + Gather yields [batch, outputDim].

Phase 3 — Interpretability (v1.2.0)

• Tree SHAP: recursive path-weight algorithm (O(TLD²), T = trees, L = leaves, D = depth) using cover weights stored in ForestIR; CPU-only path computation in pure Go, with GoMLX used for batch post-processing of SHAP matrices. API: TreeExplainer.ShapleyValues (output includes bias column) and InteractionValues.
• Feature importance: five types (weight, gain, total_gain, cover, total_cover), computed from split-feature counts plus loss_chg/sum_hess parsed from XGBoost binaries into new TreeIR fields (SplitGain, SumHess).
• Visualization: DumpText, DumpJSON, DumpDOT (Graphviz) exports from ForestIR.
• Saabas approximate contributions: a lightweight SHAP alternative distributing prediction deltas along tree paths; 10x+ faster but assumes feature independence.
• Deliverables include consistency tests against XGBoost's Python SHAP output.

Phase 4 — Training capability (v2.0.0)

• Layered training architecture: Trainer (loop, CV, early stopping, callbacks) → GradientBooster (gbtree/dart/gblinear) → Objective + Metric → TreeBuilder (histogram building, best-split search) → GoMLX compute backend.
• Booster interface (UpdateOneIter, Predict, Save/Load) and a DMatrix data container (data, labels, weights, groups).
• First six objectives with gradient/Hessian formulas: reg:squarederror, reg:logistic, binary:logistic, reg:gamma, multi:softmax, rank:pairwise (LambdaRank).
• Histogram-based tree building (XGBoost hist method) with GoMLX-accelerated histogram aggregation and split-gain computation; hyperparameters include max_depth, max_leaves, min_child_weight, L2 regularization, and column sampling.

Phase 5 — High performance & distributed (v3.0.0)

Planned: GPU/distributed training and inference, quantized inference, ONNX interop, and WASM deployment.

Risks and milestones (Sections 10–11)

The document closes with risk/countermeasure analysis and a milestone timeline keyed to the six phase releases above (truncated in the source post), with per-phase checklists: package restructuring, engine interface, IR definitions, native engine wrapper, GoMLX engine with XLA and SimpleGo benchmarks, compatibility tests for XGBoost/LightGBM/sklearn, and backend usage documentation.