Reproduce “Rebalancing Token Importance in Language Models with TF-IDF Weighted Cross-Entropy Loss” on OpenResearch

About this paper

Zhijian Li, Stefan Larson, Kevin Leach

TF-IDF-weighted cross-entropy reduces the emphasis placed on common tokens while retaining supervision for every position, offering a lightweight way to limit verbatim memorization during language-model adaptation. Across five decoder-only models, LoRA experiments show roughly a 14% average reduction in longest memorized substrings, while full fine-tuning of TinyLLaMA yields a 58% reduction with under 3% reported runtime overhead and generally comparable downstream quality.

Key insights

• The intervention changes the learning signal rather than the model architecture or training data: common boilerplate receives less influence, but every token continues to contribute to the objective.
• Its effect depends on adaptation regime. LoRA already limits memorization in some settings, whereas TF-IDF weighting has a much larger impact when all model parameters are updated.
• Lower memorization is observed throughout training checkpoints, suggesting the method does more than merely delay exact recall.

Summary by alphaXiv.