leaves 项目演进计划

状态: 草案 v1.0
日期: 2026-06-13
参考项目: XGBoost · GoMLX
当前版本: v0.8.0（纯 Go GBRT 预测库）

目录

• 〇、总纲：从「预测壳」到「可训练、可加速、可解释」的完整 GBRT 框架
• 一、现状剖析
• 二、GoMLX 技术基础评估
• 三、演进路线图概览
• 四、Phase 0：架构双轨化（兼容期）
• 五、Phase 1：GoMLX 推理引擎（加速期）
• 六、Phase 2：格式与功能对齐（补全期）
• 七、Phase 3：可解释性与高级预测（增值期）
• 八、Phase 4：训练能力（质变期）
• 九、Phase 5：高性能与分布式（成熟期）
• 十、风险与对策
• 十一、里程碑与时间线

〇、总纲

费曼式一句话

leaves 现在是个「翻译器」——读到 XGBoost/LightGBM 训练好的模型文件，逐棵树走一遍，吐出预测结果。它的骨架是对的（统一接口 + 多后端适配），缺的是肌肉（GPU 加速、SHAP 解释）和大脑（训练能力）。GoMLX 就是那块能长肌肉和大脑的培养皿——它的符号图 + XLA JIT 编译，恰好能把树遍历变成可并行、可微分的图操作。

核心命题

如何在不丢弃现有生态兼容性（LightGBM、XGBoost、scikit-learn 模型加载）的前提下，用 GoMLX 作为计算底座，分阶段将 leaves 从一个纯预测库演进为一个支持训练、GPU 加速、可解释的完整 GBRT 框架？

设计原则

1. 渐进增强，非推倒重来：每个 Phase 产出可用的中间版本
2. 接口稳定性：Ensemble 公开 API 保持兼容，内部实现可替换
3. GoMLX 优先，本地回退：GPU/XLA 可用则用 GoMLX，否则回退当前纯 Go 路径
4. 参考 XGBoost，不做克隆：学习 XGBoost 的架构思想，但用 Go 的惯用方式表达
5. 测试驱动演进：每个 Phase 有对应基准测试和兼容性测试

一、现状剖析

1.1 当前 leaves 架构

┌──────────────────────────────────────────────────────┐
│                    Ensemble (公开 API)                │
│  PredictSingle / Predict / PredictDense / PredictCSR │
├──────────────────────────────────────────────────────┤
│               ensembleBaseInterface (内部接口)         │
├─────────────────┬────────────────┬───────────────────┤
│   lgEnsemble    │   xgEnsemble   │     xgLinear      │
│  (LightGBM树)   │  (XGBoost树)   │   (XGBoost线性)    │
├─────────────────┴────────────────┴───────────────────┤
│                    lgTree (树结构体)                    │
│    逐节点比较 → 走左/右分支 → 到达叶子 → 取值           │
└──────────────────────────────────────────────────────┘
        │ IO 层
        ▼
┌───────────────┐  ┌──────────────────┐  ┌──────────────┐
│ lgEnsemble_io │  │  xgEnsemble_io   │  │ skEnsemble_io│
│ (文本/JSON)   │  │  (二进制 XGBoost) │  │ (pickle)     │
└───────────────┘  └──────────────────┘  └──────────────┘

1.2 leaves 已具备的能力

1.3 与 XGBoost 的差距（仅列可补足的）

1.4 当前性能基准

单线程下，leaves 与 C API 性能接近（差距 < 20%）。多线程下持平。瓶颈在树遍历的标量循环——每棵树每条样本走一遍分支判断，纯标量操作，无法利用 SIMD 或 GPU。

二、GoMLX 技术基础评估

2.1 GoMLX 核心模型

GoMLX 采用符号计算图 + JIT 编译范式（类似 JAX/TF2）：

阶段1: 图构建 (Graph Building)
  → 用 Go 函数描述符号操作，不执行实际计算
  
阶段2: 编译 (Compile)  
  → JIT 编译为设备可执行代码（XLA HLO → 机器码/PTX）
  
阶段3: 执行 (Execute)
  → 输入具体 Tensor 值，获取输出（自动缓存编译结果）

2.2 与树推理相关的关键原语

2.3 后端选择矩阵

2.4 GoMLX 的限制

三、演进路线图概览

Phase 0          Phase 1           Phase 2            Phase 3           Phase 4           Phase 5
架构双轨化        GoMLX推理引擎      格式与功能对齐      可解释性          训练能力           高性能分布式
(兼容期)          (加速期)           (补全期)            (增值期)          (质变期)           (成熟期)
                                                                                          
