A Minimal Tech Report for the Index-Advisor Training Framework

1. Problem Setting

本框架面向数据库索引推荐任务：输入为一个 workload、数据库 schema 以及可选的多模态侧信息，输出为一组能够降低查询代价的索引建议。核心挑战在于学习 workload 结构、谓词模式、连接图、列共现关系与索引收益之间的稳定映射。

因此，整个训练系统围绕两个目标构建：

  1. 降低对真实表名、列名和库级分布的记忆依赖。
  2. 让模型先学会表示，再学会监督模仿，最后学会按真实执行收益优化策略。

围绕这一目标，框架采用五阶段流水线：训练数据扰动、对比学习、PJ(projector)训练、SFT、RL。

2. Pipeline Overview

完整流程如下：

raw SQL/schema/workload
    -> label generation
    -> perturbation and template augmentation
    -> column/SQL embedding construction
    -> contrastive representation learning
    -> PJ(projector) training
    -> SFT supervised alignment
    -> RL policy optimization
    -> evaluation on cost reduction and format validity

其中，前两阶段主要解决泛化与表示问题，中间两阶段解决“把结构信号注入 LLM”的问题，最后一阶段直接对数据库执行收益做优化。

3. Stage I: Training Data Perturbation

3.1 Motivation

若训练 prompt 直接暴露真实表名和列名，模型很容易走记忆捷径。对于索引推荐，这会导致模型在 seen schema 上表现较好，但一旦切换到新数据库或混合 workload，其泛化能力迅速下降。

3.2 Method

本框架在 index_advisor/perturb.py 与 scripts/gen/perturb_workloads.py 中实现了 schema/workload 扰动。核心做法是：

  • 将表名替换为抽象标识，如 T1, T2, ...
  • 将列名替换为抽象标识，如 C1, C2, ...
  • 将 alias 一并扰动，避免 SQL 文本残留原始命名线索
  • 支持 workload 级全局 view，使同一 workload 内多条 SQL 共享一致映射
  • 保留从 surface name 到 canonical schema 的可逆映射，确保标签与输入严格对齐

3.3 Effect

这一阶段的作用不是制造噪声，而是将“名称不变性”强加给模型。经过扰动后，训练信号被迫更多依赖：

  • 谓词列与连接列的位置
  • 查询模板与 workload 结构
  • 列之间的组合关系
  • 索引对代价下降的真实影响

4. Stage II: Contrastive Learning

4.1 Motivation

仅靠监督标签，模型仍可能把不同 view 下的同一 workload 当作不同样本。为使表示空间具备“语义相同、名称不同”的聚合性质，框架引入对比学习。

4.2 Formulation

在 index_advisor/contrastive_loss.py 中实现了 InfoNCE 风格损失，在 scripts/sft_projector.py 中通过 contrastive_weight 启用。训练时，把同一个 workload 的不同扰动视图作为正样本，把批内其他槽位或其他样本作为负样本，从而优化：

  • 同一语义 workload 在不同 view 下的表示接近
  • 不同 workload 的表示保持可分

该机制尤其适用于在线扰动和 grouped sampling 联合训练，即在一个 batch 中有意识地放入同源 workload 的多个视角，直接约束表示层的一致性。

4.3 Role in the Pipeline

对比学习不直接生成索引答案，而是为后续 projector 与 SFT 提供更稳定的结构表示。它本质上是对“去记忆、保结构”的再强化。

5. Stage III: PJ Training

5.1 What PJ Means Here

本文中的 PJ 指 projector training。框架支持把外部结构表示映射到 LLM 隐空间，作为文本 token 的替代输入，从而把 schema/SQL 侧信息以多模态方式注入语言模型。

5.2 Multi-Modal Inputs

系统中主要有两类外部 embedding：

  1. Column embeddings，对应列语义与表结构信息。
  2. SQL embeddings，对应扰动后的 SQL 或 workload 结构表示。

在 index_advisor/mm/model.py 中，分别通过特殊 token 注入，例如 <|column|> 和 <|sql|>。这些 embedding 先经过 projector，再送入 LLM hidden space。

5.3 Why Separate PJ Training

直接让 LLM 同时学习文本、索引策略和外部 embedding 对齐，优化难度较大。将 PJ 作为独立阶段有三点好处：

  • 先把外部结构信号稳定投影到 LLM 空间
  • 降低全参数联合训练的不稳定性
  • 为后续 LoRA-SFT 或 RL 提供可复用的多模态接口

