Exploring Next Token Prediction For Optimizing Databases

TL;DR

This work investigates applying Next Token Prediction (NTP) to database optimization by introducing the Probe and Learn (PoLe) framework, which combines Decision Transformers with PMU-derived DB-tokens to learn hardware-aware policies. PoLe employs offline reinforcement-learning–style training on diverse policies and hardware/workload configurations, enabling autoregressive policy generation that can generalize to unseen environments. A NUMA-aware case study on main-memory indexing demonstrates substantial performance gains (up to about $3\times$) and cross-hardware generalization, illustrating PoLe's potential to drive data-driven DBMS optimization across heterogeneous platforms. The approach advances the notion of generalizable, hardware-informed optimization in databases, with broad implications for adaptive tuning in cloud and multi-architecture environments.

Abstract

The Next Token Prediction paradigm (NTP, for short) lies at the forefront of modern large foundational models that are pre-trained on diverse and large datasets. These models generalize effectively, and have proven to be very successful in Natural Language Processing (NLP). Inspired by the generalization capabilities of Large Language Models (LLMs), we investigate whether the same NTP paradigm can be applied to DBMS design and optimization tasks. Adopting NTP directly for database optimization is non-trivial due to the fundamental differences between the domains. In this paper, we present a framework, termed Probe and Learn (PoLe), for applying NTP to optimize database systems. PoLe leverages Decision Transformers and hardware-generated tokens to effectively incorporate NTP into database systems. As a proof of concept, we demonstrate PoLe in the context of the index scheduling task over NUMA servers in main-memory database systems. Preliminary results for this scheduling task demonstrate that adopting NTP and PoLe can improve both performance and generalizability.

Figures (5)

• Figure 1: The training and the inference phases of NTP.
• Figure 2: The Probe and Learn (PoLe) framework.
• Figure 3: Performance of a main-memory B$^+$-Tree under various scheduling policies.
• Figure 4: The scheduling policies learned by PoLe. Each grid corresponds to the learned scheduling policy for the corresponding server for the YCSB-A workload. The top-left grid cell corresponds to the 1st index chunk. The first cell in $i$-th row corresponds to the $\langle i*16+1\rangle$-th index chunk. The corresponding text of a grid cell signifies the hardware core where the query for that particular index chunk is scheduled. Cells sharing the same color indicate that the associated index chunks are scheduled on the same NUMA server.
• Figure 5: Top performing scheduling policies present in the training dataset for the Intel Sandy Bridge and NVIDIA Grace Hopper servers. Each grid corresponds to a scheduling policy for the corresponding server. The same scheduling policy can be present multiple times in the training dataset for different workloads. Cells sharing the same color indicate that the associated index chunks are scheduled on the same NUMA server.