v0.8.0 ──────► v0.9.0 ──────► v1.0.0 ──────► v1.1.0 ──────► v2.0.0 ──────► v3.0.0
   │                │                │                 │                │                │
   │ 拆分monolith   │ 图编译树推理    │ JSON模型         │ Tree SHAP      │ 基础训练循环    │ GPU/分布式
   │ 接口抽象       │ XLA/GPU加速    │ 评估指标         │ 特征重要性     │ GBDT训练器     │ 量化推理
   │ 引入gomlx dep  │ SimpleGo后端   │ categorical feat │ 近似贡献值     │ objective/metric│ ONNX互操作
   │                │                │ 多输出树          │ 模型可视化     │ CV/early stop  │ WASM部署

四、Phase 0：架构双轨化（兼容期） 目标版本 v0.9.0

目标

在不改变任何公开 API 的前提下，重构内部架构，为引入 GoMLX 做准备。

4.1 引入 GoMLX 依赖

go get github.com/gomlx/gomlx@latest

注意：GoMLX 要求 Go ≥ 1.25.0，当前 leaves 使用 Go 1.12。必须分两步：

1. 先升级 Go 版本（leaves 代码本身兼容 Go 1.18+ 即可）
2. 再添加 GoMLX 依赖

4.2 拆分 monolithic 根包

当前 leaves 根包混合了所有逻辑。改为按职责拆包：

leaves/                    # 根包：仅公开 API（向后兼容的类型别名）
├── model/
│   ├── ensemble.go        # Ensemble 公开类型（当前 leaves.go 核心）
│   ├── interface.go       # ensembleBaseInterface（保持不变）
│   └── options.go         # 预测选项（nEstimators, nThreads 等）
├── tree/
│   ├── lgtree.go          # 树数据结构（当前 lgtree.go）
│   ├── native.go          # 原生纯 Go 树推理引擎（当前逻辑）
│   └── gomlx.go           # GoMLX 图推理引擎（Phase 1 实现）
├── io/
│   ├── lgb.go             # 当前 lgensemble_io.go
│   ├── xgb.go             # 当前 xgensemble_io.go + xgblinear_io.go
│   ├── skl.go             # 当前 skensemble_io.go
│   └── format.go          # 格式检测、版本协商
├── ensemble/
│   ├── lgb.go             # 当前 lgensemble.go
│   ├── xgb.go             # 当前 xgensemble.go
│   └── xgblinear.go       # 当前 xgblinear.go
├── transform/
│   ├── transform.go       # 当前 transformation/ 包内容
│   └── metrics.go         # Phase 2 评估指标
├── mat/
│   ├── dense.go           # 当前 mat/dense.go
│   ├── csr.go             # 当前 mat/csr.go
│   └── util.go            # 当前 mat/util.go
├── internal/              # 内部工具（当前 internal/ 内容，保持不变）
└── util/                  # 通用工具（当前 util/ 内容）

迁移策略：根包使用类型别名（type Ensemble = model.Ensemble），确保旧代码零改动编译通过。在 v1.0.0 时移除别名。

4.3 定义推理引擎接口

// tree/engine.go
package tree

// Engine 是树推理的可插拔后端
// 两种实现：NativeEngine（当前纯 Go）、GoMLXEngine（Phase 1）
type Engine interface {
    // PredictDense 在稠密矩阵上执行预测
    // 输入: vals[row*cols + col] 格式
    // 输出: predictions[row*outputGroups + group]
    PredictDense(
        vals []float64, nrows int, ncols int,
        predictions []float64,
        nEstimators int,
    ) error

    // PredictCSR 在 CSR 稀疏矩阵上执行预测
    PredictCSR(
        indptr []int, cols []int, vals []float64,
        predictions []float64,
        nEstimators int,
    ) error

    // PredictSingle 单样本预测
    PredictSingle(fvals []float64, nEstimators int) float64

    // NOutputGroups 输出维度
    NOutputGroups() int

    // Close 释放资源（如 GoMLX 的 Exec 缓存）
    Close() error
}

// Factory 根据配置创建引擎
type Factory func(treeData *TreeData, transform transform.Transform) (Engine, error)

4.4 提取树数据为中间表示 (IR)

当前 lgTree 是 LightGBM-specific 的结构。需要定义统一的树中间表示，让 GoMLX 能一次加载，多后端执行：

