Next Token Prediction在数据库优化中的创新应用
1. 从语言模型到数据库优化：Next Token Prediction的跨界应用
在自然语言处理（NLP）领域，Next Token Prediction（NTP）技术已经成为大型语言模型（如GPT系列）的核心机制。这项技术通过预测序列中的下一个token（可以是词、子词或字符），实现了文本生成、机器翻译等任务的突破性进展。但鲜为人知的是，这一看似专属于NLP的技术范式，正在数据库系统优化领域展现出惊人的潜力。
传统数据库管理系统（DBMS）面临着硬件和工作负载双重异构性带来的严峻挑战。现代服务器架构日趋复杂，NUMA（非统一内存访问）架构、多级缓存层次、混合内存技术（如DRAM与PMEM）的普及，使得数据访问延迟差异可达4倍之多。同时，工作负载也从传统的OLTP、OLAP扩展到HTAP乃至新兴的LLM负载。这种复杂性使得传统的基于规则的优化器往往难以应对。
普渡大学的研究团队提出了一个大胆的设想：能否将NTP范式迁移到数据库优化领域？他们的研究成果——Probe and Learn（PoLe）框架，给出了肯定的答案。通过结合决策变换器（Decision Transformers）和硬件生成的DB-tokens，PoLe在NUMA服务器上的索引调度任务中实现了高达3倍的性能提升，更令人振奋的是，这种优化能力还能泛化到训练阶段未见过的硬件环境。
2. NTP在数据库优化中的核心挑战
2.1 领域差异与目标不匹配
直接将NTP应用于数据库优化面临几个根本性挑战。首先，NLP中的token具有相对固定的语义和语法规则。例如在句子"SIGMOD'25将在德国举行"中，"德国"作为地点名词的出现，自然排除了其他地点选项。而数据库中的"token"概念则高度依赖上下文——在连接顺序选择中，表名可作为token；在查询调度中，查询本身又成为token。这种多变性使得模型难以学习统一的表示。
更重要的是目标函数的差异。NLP任务通常没有明确的优化目标，而数据库优化则具有明确的指标导向——最小化查询延迟、最大化吞吐量或提高资源利用率。简单的token预测无法自动满足这些性能目标。以四表连接（A⋈B⋈C⋈D）为例，NTP可以预测下一个连接的表，但无法保证生成的连接顺序是最优的。
2.2 缺乏有效的token表示
数据库领域缺乏像NLP那样统一的token表示方法。不同数据库实例的表结构完全不同，使得模型难以学习通用的表表示。即使对于相同的表集合（如A、B、C、D），不同的连接顺序（A⋈B⋈C⋈D vs D⋈C⋈A⋈B）可能都是语法上有效的，但性能差异可能很大。传统的NTP范式无法捕捉这种性能差异。
3. PoLe框架的核心构建块
3.1 决策变换器（Decision Transformers）
为解决目标导向的问题，PoLe采用了决策变换器这一强化学习架构。与标准NTP不同，决策变换器在预测下一个动作（token）时，会考虑期望的累积奖励（Reward-to-Go，RTG）。这使得模型能够将token预测与具体的优化目标（如查询吞吐量）对齐。
决策变换器的输入序列由三种token交替组成：
- RTG token（ˆr）：从当前时刻开始的预期累积奖励
- 状态token（s）：系统当前状态表示
- 动作token（a）：要执行的动作
训练时，模型通过监督学习最大化动作预测的准确性；推理时，则以自回归方式生成动作序列，同时根据实际获得的奖励动态调整RTG。
3.2 DB-tokens：硬件性能监控单元
PoLe最具创新性的贡献是提出了DB-tokens的概念——直接从硬件性能监控单元（PMU）采集的硬件性能指标。这些指标包括：
- 各级缓存（L1/L2/LLC）的访问与缺失次数
- 分支预测错误数
- 本地与远程内存访问次数
- 执行指令数
- 核心时钟周期数
关键洞见：硬件PMU指标虽然底层，但能精确反映查询执行时的系统行为。例如频繁访问的数据区域会导致特定核心的缓存活动激增，而分析型查询则表现为高指令数和内存访问。
DB-tokens具有三大优势：
- 计算开销低：直接从硬件寄存器读取，几乎不影响查询性能
- 通用性强：不同硬件平台都提供类似的PMU计数器
- 信息丰富：能捕捉工作负载特征和硬件状态
4. PoLe框架的运作机制
4.1 探测（Probe）阶段
PoLe首先通过广泛的探测收集训练数据：
- 策略选择：从策略池中采样不同的调度策略（包括最优和次优策略）
- 硬件配置：在不同硬件（如Intel Sandy Bridge、NVIDIA Grace Hopper）上运行
- 工作负载：执行多样化的负载（OLTP、OLAP、HTAP等）
在执行过程中，系统会：
- 定期采集PMU指标和性能数据（吞吐量、延迟等）
- 将策略编码为动作token，与DB-tokens一起存入离线数据集
- 控制数据采集频率，避免引入过大开销
4.2 学习（Learn）阶段
4.2.1 离线训练
使用收集的离线数据集训练决策变换器，关键步骤包括：
- 序列建模：将每个训练样本组织为(RTG, 状态, 动作)的序列
- 教师强制（Teacher Forcing）：使用真实动作进行训练，避免误差累积
- 目标优化：最小化动作预测的交叉熵损失
训练后的模型能够理解不同策略在各种硬件和工作负载下的性能表现。
4.2.2 在线推理
面对新环境时，PoLe的推理过程如下：
- 初始化：输入初始硬件状态和期望性能目标（RTG）
- 自回归生成：
- 模型预测下一个最佳动作（如将查询调度到特定核心）
- 执行动作并观察新状态和实际奖励
- 更新RTG并反馈给模型
- 终止：当生成完整策略序列或达到停止条件时结束
值得注意的是，推理时模型对实际硬件和工作负载一无所知，仅依靠DB-tokens和学到的知识进行决策。这使得PoLe能够泛化到训练阶段未见过的硬件环境。
5. 案例研究：NUMA服务器上的索引调度
5.1 问题背景
现代NUMA服务器的复杂性给数据库性能带来巨大挑战：
- 多插槽架构：数据访问延迟差异可达4倍（本地vs远程DRAM）
- 核心数量庞大：单服务器可达上千核心
- 资源竞争：并发查询会导致互连带宽、缓存和内存控制器饱和
传统调度策略（如操作系统默认的本地分配、交错分配等）往往无法充分利用硬件特性。PoLe将索引调度建模为NTP任务：
- 将B+树索引划分为多个chunk（如256个）
- 动作token：核心ID（决定每个chunk的查询在哪个核心执行）
- DB-tokens：PMU指标+吞吐量
- 目标：最大化整体吞吐量
5.2 实验设置
研究团队在三种硬件平台上进行评估：
- Intel Sandy Bridge：4插槽/4 NUMA节点/32核，DDR3内存
- NVIDIA Grace Hopper：单插件/72核，LPDDR5X内存（ARM架构）
- Intel Skylake X（未见过的硬件）：4插槽/92核，用于测试泛化能力
对比基线包括：
- OS:D - 操作系统默认本地分配策略
- OS:I - 交错分配策略
- SE:N - Shared Everything NUMA策略
- SN - Shared Nothing策略
5.3 性能结果
实验数据显示：
在训练见过的硬件（Sandy Bridge和Grace Hopper）上：
- PoLe相比最佳基线提升2.78倍吞吐量
- 生成的调度策略与训练集中的最优策略都不同（图4 vs 图5）
在未见过的Skylake X硬件上：
- 仍能实现3倍性能提升
- 证明NTP范式具有强大的泛化能力
实战建议：当部署数据库到新型硬件时，即使没有该平台的训练数据，PoLe仍可能提供优于传统策略的性能。这对于云环境特别有价值，因为实例类型经常更新。
6. 实施考量与优化方向
6.1 成本分析
PoLe的实施涉及两类主要成本：
离线数据集准备：
- 需要在不同硬件上执行多种策略和负载
- 实验中使用445个Sandy Bridge样本（约$614）和49个Grace样本（约$73）
- 每个样本运行约5分钟
模型训练：
- 使用3.37M参数的决策变换器
- 训练时间约26 CPU小时（约$32）
- 推理时间短（约1 CPU分钟）
6.2 优化策略
基于实践经验，我们总结出以下优化方向：
PMU指标选择：根据优化目标调整采集的指标。例如：
- NUMA优化：侧重本地/远程内存访问
- 连接排序：关注写操作（反映中间结果大小）
采样粒度：
- 调度任务：查询级别采样
- 连接排序：算子级别采样
- 数据分区：分区级别采样
动作空间设计：
- 保持动作（如核心ID）的离散性和合理性
- 过大动作空间会导致训练困难
数据集多样性：
- 包含不同性能水平的策略（不仅是最优策略）
- 覆盖尽可能多的工作负载和硬件配置
7. 技术展望与应用场景
NTP范式在数据库优化中的应用远不止于索引调度。潜在的应用场景包括：
查询优化：
- 连接顺序选择
- 访问路径选择
- 并行度调整
资源管理：
- 内存分配策略
- 线程调度
- 缓存管理
新兴硬件适配：
- CXL内存的利用
- 异构计算（CPU+GPU+加速器）调度
- 持久内存（PMEM）的使用策略
特别值得关注的是，随着云数据库的普及，数据库实例需要频繁迁移和扩展。PoLe的泛化能力使其成为云环境的理想选择——一个训练好的模型可以适应各种实例类型，无需为每种硬件配置重新优化。
从更长远看，这项研究代表了一个重要趋势：数据库系统正从基于规则的静态优化，转向基于学习的自适应优化。NTP范式与数据库的结合，可能催生出新一代的"自优化"数据库系统，能够自动适应不断变化的硬件和工作负载环境。