Rebalancing Token Importance in Language Models with TF-IDF Weighted Cross-Entropy Loss

2026-09-09

AI Summary

本文通过TF-IDF加权的交叉熵损失方法重新平衡语言模型中的词元重要性，以减少模型对表面文本的记忆，同时保持困惑度和下游任务性能。

Abstract

Large language models are typically trained under uniform token weighting, which allows frequent and low-information tokens to dominate learning and can increase the tendency to memorize surface-level text spans. To address this, we present an information-weighted cross-entropy loss that rescales token-level contributions using TF-IDF statistics, emphasizing semantically informative tokens while down-weighting ubiquitous ones. Experiments on five decoder-only LLMs ranging from 1.1B to 13B parameters show consistent reductions in memorized substring length while preserving perplexity and downstream task performance. Under LoRA fine-tuning, TF-IDF reduces average substring memorization length by 14% across all five models. Under full-weight fine-tuning on TinyLLaMA 1.1B, the reduction reaches 58%. Our approach is architecture-agnostic and can be incorporated into existing training pipelines with less than 3% computational overhead, offering a lightweight and principled way to mitigate memorization without disrupting standard training dynamics.

Problem

• large language models
• uniform token weighting
• memorization
• surface-level text spans

Innovation

• TF-IDF weighted cross-entropy loss
• token importance rebalancing
• memorization reduction