// tree/ir.go
package tree

// ForestIR 是所有树模型的通用中间表示
// 与框架来源无关（LightGBM/XGBoost/scikit-learn 加载后统一转为此格式）
type ForestIR struct {
    NumFeatures     int
    NumOutputGroups int  // 1 for binary/regression, N for multiclass
    Trees           []TreeIR
    BaseScore       float64
}

// TreeIR 单棵决策树的中间表示
type TreeIR struct {
    NumLeaves     int
    MaxDepth      int
    // 以完整二叉树方式存储（非叶节点有 left/right child）
    SplitFeature  []int32   // [numNodes] 分裂特征索引，-1 for leaf
    SplitThreshold []float64 // [numNodes] 分裂阈值
    DefaultLeft   []bool    // [numNodes] 缺失值默认方向
    LeftChild     []int32   // [numNodes] 左子节点索引
    RightChild    []int32   // [numNodes] 右子节点索引
    LeafValue     []float64 // [numLeaves] 叶子值
    // 分类特征（LightGBM 特有）
    CatBoundaries []uint32  // 分类特征边界
    CatThresholds []uint32  // 分类特征阈值集合
}

4.5 Phase 0 产出

• go.mod 更新到 Go 1.21+（兼容 GoMLX 要求的 Go 1.25 之前，先用最小可行版本）
• 包结构重组，根包保留类型别名
• tree.Engine 接口定义
• tree.ForestIR / tree.TreeIR 通用 IR 定义
• NativeEngine 实现（封装当前预测逻辑）
• 所有现有测试通过
• 性能回归测试基准（benchmark 对比）

五、Phase 1：GoMLX 推理引擎（加速期） 目标版本 v1.0.0

目标

用 GoMLX 符号图实现树推理，获得 GPU 加速和 XLA 编译优化。作为可选后端，用户可通过选项切换。

5.1 树遍历的图实现

核心挑战：如何把「逐节点 if/else 判断」表示成可并行计算的图？

方案 A：完全循环展开（推荐起步方案）

XGBoost 树的 max_depth 通常为 6（63 个非叶节点）。对于固定深度的树，可以将循环完全展开为静态图：

// tree/gomlx_engine.go
// 为 max_depth=6 的树构建展开后的推理图
func buildTreeGraph(g *graph.Graph, features *graph.Node, tree *TreeIR) *graph.Node {
    // 把树编码为常量
    splitFeat := graph.Const(g, tree.SplitFeature)     // [numNodes]int32
    thresholds := graph.Const(g, tree.SplitThreshold)   // [numNodes]float64
    leftChild  := graph.Const(g, tree.LeftChild)        // [numNodes]int32
    rightChild := graph.Const(g, tree.RightChild)       // [numNodes]int32
    leafValues := graph.Const(g, tree.LeafValue)        // [numLeaves]float64
    
    // 从根节点(索引0)开始
    currentNode := graph.Scalar(g, dtypes.Int32, 0)
    
    // 展开 max_depth 层
    for depth := 0; depth < tree.MaxDepth; depth++ {
        // 取当前节点的分裂特征和阈值
        featIdx := graph.Gather(splitFeat, currentNode)
        featVal := graph.Gather(features, featIdx)     // 从输入特征取值
        thr     := graph.Gather(thresholds, currentNode)
        
        // 比较: feature[splitIdx] < threshold ?
        goLeft := graph.LessThan(featVal, thr)
        
        // 路由: goLeft ? leftChild[current] : rightChild[current]
        leftIdx  := graph.Gather(leftChild, currentNode)
        rightIdx := graph.Gather(rightChild, currentNode)
        currentNode = graph.Where(goLeft, leftIdx, rightIdx)
    }
    
    // currentNode 现在是叶子索引（负值），取反得到正索引
    leafIdx := graph.Neg(currentNode)
    return graph.Gather(leafValues, leafIdx)
}

方案 B：While 循环（通用方案）

对于深度不确定的树（如 LightGBM 的 num_leaves 优先策略），使用 While 循环：

