Token Weighting for Long-Range Language Modeling

TL;DR

The paper tackles the difficulty of long-context understanding in large language models by introducing non-uniform token weights in the training loss. It presents a two-step framework—token scoring that contrasts short- and long-context confidences and subsequent postprocessing to produce dense or sparse weights—and investigates multiple base-model configurations and data scenarios. Empirical results show that dense weighting improves long-context performance across tasks, while sparse weighting specializes the model for long-range, retrieval-heavy tasks, with notable differences between frozen and unfrozen scoring approaches. The work provides practical guidelines for steering model behavior through loss weighting, demonstrates generalizability to different model sizes and datasets, and discusses trade-offs and limitations for future research in long-context language modeling.

Abstract

Many applications of large language models (LLMs) require long-context understanding, but models continue to struggle with such tasks. We hypothesize that conventional next-token prediction training could contribute to this, because each token is assigned equal weight. Yet, intuitively, the amount of context needed to predict the next token accurately varies greatly across different data. To reflect this, we propose various novel token-weighting schemes that assign different weights to each training token in the loss, thereby generalizing existing works. For this, we categorize token-weighting methods using a two-step framework which compares the confidences of a long-context and short-context model to score tokens. We evaluate all methods on multiple long-context understanding tasks and show that non-uniform loss weights are helpful to improve the long-context abilities of LLMs. Different short-context models can be used effectively for token scoring, including models that are much smaller than the long-context model that is trained. All in all, this work contributes to a better understanding of the trade-offs long-context language modeling faces and provides guidelines for model steering via loss-weighting based on empirical evidence. The code can be found on Github.

Figures (5)

• Figure 1: We compare different $\lambda$ for the dense unfrozen setting with Llama-3 8B (\ref{['eq:interpolation']}). $\lambda=1$ coincides with the baseline (dotted). While short-context benefits from large $\lambda$, for long-context $\lambda = 0.25$ is best.
• Figure 2: We compare sparsity values $\kappa$ for the sparse unfrozen setting with Llama-3 8B (\ref{['subsec: loss weighting']}). $\kappa=1$ coincides with the baseline (dotted). Increasing $\kappa$ benefits short-context tasks but hurts long-context capabilities.
• Figure 3: A kernel density estimate of the PMI distribution of all tokens in PG19-eval, measured with the uniform 32k model at the end of training.
• Figure 4: We compare different interpolation values $\lambda$ for the dense frozen setting with Llama-3 8B (\ref{['eq:interpolation']}). $\lambda=1$ coincides with the baseline (dotted). While short-context benefits from large $\lambda$, for long-context it is not clear how to choose $\lambda$.
• Figure 5: We compare different sparsity values $\kappa$ for the sparse frozen setting with Llama 3 8B (\ref{['subsec: loss weighting']}). $\kappa=1$ coincides with the baseline (dotted). Performance is robust to the choice of $\kappa$.

Theorems & Definitions (2)

• proof
• proof