Exploring Next Token Prediction For Optimizing Databases

2025-03-25

AI Summary

This work investigates adapting the next-token prediction (NTP) paradigm from large language models to database management system (DBMS) optimization tasks to improve both performance and generalization. To this end, we propose PoLe (Probe and Learn), an end-to-end framework that systematically introduces NTP into core DBMS optimization—its first such application. Our key contributions are: (1) a hardware-aware tokenization scheme that encodes low-level operations (e.g., in-memory index scheduling) as learnable sequences; and (2) a Decision Transformer–based policy learning architecture that unifies reinforcement learning–inspired heuristics with sequential decision modeling. Evaluated on main-memory index scheduling, PoLe achieves significantly higher throughput and improved latency stability compared to baselines. It also demonstrates superior generalization over traditional rule-based engines and supervised learning approaches, empirically validating the feasibility and advantages of applying the NTP paradigm to DBMS optimization.

Technology Category

• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP
• Planning, Routing, and Scheduling: Learning for Planning and Scheduling

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

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 also 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. Preliminary results from the main-memory index scheduling task demonstrate that adopting NTP can improve both performance and generalizability.

Problem

• Applying Next Token Prediction to database optimization tasks
• Addressing domain differences between NLP and DBMS for NTP adoption
• Improving performance and generalizability in main-memory index scheduling

Innovation

• Applies Next Token Prediction to databases
• Uses Probe and Learn framework (PoLe)
• Leverages Decision Transformers and hardware tokens