func buildTreeGraphWhile(g *graph.Graph, features *graph.Node, tree *TreeIR) *graph.Node {
    // ...常量设置同上...
    
    initState := []*graph.Node{
        graph.Scalar(g, dtypes.Int32, 0),  // current node index
    }
    
    // 条件：currentNode >= 0（叶子节点编码为负）
    condFn := func(g *graph.Graph, state []*graph.Node) *graph.Node {
        return graph.GreaterOrEqual(state[0], graph.Scalar(g, dtypes.Int32, 0))
    }
    
    // 主体：走一步
    bodyFn := func(g *graph.Graph, state []*graph.Node) []*graph.Node {
        currentNode := state[0]
        featIdx := graph.Gather(splitFeat, currentNode)
        featVal := graph.Gather(features, featIdx)
        thr := graph.Gather(thresholds, currentNode)
        goLeft := graph.LessThan(featVal, thr)
        next := graph.Where(goLeft,
            graph.Gather(leftChild, currentNode),
            graph.Gather(rightChild, currentNode))
        return []*graph.Node{next}
    }
    
    result := graph.While(condFn, bodyFn, initState...)
    leafIdx := graph.Neg(result[0])
    return graph.Gather(leafValues, leafIdx)
}

方案 C：分类特征（LightGBM one-hot 决策）

LightGBM 的分类特征不需要数值比较，而是用 Contains 检查：

// 对 LightGBM 分类节点的判断
// 等价于: catThresholds[boundaries[idx]:boundaries[idx+1]] 中是否包含 featVal
func buildCatDecision(g *graph.Graph, featVal *graph.Node, 
    catBoundaries, catThresholds *graph.Node) *graph.Node {
    // 对每个可能的分类值做 Equal 比较，然后 ReduceOr
    // 这比循环高效，因为可以向量化
}

5.2 多棵树并行

将所有树同时展开，用 [batch, numTrees] 张量并行执行：

func buildForestGraph(g *graph.Graph, features *graph.Node, forest *ForestIR) *graph.Node {
    // features 形状: [batchSize, numFeatures]
    
    treeOutputs := make([]*graph.Node, len(forest.Trees))
    for i, tree := range forest.Trees {
        // 每棵树的图独立构建
        treeOutputs[i] = buildTreeGraph(g, features, &tree)
    }
    
    // 堆叠所有树输出: [batchSize, numTrees]
    stacked := graph.Stack(treeOutputs, 1) // axis=1
    
    // 求和得到最终预测: [batchSize, 1]
    return graph.ReduceSum(stacked, 1)
}

5.3 缺失值处理

嵌入 IsNaN 检测：

// 在比较之前插入缺失值检查
func buildNodeDecision(g *graph.Graph, featVal, thr *graph.Node, 
    defaultLeft *graph.Node) *graph.Node {
    isMissing := graph.IsNaN(featVal)
    goLeft := graph.LogicalOr(
        isMissing,
        graph.LessThan(featVal, thr),
    )
    // 如果 defaultLeft 为 false，取反
    if !tree.DefaultLeft[currentNode] {
        goLeft = graph.LogicalNot(goLeft)
    }
    return goLeft
}

5.4 变换函数集成

GoMLX 已内置 Softmax 融合操作、Logistic(Sigmoid)、Exp 等，直接使用，无需自己实现：

func applyTransform(g *graph.Graph, raw *graph.Node, t transform.Type) *graph.Node {
    switch t {
    case transform.Raw:
        return raw
    case transform.Logistic:
        return graph.Logistic(raw)
    case transform.Softmax:
        // GoMLX 内置融合 Softmax，包含数值稳定处理
        return graph.Softmax(raw)
    case transform.Exponential:
        return graph.Exp(raw)
    case transform.LeafIndex:
        // 在树遍历时直接返回叶子索引，不用额外处理
        return raw
    }
    return raw
}

5.5 引擎切换 API

// 新选项
type PredictOption func(*predictOptions)

type predictOptions struct {
    backend tree.Backend  // tree.NativeBackend | tree.GoMLXBackend
}

// BackendAuto: 自动选择 GoMLX(XLA)，失败时回退 Native
// BackendGoMLX: 强制 GoMLX，无 GPU 则报错
// BackendNative: 强制当前纯 Go 实现
// BackendSimpleGo: GoMLX SimpleGo 后端（WASM 兼容）
func WithBackend(b tree.Backend) PredictOption { ... }

// 用法
model.PredictDense(vals, rows, cols, preds, nEstimators, nThreads,
    WithBackend(tree.BackendGoMLX),
)

5.6 Phase 1 产出

• tree/gomlx_engine.go — 基于 GoMLX 图计算的推理引擎
• go.mod 包含 github.com/gomlx/gomlx 依赖
• tree/ir.go — ForestIR/TreeIR 通用 IR 定义
• IO 层输出统一转为 ForestIR
• WithBackend() 选项
• XLA 后端 benchmark：GPU 推理性能对比
• SimpleGo 后端 benchmark：纯 CPU 性能对比
• 所有现有兼容性测试通过（XGBoost / LightGBM / sklearn）
• 文档：GoMLX 后端使用指南