在实现上，scripts/sft_projector.py 支持：

  • projector-only 训练
  • projector + LoRA-LM 联合训练
  • projector checkpoint 保存与恢复
  • 与对比学习、在线扰动同时启用

因此，PJ 阶段起到承上启下的作用：它把“结构表示学习”变成“可被 LLM 消化的输入通道”。

6. Stage IV: SFT

6.1 Objective

SFT 阶段解决的是监督对齐问题：让模型学会从 workload/schema/prompt 生成格式正确、语义合理的索引建议。

训练样本通常由以下信息组成：

  • workload 与 SQL 集合
  • schema 信息及列统计
  • labels，即离线标注或专家系统生成的推荐索引
  • 可选的 column embeddings 与 SQL embeddings

6.2 Training Pattern

在本框架中，SFT 有两种典型路径：

  1. 纯文本或 LoRA-SFT，主要学习输出格式和索引选择模式。
  2. 带 projector 的多模态 SFT，同时利用文本 prompt 与结构 embedding。

此阶段的主要收益是把模型从“知道结构相似”推进到“能产出正确动作”。从学习角度看，SFT 是行为克隆；从系统角度看，它是 RL 的 warm start。

6.3 Why SFT Is Still Necessary

即便已经完成扰动和对比学习，如果没有 SFT，模型也只是在表示层更稳，但并不一定会输出合法、紧凑且符合任务协议的索引集合。SFT 因而负责：

  • 约束输出格式
  • 学习索引推荐的初始策略
  • 缩小 RL 搜索空间
  • 提高 reward 稀疏场景下的训练稳定性

7. Stage V: RL

7.1 Motivation

SFT 依赖静态标签，但静态标签不一定对应真实执行最优策略。索引推荐本质上是一个带环境反馈的决策问题，因此最终需要用数据库代价收益来训练策略。

7.2 Reward Design

在 lmf_hooks/reward.py 中，框架提供了多种 reward：

  • format：只检查输出格式是否合法
  • cost：根据索引带来的执行代价下降计算奖励
  • div：在收益基础上对索引数量做归一化，抑制冗余索引

核心思想是把模型输出的索引集合送入数据库执行器，比较加索引前后的 workload cost，再将收益映射为训练奖励。这样优化目标从“模仿标签”转变为“提升真实系统指标”。

7.3 Optimization

scripts/train.py 支持 PPO 与 GRPO 两类策略优化器：

  • PPO 适合以 SFT/LoRA checkpoint 为初始化，稳定做 reward-driven 微调
  • GRPO 适合在组内比较多个生成结果，以相对优势更新策略

RL 阶段通常接在 SFT 之后，因为：

  • SFT 提供合法且可用的初始策略
  • RL 直接优化 cost reduction，而不是仅优化 token likelihood
  • 两阶段联合可兼顾输出稳定性与真实收益

8. End-to-End Interpretation

将整个框架串起来看，其逻辑可以概括为：

  • 扰动阶段解决“不要记住名字”
  • 对比学习解决“同一语义应当有相近表示”
  • PJ 训练解决“如何把结构侧信息稳定注入 LLM”
  • SFT 解决“如何输出一个像样的索引策略”
  • RL 解决“如何把策略继续推向真实成本最优”

这五个阶段并非孤立堆叠，而是严格前后依赖：没有扰动，对比学习与泛化目标会被削弱；没有 PJ，多模态结构信息难以有效进入 LLM；没有 SFT，RL 初始策略过弱；没有 RL，模型最终仍停留在标签模仿而非系统最优。

9. Implementation Mapping

当前仓库中的关键实现对应关系如下：

  • 数据扰动：index_advisor/perturb.py，scripts/gen/perturb_workloads.py
  • workload 加载与在线扰动恢复：index_advisor/workload.py
  • 对比损失：index_advisor/contrastive_loss.py
  • PJ/SFT 主训练入口：scripts/sft_projector.py
  • 多模态模型注入：index_advisor/mm/model.py，lmf_hooks/model.py
  • RL 训练入口：scripts/train.py
  • reward 计算：lmf_hooks/reward.py

这些模块共同组成了一个从数据增强、表示学习、监督对齐到策略优化的完整实验框架。

10. Conclusion

本实验框架的核心贡献不只是“用 LLM 做索引推荐”，而是构建了一条面向泛化与真实收益优化的完整训练链路。通过将训练流程拆分为扰动、对比学习、PJ、SFT 与 RL 五个阶段，系统能够逐步解决 schema memorization、表示不稳定、结构信号注入困难、监督偏差以及静态标签次优等问题。