Tracing LLM Behavior to the Training Data with Empirical Next-Token Distributions

2026-07-15

AI Summary

This study investigates the alignment between the next-token prediction distributions of large language models and the empirical next-token distribution (ENTD) observed in their training data. By constructing ENTD as an interpretable reference, the authors systematically compare output distributions across models of varying scales under diverse contextual conditions and perform attribution analyses grounded in Transformer architecture and training dynamics. The findings reveal that model predictions closely align with ENTD for most inputs, with alignment improving as model scale and training compute increase. The work also identifies and dissects significant deviations in long-tail contexts. Advocating a “data-centric mechanistic interpretability” paradigm, this research offers a data-provenance perspective for understanding model behavior.

Abstract

In this paper, we study the connection between an LLM's output distribution and the data used to train it. Specifically, we study the degree to which an LLM's next-token distribution agrees with the empirical next-token distribution (ENTD) given the context in the training data. The ENTD is an appealing target because it is the unrestricted global minimizer of the next-token cross entropy loss used for pretraining, as well as an easily interpretable function of the pretraining corpus. We find that for a significant fraction of inputs, the LLM's distribution agrees with the ENTD almost perfectly, and the average agreement increases with model scale and training compute. Nevertheless, there is a long tail of input sequences where the LLM and ENTD differ significantly, and we examine several possible sources of this discrepancy across the transformer architecture, training procedure, and finite-sample noise in the ENTD estimate itself. More broadly, we hope our findings will encourage more work on ``data-centric mechanistic interpretability,'' a complement to standard mechanistic interpretability that opens the black box of how model behaviors arise from the data, rather than how they are encoded in the learned weights.

Problem

• LLM behavior
• training data
• next-token distribution
• empirical distribution
• data-centric interpretability

Innovation

• Empirical Next-Token Distribution
• data-centric mechanistic interpretability
• training data tracing
• next-token prediction
• large language models