5.7 性能预期

六、Phase 2：格式与功能对齐（补全期） 目标版本 v1.1.0

目标

对标 XGBoost 的主要模型格式和评估功能，消除关键功能缺口。

6.1 XGBoost JSON 模型格式

参考：C:\GitHub\xgboost\src\c_api\c_api_json.cc

XGBoost 自 1.0 起支持 JSON 格式序列化。JSON 格式比二进制格式有以下优势：

• 人类可读，便于调试
• 跨版本兼容性更好
• 包含更丰富的元数据

实现要点：

• 解析 gbtree_model_param（树数、特征数等）
• 解析 trees[] 数组（每棵树的 JSON 表示，含 split_indices、split_conditions、left_children、right_children、base_weights 等）
• 解析 train_param 获取 base_score、num_class 等
• 输出统一转 ForestIR。

// io/xgb_json.go
func XGEnsembleFromJSON(reader io.Reader, loadTransformation bool) (*Ensemble, error)

6.2 评估指标 (Metrics)

参考 XGBoost 的 24 个内置指标，实现最常用的核心子集：

// metrics/metrics.go
package metrics

type Metric interface {
    Name() string
    // Evaluate 计算指标值
    Evaluate(yTrue, yPred []float64) (float64, error)
    // EvaluatePerGroup 按 group 评估（排序任务）
    EvaluatePerGroup(yTrue, yPred []float64, groups []int) (float64, error)
    // HigherIsBetter 是否为越高越好的指标
    HigherIsBetter() bool
}

6.3 XGBoost 原生分类特征支持

参考：C:\GitHub\xgboost\src\encoder\categorical.h

XGBoost 1.6+ 引入了原生分类特征。与 LightGBM one-hot 方式不同，XGBoost 的分类节点使用分区比较：

// tree/ir.go 扩展
type TreeIR struct {
    // ...现有字段...
    
    // 分类特征节点（XGBoost 原生格式）
    // CategoryNode: 检查特征值是否在 categories 集合中
    IsCategorical []bool     // [numNodes] 该节点是否为分类分裂
    Categories    [][]uint32 // [numNodes] 分类值集合（仅 IsCategorical=true 的节点）
}

在二进制模型解析中读取 ftype（feature type）字段，ftype == 1 表示分类特征。

6.4 多输出树支持

参考：C:\GitHub\xgboost\src\tree\tree_model.h -> MultiTargetTree

XGBoost 的 multi:softmax 模式使用多棵独立树（每个类一棵），leaves 已支持。但 XGBoost 3.x 新增了 MultiOutputTree 策略——一棵树输出多个值（向量叶）：

// tree/ir.go
type TreeIR struct {
    // ...现有字段...
    
    // 多输出叶：LeafValues[numLeaves * outputDim]
    // 按 [leaf_i*dim : (leaf_i+1)*dim] 排列
    OutputDim  int
    LeafValues []float64 // 长度 = numLeaves * OutputDim
}

GoMLX 侧用 Reshape 将 [numLeaves * outputDim] 转为 [numLeaves, outputDim]，Gather 后得到 [batch, outputDim]。

6.5 Phase 2 产出

• JSON 格式模型加载器 io/xgb_json.go
• metrics/ 包：12+ 核心评估指标
• XGBoost 原生分类特征解析与推理
• 多输出树支持
• 兼容性测试覆盖上述新功能
• 评估指标 benchmark

七、Phase 3：可解释性与高级预测（增值期） 目标版本 v1.2.0

目标

实现 XGBoost 的核心可解释性功能：Tree SHAP、特征重要性，填补当前最大的功能缺口。

7.1 Tree SHAP 贡献值

参考：C:\GitHub\xgboost\src\tree\tree_model.cc (CalcContrib)
论文：Lundberg & Lee, 2017. "A Unified Approach to Interpreting Model Predictions"

Tree SHAP 是 XGBoost 最强大的可解释性功能。算法本质是在树上计算每个特征的 Shapley 值：

核心公式（树 SHAP 的递归算法）：

对于节点 i，定义两条路径的权重：

• w_left = cover[left] / cover[i]（训练样本落入左子树的比例）
• w_right = cover[right] / cover[i]

