ONNX项目中TreeEnsemble算子的编码优化与功能增强

背景与现状

ONNX作为开放的神经网络交换格式，其TreeEnsemble系列算子（包括TreeEnsembleClassifier和TreeEnsembleRegressor）长期以来存在一些编码限制和效率问题。当前实现主要面临三个核心挑战：

1. 集合成员关系表达能力不足：现有算子无法直接表示集合成员关系（SET_MEMBERSHIP）这一常见操作，特别是在处理类别型变量时。上游框架如LightGBM经常产生这类操作，而当前转换器只能通过串联相等比较来模拟，导致计算图结构复杂化。

2. 编码冗余问题：现有实现包含多个冗余属性，如node_hitrates和nodes_missing_value_tracks_true等，这些属性要么未被实际使用，要么可以通过更简洁的方式表达。

3. 精度支持局限：当前算子仅支持32位浮点输出，与主流机器学习框架如XGBoost和LightGBM的双精度支持不匹配，导致数值精度差异。

技术改进方案

集合成员关系支持

新增SET_MEMBERSHIP节点类型，通过专用属性存储可能的成员集合。这种直接编码方式相比当前通过多个EQ节点串联的实现具有明显优势：

• 减少计算图复杂度
• 提高运行时效率
• 增强模型可解释性
• 确保不同框架间转换一致性

编码优化

针对冗余属性进行精简：

1. 移除未使用属性：node_hitrates和node_hitrates_as_tensor等未被实际使用的属性将被移除，简化算子定义。

2. 简化节点模式：将节点模式缩减为核心正交集（BRANCH_LEQ、BRANCH_LT、BRANCH_EQ和LEAF），保持表达能力的同时提高实现效率。

3. 优化缺失值处理：nodes_missing_value_tracks_true属性可通过分支重排实现相同语义，减少运行时分支判断开销。

双精度支持扩展

新增对64位浮点输出的支持，解决与上游框架的数值精度差异问题。这一改进将：

• 确保数值计算一致性
• 满足高精度应用场景需求
• 保持向后兼容性

架构演进建议

基于对现有算子的分析，提出更根本性的架构改进：

1. 统一算子设计：将TreeEnsembleClassifier和TreeEnsembleRegressor合并为单一TreeEnsemble算子，通过后续标准操作实现分类功能。这种设计具有以下优势：

   • 减少算子维护成本
   • 提高组合灵活性
   • 简化运行时实现

2. 标签编码外置：将classlabels_strings等属性移除，改为通过LabelEncoder等标准操作实现，增强模型模块化。

3. 多目标输出优化：支持向量化叶节点输出，避免为多目标场景复制整个树结构，提高模型紧凑性。

实现考量与性能影响

这些改进需要平衡表达力与性能：

1. 运行时优化：分支模式简化和冗余属性移除可直接提升推理速度，特别是在大规模树集成场景。

2. 内存效率：更紧凑的编码格式减少模型体积，改善加载和缓存效率。

3. 转换器兼容性：通过渐进式改进路径（如先引入新算子再弃用旧版）确保生态平稳过渡。

未来展望

本次TreeEnsemble算子的改进为ONNX在传统机器学习领域的持续优化奠定了基础。类似的设计理念可扩展至其他算子（如LinearClassifier/SVM等），推动ONNX成为更统一高效的模型交换标准。

随着多目标学习和高精度计算需求的增长，ONNX在保持性能的同时增强表达力的努力，将使其在工业部署和学术研究中发挥更大价值。