Jane Street Market Prediction
XGBoost + Treelite Notebook 分析

这份 notebook 的重点不是重新发明预测模型,而是把一个 XGBoost 树模型编译成原生动态库,在 Jane Street 的逐批在线预测环境中降低推理延迟。

一、核心结论

Notebook 走的是一个清晰的 baseline 流程:读取 Feather 数据,筛选较新的有效样本,训练 XGBoost 二分类器,用 Treelite 编译模型,再通过比赛 API 流式提交动作。

130
使用的特征数量
1,571,415
最终训练样本数
约 5.2x
本次测速中的墙钟加速比
读取数据→日期与权重筛选→构造 action 标签→训练 XGBoost→Treelite 编译→流式提交

二、数据与预处理

Notebook 从 ../input/janestreet-save-as-feather/ 读取四个 Feather 文件:

文件用途
train.feather训练数据,原始形状为 (2,390,491, 138)
features.feather特征相关元数据,随后未使用
example_test.feather测试样例,随后未使用
example_sample_submission.feather提交格式样例,随后未使用
START_DATE = 85
train = train.query(f'date > {START_DATE}')
train = train.query('weight > 0').reset_index(drop=True)

日期过滤后样本从 2,390,491 变为 1,862,597;再删除零权重样本,最终保留 1,571,415 行。零权重记录不会对交易收益产生贡献,因此删除它们可以节省训练资源。

特征通过列名选择:

feats = train.columns[train.columns.str.startswith('feature')].values.tolist()

最终使用 feature_0 到 feature_129,没有把 date、weight、resp 或 ts_id 直接作为输入,避免明显的目标泄漏。

三、标签与模型

train['action'] = train['resp'] * train['weight']
train['action'] = 1 * (train['action'] > 0)

由于前面已经保证 weight > 0,这个标签实际等价于:

action = 1 if resp > 0 else 0

因此模型学习的是“未来响应变量是否为正”,输出为交易概率,而不是直接预测收益大小。

参数值含义
objectivebinary:logistic二分类概率输出
learning_rate0.08每棵树的学习步长
max_depth7单棵树最大深度
subsample0.8每棵树使用 80% 样本
colsample_bytree0.72每棵树使用 72% 特征
eval_metriclogloss训练期间观察对数损失
关键问题:树数量配置不一致。 参数字典中写了 n_estimators = 480,但实际调用 xgb.train(params, dtrain, 100, ...),第三个参数才是 boosting 轮数。因此实际只训练了 100 棵树。输出从 [0] 到 [99] 也证实了这一点。若目标是 480 棵树,应使用 num_boost_round=480。

四、Treelite 编译与推理

model = treelite.Model.from_xgboost(bst)
model.export_lib(
    toolchain='gcc',
    libpath='./mymodel.so',
    params={'parallel_comp': 32},
    verbose=True
)
predictor = treelite_runtime.Predictor('./mymodel.so', verbose=True)

treelite 负责把 XGBoost Booster 转为可编译模型,生成 mymodel.so;treelite_runtime 负责加载这个共享库并执行预测。parallel_comp=32 主要加速模型代码的编译,不等同于在线预测一定使用 32 个线程。

推理方式测试设置Wall time
普通 XGBoost10,000 行 × 130 特征75.4 ms
Treelite同一批数据14.5 ms

本次单次测试约为 75.4 / 14.5 = 5.2 倍加速。随后 notebook 用逐元素比较检查两套模型的结果完全一致:

predicted_normal == predicted_treelite

更稳健的生产检查应使用 np.allclose,并额外用包含缺失值的真实样本验证数值一致性。

五、在线提交逻辑

for test_df, pred_df in tqdm(iter_test):
    if test_df['weight'].item() > 0:
        batch = treelite_runtime.Batch.from_npy2d(test_df[feats].values)
        pred_df.action = (
            predictor.predict(batch) > TRADING_THRESHOLD
        ).astype('int')
    else:
        pred_df.action = 0
    env.predict(pred_df)

提交阶段的决策规则是:

实际处理约 15,219 个测试批次,速度约 90 batches/秒,总耗时约 2 分 49 秒。

六、方法论风险

没有验证集

只在训练集上观察 logloss,无法判断过拟合,也无法证明模型对未来日期有效。

指标与比赛目标不完全一致

训练的是方向分类,比赛关心的是加权交易收益,不是单纯准确率或 logloss。

阈值未经验证

0.5 是默认分类阈值,不一定是最大化收益的阈值。

速度测试不充分

只测了一次,未报告均值、方差、P95 延迟,也未拆分数据转换成本。

依赖版本较旧

代码依赖 Python 3.7、Treelite 0.93 和旧版 XGBoost,现代环境可能存在 API 或二进制兼容问题。

部署元数据不完整

没有单独保存特征顺序、训练参数、阈值、模型版本和环境版本。

七、推荐改进顺序

  1. 修正树数量配置,明确使用 num_boost_round=480,或删除无效的 n_estimators。
  2. 按日期做时间切分验证,不能只使用训练集指标。
  3. 在验证集上计算比赛风格的收益,例如 action × weight × resp。
  4. 扫描多个交易阈值,如 0.40 到 0.60,选择收益和稳定性更好的阈值。
  5. 按日期检查累计收益、波动率、最大回撤、交易数量和不同市场阶段的表现。
  6. 用多次重复实验 benchmark 普通 XGBoost 与 Treelite,并报告均值、中位数和 P95。
  7. 保存模型、特征列表、阈值、依赖版本和训练配置,保证线上特征顺序完全一致。
总体评价:这是一个偏工程部署优化的 Jane Street baseline。Treelite 部分设计清楚且测速结果明显,但模型研究部分仍较简化。若要提升竞赛可靠性,优先级应是“时间验证与收益评估”,其次才是继续调整模型或推理速度。