则从根到叶，沿路径累积每个特征贡献。算法复杂度 O(TLD²)，其中 T=树数，L=叶数，D=深度。对于典型 XGBoost 模型（depth=6，leaves≤64），这是可接受的。

GoMLX 实现方案

不能直接做符号图（因为 SHAP 的路径依赖逻辑太复杂），采用两阶段：

阶段 1：CPU 端做 SHAP 路径计算（纯 Go），利用 ForestIR 中已存储的 cover（覆盖率）数据。

阶段 2：如果输入是大 batch，将已算出的 SHAP 矩阵用 GoMLX 做批量后处理。

// shap/shap.go
package shap

// TreeExplainer 解释器
type TreeExplainer struct {
    forest *tree.ForestIR
}

// ShapleyValues 计算 SHAP 值
// 输入: featureValues [nSamples][nFeatures]float64
// 输出: contributions [nSamples][nFeatures + 1]float64
//       最后一列为 bias（base_score）
func (e *TreeExplainer) ShapleyValues(featureValues [][]float64) ([][]float64, error)

// InteractionValues 计算 SHAP 交互值
// 输出: [nSamples][nFeatures][nFeatures]float64
func (e *TreeExplainer) InteractionValues(featureValues [][]float64) ([][][]float64, error)

7.2 特征重要性

// importance/importance.go
package importance

type ImportanceType int
const (
    ImportanceWeight     ImportanceType = iota // 特征被使用的次数
    ImportanceGain                             // 平均增益
    ImportanceTotalGain                        // 总增益
    ImportanceCover                            // 平均覆盖
    ImportanceTotalCover                       // 总覆盖
)

// FeatureImportance 特征重要性结果
type FeatureImportance struct {
    Type   ImportanceType
    Scores []float64 // [numFeatures]
    Names  []string  // 特征名（如果有）
}

// Compute 从 ForestIR 计算特征重要性
func Compute(forest *tree.ForestIR, impType ImportanceType, featureNames []string) *FeatureImportance

计算来源：

• Weight: 遍历所有树的 SplitFeature 数组，统计每个特征被用作分裂的次数
• Gain/TotalGain: 读取 XGBoost 二进制中的 loss_chg（损失变化）/ split_gain
• Cover: 读取 sum_hess（Hessian 和 = 覆盖率）

这需要在 IO 层将 loss_chg 和 sum_hess 解析进 TreeIR：

// tree/ir.go 扩展
type TreeIR struct {
    // ...
    SplitGain []float64 // [numNodes] 分裂增益
    SumHess   []float64 // [numNodes] Hessian 和
}

7.3 模型可视化

// viz/viz.go

// DumpText 输出类似 XGBoost dump 的文本树结构
func DumpText(forest *tree.ForestIR, featureNames []string) string

// DumpJSON 输出可被前端渲染的 JSON 树结构
func DumpJSON(forest *tree.ForestIR, featureNames []string) ([]byte, error)

// DumpDOT 输出 Graphviz DOT 格式
func DumpDOT(forest *tree.ForestIR, featureNames []string) string

7.4 近似贡献值（Saabas 方法）

作为 SHAP 的轻量替代方案：

// shap/approximate.go

// ApproximateContributions 使用 Saabas 方法近似特征贡献
// 原理：沿树路径，每个节点将分裂前后的预测差分配给对应特征
// 比 Tree SHAP 快 10x+，但仅适用于独立特征假设
func (e *TreeExplainer) ApproximateContributions(featureValues [][]float64) ([][]float64, error)

7.5 Phase 3 产出

• shap/ 包：Tree SHAP 贡献值
• shap/ 包：交互值
• importance/ 包：5 种特征重要性
• viz/ 包：模型文本/JSON/DOT 输出
• 与 XGBoost Python SHAP 结果的一致性测试
• 文档：SHAP 使用示例

八、Phase 4：训练能力（质变期） 目标版本 v2.0.0

目标

实现 GBDT 的基本训练能力，使 leaves 从纯预测库变为可训练框架。

8.1 训练架构设计

参考 XGBoost 的分层架构：

┌─────────────────────────────────────────────────┐
│                   Trainer                       │
│  (训练循环、CV、early stopping、回调)             │
├─────────────────────────────────────────────────┤
│            GradientBooster (GBM)                 │
│  gbtree / dart / gblinear                        │
├──────────────────┬──────────────────────────────┤
│   Objective      │       Metric                 │
│  (损失+梯度/hess) │  (评估指标)                  │
├──────────────────┴──────────────────────────────┤
│              TreeBuilder                         │
│  (直方图构建、最佳分裂搜索)                        │
├─────────────────────────────────────────────────┤
│              GoMLX 计算后端                       │
│  (梯度计算、直方图聚合、张量加速)                   │
└─────────────────────────────────────────────────┘

