# Token Weighting for Long-Range Language Modeling

arXiv: https://arxiv.org/abs/2503.09202 · HTML: https://mlsanity.com/paper/2503.09202/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 50 references that are in the map.
- First posted on arXiv 12 Mar 2025.

## Builds on (most cited first, 30 of 35)

- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Extending Context Window of Large Language Models via Positional Interpolation](https://mlsanity.com/paper/2306.15595/) (2023)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [Learning to Reweight Examples for Robust Deep Learning](https://mlsanity.com/paper/1803.09050/) (2018)
- [Textbooks Are All You Need](https://mlsanity.com/paper/2306.11644/) (2023)
- [Weak-to-Strong Generalization: Eliciting Strong Capabilities With Weak Supervision](https://mlsanity.com/paper/2312.09390/) (2023)
- [Effective Long-Context Scaling of Foundation Models](https://mlsanity.com/paper/2309.16039/) (2023)
- [Pretraining Language Models with Human Preferences](https://mlsanity.com/paper/2302.08582/) (2023)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [Landmark Attention: Random-Access Infinite Context Length for Transformers](https://mlsanity.com/paper/2305.16300/) (2023)
- [Data Engineering for Scaling Language Models to 128K Context](https://mlsanity.com/paper/2402.10171/) (2024)
- [QuRating: Selecting High-Quality Data for Training Language Models](https://mlsanity.com/paper/2402.09739/) (2024)
- [D4: Improving LLM Pretraining via Document De-Duplication and Diversification](https://mlsanity.com/paper/2308.12284/) (2023)
- [How to Train Long-Context Language Models (Effectively)](https://mlsanity.com/paper/2410.02660/) (2024)
- [At Which Training Stage Does Code Data Help LLMs Reasoning?](https://mlsanity.com/paper/2309.16298/) (2023)
- [Can Perplexity Reflect Large Language Model's Ability in Long Text Understanding?](https://mlsanity.com/paper/2405.06105/) (2024)
- [A Controlled Study on Long Context Extension and Generalization in LLMs](https://mlsanity.com/paper/2409.12181/) (2024)

## Built on by (most cited first, 14 of 14)

- [TAB-PO: Preference Optimization with a Token-Level Adaptive Barrier for Token-Critical Structured Generation](https://mlsanity.com/paper/2603.00025/) (2026)
- [GoLongRL: Capability-Oriented Long Context Reinforcement Learning with Multitask Alignment](https://mlsanity.com/paper/2605.19577/) (2026)
- [Revisiting Long-context Modeling from Context Denoising Perspective](https://mlsanity.com/paper/2510.05862/) (2025)
- [Rethinking Supervised Fine-Tuning: Emphasizing Key Answer Tokens for Improved LLM Accuracy](https://mlsanity.com/paper/2512.21017/) (2025)
- [Where Does Long-Context Supervision Actually Go? Effective-Context Exposure Balancing](https://mlsanity.com/paper/2605.10544/) (2026)
- [Turbo-ICL: In-Context Learning-Based Turbo Equalization](https://mlsanity.com/paper/2505.06175/) (2025)
- [Entropy-Guided Loop: Achieving Reasoning through Uncertainty-Aware Generation](https://mlsanity.com/paper/2509.00079/) (2025)
- [Short-Context Dominance: How Much Local Context Natural Language Actually Needs?](https://mlsanity.com/paper/2512.08082/) (2025)
- [Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models](https://mlsanity.com/paper/2604.02340/) (2026)
- [Supervised Fine-Tuning Needs to Unlock the Potential of Token Priority](https://mlsanity.com/paper/2602.01227/) (2026)
- [QAQ: Bidirectional Semantic Coherence for Selecting High-Quality Synthetic Code Instructions](https://mlsanity.com/paper/2603.12165/) (2026)
- [ARCA: Adapter-Residual Credit Assignment When Token Signals Degenerate](https://mlsanity.com/paper/2606.00257/) (2026)
- [Decoupling Token Roles in Autoregressive Pretraining](https://mlsanity.com/paper/2609.33405/) (2026)
- [What Shared Prefixes Hide: Trajectory Dropout for On-Policy Distillation](https://mlsanity.com/paper/2609.33455/) (2026)

## Abstract (arXiv)

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.