// train/trainer.go
package train

type Trainer struct {
    booster  Booster
    objective objective.Objective
    metrics  []metrics.Metric
    callbacks []Callback
}

// Booster 接口
type Booster interface {
    // UpdateOneIter 一轮迭代
    UpdateOneIter(data *DMatrix, gradients, hessians []float64) error
    
    // Predict 当前模型预测
    Predict(data *DMatrix) ([]float64, error)
    
    // Save/Load 模型序列化
    Save(writer io.Writer) error
    Load(reader io.Reader) error
}

// DMatrix 训练数据容器
type DMatrix struct {
    Data    *mat.DenseMat
    Labels  []float64
    Weights []float64
    Groups  []int // 排序任务
}

8.2 Objective 目标函数

// objective/objective.go
package objective

type Objective interface {
    // GetGradient 计算梯度和 Hessian
    // pred: 当前预测值 [nSamples]
    // label: 真实标签 [nSamples]
    // 返回: grad [nSamples], hess [nSamples]
    GetGradient(pred, label []float64) (grad, hess []float64, err error)
    
    // PredTransform 预测值到最终输出
    PredTransform(rawPred []float64) []float64
}

首批实现（共 6 个最常用目标函数）：

8.3 TreeBuilder — 直方图方法

参考 XGBoost 的 hist 树构建方法（最常见、最高效的 CPU 方法）：

// train/treebuilder.go
package train

type TreeBuilder struct {
    maxDepth     int
    maxLeaves    int
    minChildWeight float64
    regLambda    float64 // L2 正则
    colsampleBytree float64
    // GoMLX 加速：用 GoMLX 做直方图聚合 + 分裂增益计算
    useGoMLX     bool
}

// BuildTree 构建一棵回归树
// gradients, hessians: 样本梯度/二阶梯度 [nSamples]
// 返回: TreeIR
func (tb *TreeBuilder) BuildTree(
    data *mat.DenseMat,
    gradients, hessians []float64,
) (*tree.TreeIR, error)

GoMLX 加速点：

• 直方图构建：ReduceSum 按特征值分箱聚合梯度
• 分裂增益计算：向量化 gain = G²/(H+λ) 评估所有候选分裂点
• 最佳分裂搜索：ReduceMax 找最大增益

8.4 DART Booster

DART (Dropouts meet Multiple Additive Regression Trees) 的推理已在 leaves 中支持（从 XGBoost 训练好的模型加载）。现在补充训练逻辑：

// train/dart.go

type DartBooster struct {
    trees    []*tree.TreeIR
    weights  []float64  // 每棵树的权重
    
    // DART 特有
    dropRate     float64  // 每次迭代丢弃比例
    skipDrop     float64  // 跳过丢弃概率
    normalizeType string   // tree / forest
}

8.5 训练 API

// 基础用法
trainer := train.NewTrainer(train.Config{
    Objective:  objective.RegSquareError,
    Booster:    train.BoosterGBTree,
    NumRound:   100,
    MaxDepth:   6,
    LearningRate: 0.3,
    Metrics:    []metrics.Metric{metrics.RMSE, metrics.MAE},
})

// 训练
err := trainer.Fit(dtrain)
if err != nil { ... }

// 保存模型
err = trainer.Save("output.model")

// 加载预测（使用现有 Ensemble API）
model, err := leaves.XGEnsembleFromFile("output.model", true)
pred := model.PredictSingle(fvals, 0)

8.6 交叉验证与早停

// train/cv.go

type CVResult struct {
    TestMetricMean []float64  // 每轮验证集均值
    TestMetricStd  []float64  // 每轮验证集标准差
    BestIteration  int
    BestScore      float64
}

// CrossValidate K 折交叉验证
func CrossValidate(config train.Config, data *DMatrix, folds int) (*CVResult, error)

// EarlyStopping 回调
type EarlyStopping struct {
    MetricName   string
    StoppingRounds int
    Maximize     bool
}

8.7 Phase 4 产出

• train/ 包：训练循环、Trainer、Booster 接口
• objective/ 包：6 个核心目标函数
• train/treebuilder.go：基于直方图的树构建器
• GoMLX 加速的直方图聚合和分裂搜索
• 与 XGBoost 训练结果的等价性测试（相同数据、相同参数）
• train/dart.go：DART 训练器
• train/cv.go：交叉验证 + 早停
• 文档：训练使用指南

九、Phase 5：高性能与分布式（成熟期） 目标版本 v3.0.0

目标

实现 GPU 原生推理、量化推理、分布式部署，使 leaves 成为生产级可部署的 GBRT 框架。

9.1 量化推理

将树模型量化为 Int8/Int4 精度，减少模型体积和推理内存：

// quantize/quantize.go
package quantize

// QuantizedForest 量化森林模型
type QuantizedForest struct {
    // 阈值和叶子值量化为 int8
    Thresholds []int8
    LeafValues []float64  // 保留 float 精度（叶子值影响大）
    Scales     []float64  // 量化的 scale/zero_point
}

// 利用 GoMLX 的 QuantizedDense 融合操作模式

9.2 ONNX 互操作

利用 onnx-gomlx 将 XGBoost ONNX 模型转为 GoMLX 图：

XGBoost 训练 → ONNX-ML 导出 → onnx-gomlx 解析 → GoMLX 图执行

这为 leaves 提供了一条获取 GPU 推理的捷径——只要 XGBoost 模型能导出 ONNX，就能通过 GoMLX 的 XLA 后端获得 GPU 加速。

// io/onnx.go

// XGEnsembleFromONNX 从 ONNX 文件加载 XGBoost 模型
func XGEnsembleFromONNX(path string) (*Ensemble, error)

9.3 GPU 直方图训练

利用 GoMLX XLA GPU 后端，将 Phase 4 的直方图构建和分裂搜索放到 GPU 上：

// train/gpu_builder.go

type GPUTreeBuilder struct {
    // 使用 GoMLX XLA 后端
    // 参考 XGBoost device_hist.cu 的 GPU 直方图算法
    // 核心：并行归约构建直方图 → 并行扫描找到最佳分裂
}

关键优化：

• GPU 上多特征直方图并行构建
• 使用 ReduceScatter 做分布式直方图归约
• 利用 XLA 的 kernel fusion 减少内存带宽

9.4 WASM 部署

利用 GoMLX 的 SimpleGo 后端（纯 Go，可编译到 WASM）：

// tree/wasm_engine.go

// Go 编译: GOOS=js GOARCH=wasm go build
// 结果可在浏览器中运行 leaves 推理
func NewWASMEngine(forest *ForestIR) (Engine, error) {
    return NewGoMLXEngine(forest, gomlx.WithBackend(gomlx.SimpleGoBackend))
}

9.5 分布式推理

基于 GoMLX 的 SPMD / AutoSharding：

• 模型分片：不同 GPU 加载不同的树子集
• 数据分片：不同节点处理不同 batch
• 结果集会：AllReduce 汇总

// distributed/distributed.go

type DistributedPredictor struct {
    // 使用 gomlx/pkg/core/distributed
    // 自动将 ForestIR 分片到多个设备
    // 利用 GoMLX 的 SPMD 模型统一管理数据并行和模型并行
}

9.6 Phase 5 产出

• quantize/ 包：Int8 量化推理
• io/onnx.go：ONNX 模型加载
• train/gpu_builder.go：GPU 直方图训练
• WASM 部署指南和 Dockerfile
• distributed/ 包：分布式推理
• 性能报告：GPU vs CPU、量化前后对比
• 部署文档完整

十、风险与对策

十一、里程碑与时间线

附录 A：与 XGBoost 的功能对应表

附录 B：GoMLX 集成关键代码路径

leaves API 调用
    │
    ▼
Ensemble.PredictDense(...)
    │
    ├── [backend=Native] ──► tree.NativeEngine.predictInner()
    │                              │
    │                              └── 逐树逐样本循环（当前逻辑）
    │
    └── [backend=GoMLX] ──► tree.GoMLXEngine.predictDense()
                                   │
                                   ├── 1. 构造 ForestIR → GoMLX 常量
                                   ├── 2. 构建计算图 (buildForestGraph)
                                   │      ├── For each tree: buildTreeGraph
                                   │      │     ├── 编码树结构为 Const
                                   │      │     ├── 展开遍历循环 (max_depth 层)
                                   │      │     └── Gather 叶子值
                                   │      └── Stack + ReduceSum
                                   ├── 3. graph.Exec.Compile() (首次/缓存命中)
                                   └── 4. exec.Call(params) → Tensor → []float64