PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhijian Li, Stefan Larson, Kevin Leach

arXiv:2609.11029v1cs.CLcs.LG

TL;DR

Uniform token weighting can encourage memorization of surface-level text, motivating an objective that gives different importance to tokens. The paper applies TF-IDF-weighted cross-entropy, which preserves supervision while emphasizing informative tokens, and reports lower memorization with preserved perplexity and downstream performance.

  • Problem

    Uniform token weighting allows frequent, low-information tokens to dominate learning and can increase memorization of surface-level text spans.

  • Method

    TF-IDF-weighted cross-entropy rescales token-level gradients using corpus statistics, upweighting distinctive tokens and down-weighting ubiquitous ones while retaining supervision for all tokens.

  • Results

    Across five decoder-only LLMs, TF-IDF consistently reduces memorization while preserving perplexity and downstream summarization and QA performance; average LMS reduction is 14% with LoRA and 58% for full-weight TinyLLaMA 1.1B.

  • Takeaways & Limitations

    The architecture-agnostic objective offers a lightweight way to mitigate memorization without changing standard training procedures or downstream task performance.

  • Takeaways & Limitations

    The evaluation covers verbatim exact-match recall, not semantic paraphrasing or adversarial extraction, and does not test training from scratch.

Abstract

from arXiv · show

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.

1 Introduction

Uniform token weighting can make language models memorize repeated surface patterns, so this work introduces TF-IDF weighting to redirect learning toward informative tokens while preserving standard supervision. Across five decoder-only models, the method reduces memorization while preserving perplexity and downstream performance.

  • Existing alternatives include costly deduplication, utility-trading differential privacy, post-hoc interventions, and token dropping that removes supervision rather than reweighting all tokens.The proposed objective instead modifies learning during continued pretraining and fine-tuning.
  • TF-IDF-weighted cross-entropy rescales token-level gradients, upweighting distinctive tokens and down-weighting globally ubiquitous ones while retaining supervision for every token.The method is architecture-agnostic and integrates into existing training pipelines.
  • Across five decoder-only LLMs from 1.1B to 13B parameters, TF-IDF preserves perplexity and summarization and QA performance while consistently reducing substring memorization and ROUGE-L.The evaluation spans models with different parameter scales.
  • 58%: TinyLLaMA 1.1B achieves the largest average LMS reduction under full-weight fine-tuning, while LoRA produces a 14% average reduction across five models.The reduction is larger when the CE baseline memorizes more.

2 Related Work

Prior memorization mitigations act on data, optimization, post-hoc editing, or selective supervision, while this work reweights every token using static corpus statistics. TF-IDF therefore targets token informativeness without auxiliary models or iterative meta-optimization.

  • Existing mitigation strategies include deduplication, differential privacy, unlearning, model editing, and token dropping, each operating at different stages or with different supervision trade-offs.These approaches span data-level, optimization-level, post-hoc, and objective-level interventions.
  • TF-IDF weighting differs from token dropping by retaining supervision on every token while redirecting gradient pressure away from common surface tokens.This preserves full-sequence modeling rather than excluding selected positions.
  • Regularization methods constrain models or soften objectives but generally treat the loss landscape uniformly instead of accounting for token-specific memorization tendencies.The proposed approach rebalances what the model prioritizes rather than only constraining how it learns.
  • Related token-weighting methods use dynamic model-dependent metrics, reference models, or meta-optimization, whereas TF-IDF uses static corpus statistics.The paper presents this as a computationally efficient alternative without secondary models or iterative meta-optimization.

3 TF-IDF Weighted Loss

The proposed loss multiplies each token’s autoregressive cross-entropy term by a TF-IDF-derived weight, normalized to preserve the mean learning scale. Local frequency and corpus-level rarity concentrate gradient pressure on informative content while common tokens receive less emphasis.

  • The TF-IDF objective retains supervision for every token but rescales each cross-entropy term by a token-specific weight w_i.Standard cross-entropy is recovered when w_i = 1.
  • 3.1 Buffer-Averaged TF-IDF Statistics: Weights combine local term frequency with smoothed inverse document frequency, emphasizing tokens frequent within a sequence but rare across the accumulation buffer.Smoothed IDF keeps weights positive for near-universal tokens.
  • 3.1 Buffer-Averaged TF-IDF Statistics: K = 16 mini-batches provide the statistics buffer, balancing reweighting strength and frequency-estimation stability without global pre-computation.The buffer contains N = B × K sequences.
  • 3.1 Buffer-Averaged TF-IDF Statistics: Raw TF-IDF weights are normalized to unit mean across non-padding tokens, ensuring E[w] = 1 so standard learning rates remain applicable.The raw weight is defined as w′_i = tf_i · idf(x_i).
  • 3.2 Weighting Dynamics: In context, tokens such as relativity and gravity receive higher weights, whereas ubiquitous function words, formatting symbols, and boilerplate receive lower IDF weights.This shifts gradient mass toward semantically meaningful content without removing token supervision.

4 Experimental Setup

Experiments compare CE and TF-IDF objectives across five pretrained decoder-only models, using LoRA as the primary adaptation method and full-weight training on TinyLLaMA as a check. Evaluation covers memorization, perplexity, summarization, and QA under matched protocols.

  • Models and Training Regimes: The study spans five decoder-only LLMs from 1.1B to 13B parameters, enabling evaluation across model scales from TinyLLaMA through LLaMA-2 13B.All experiments initialize from publicly released pretrained checkpoints.
  • Models and Training Regimes: Primary experiments use LoRA with frozen original weights, while full-weight fine-tuning updates all parameters on TinyLLaMA 1.1B to test whether findings depend on parameterization.Full-weight training uses a reduced learning rate because the LoRA rate caused catastrophic memorization.
  • Evaluation Categories: Memorization is measured using prefix matches, Longest Memorized Substring, and ROUGE-L, with lower values indicating less verbatim recall.Models are probed from prefixes of 32, 50, and 100 tokens across checkpoints from 10% to 100% training progress.
  • Evaluation Categories: Perplexity is measured on the disjoint WikiText-2 validation split, while summarization uses CNN/DailyMail and QA uses SQuAD v1.1.Summarization reports ROUGE and BERTScore-F1; QA reports Exact Match and F1.
  • Evaluation Categories: Figure 2 tracks average LMS over training progress for smaller, mid-size, and larger models under CE and TF-IDF objectives.The figure is designed to compare memorization trajectories across model groups and objectives.

5 Does TF-IDF-Weighted Loss Reduce Memorization?

TF-IDF weighting consistently reduces substring memorization, with the largest effect when full-weight fine-tuning allows greater memorization pressure. The method retains supervision for all tokens while lowering memorization relative to standard cross-entropy.

  • 14%: Across all five LoRA models, TF-IDF reduces average LMS while maintaining a stable lower memorization trajectory.TF-IDF produces lower LMS at every training checkpoint, with the gap widening for larger models as training progresses.
  • Over 20%: GPT-J 6B’s average LMS falls from 4.46 to 3.55, while its maximum memorized span drops from 21 to 10 tokens.LLaMA-2 7B, LLaMA-2 13B, and Pythia 1.4B also show reductions in average LMS.
  • 58%: Full-weight TinyLLaMA 1.1B reduces average LMS from 8.30 to 3.50, with maximum span falling from 128 to 18 tokens.The relative reduction is statistically significant, with a 95% bootstrap confidence interval of [34.9%, 71.1%].
  • A qualitative example shows CE reproducing an injected continuation verbatim, whereas TF-IDF generates a related but non-memorized continuation.The comparison uses a 32-token prefix and 128 generated tokens under full-weight fine-tuning.
  • TF-IDF retains supervision for every token, redirecting gradient pressure toward distinctive content rather than removing token supervision.This design contrasts with token-dropping approaches and is intended to reduce exact surface-form memorization.

6 Impact on Language Modeling and Downstream Performance

TF-IDF weighting preserves language-modeling and downstream utility while reducing memorization. Across perplexity, summarization, and QA, results remain comparable to or better than the cross-entropy baseline.

  • Perplexity improves for five of six configurations, while remaining stable or slightly improved under TF-IDF weighting.The sole exception is Pythia 1.4B, which increases from 17.23 to 18.36; full-weight TinyLLaMA improves from 20.94 to 16.06.
  • 6.2 Summarization: Summarization scores diverge by at most 0.5 ROUGE points, with TF-IDF and CE producing nearly indistinguishable results across models.LLaMA-2 7B achieves 33.0 versus 32.9 ROUGE-1, and ROUGE agrees with BERTScore-F1 across models.
  • 6.3 QA: QA differences stay within 1.5% EM and 0.6% F1, preserving performance across model scales.For full-weight TinyLLaMA, CE and TF-IDF achieve EM/F1 scores of 79.0/84.46 and 78.4/84.18, respectively.
  • TF-IDF weighting produces results within noise of CE across perplexity, summarization, and QA while consistently reducing memorization.The reported utility preservation accompanies reductions in verbatim memorization rather than a measurable downstream-performance penalty.

7 Conclusion

The paper introduces TF-IDF-weighted cross-entropy to prioritize informative tokens and reports reduced memorization without compromising downstream performance. The objective is architecture-agnostic and incurs less than 3% computational overhead.

  • TF-IDF-weighted cross-entropy rebalances learning toward semantically rich tokens instead of frequent, low-information tokens.The objective changes token importance while retaining the paper’s focus on mitigating verbatim memorization.
  • Across five decoder-only LLMs, reweighting reduced Longest Memorized Substring and ROUGE-L without compromising linguistic fluency or summarization and QA performance.Training-checkpoint analysis showed a stable downward shift rather than merely delayed memorization onset.
  • The architecture-agnostic objective adds less than 3% computational overhead, supporting integration into standard training pipelines.The reported TinyLLaMA comparison was 256 ms versus 262 ms per step.

Limitations

The evaluation is limited to specific memorization, training, tokenization, and context-length settings. These boundaries leave semantic and adversarial extraction, pretraining-from-scratch behavior, subword weighting, and longer-range dynamics for future study.

  • Exact-match memorization evaluation excludes semantic paraphrasing and adversarial prompting for training-data extraction.The scope prioritizes verbatim recall as the direct privacy and copyright risk examined here.
  • The experiments use LoRA fine-tuning predominantly and validate full-weight training only on TinyLLaMA 1.1B, not training from scratch.The authors note that down-weighting function words could potentially hinder early grammatical-structure acquisition from scratch.
  • BPE subword weighting may dilute intended upweighting when rare words split into individually common subword units.The paper identifies empirical subword-weight analysis as future work.
  • A 256-token block size may miss long-range memorization dynamics in models with larger context windows.The block-size constraint arose from hardware limitations.

Ethics Statement

The experiments follow the ACL Ethics Policy and use publicly available datasets according to their intended research use. The proposed loss function is aimed at enhancing LLM privacy and safety by reducing verbatim memorization.

  • The study uses publicly available CNN/DailyMail, SQuAD, and WikiText-2 datasets in accordance with their intended research use.
  • The proposed loss function is intended to enhance LLM privacy and safety by reducing verbatim memorization.

AI Usage Disclosure

The authors used ChatGPT and Claude for linguistic polishing, grammatical refinement, and debugging custom Python scripts. Human authors performed the core research design, derivation, experiments, and interpretation.

  • ChatGPT and Claude assisted with linguistic polishing, grammatical refinement, and debugging the TF-IDF loss and memorization-evaluation scripts.
  • Human authors solely performed the core research objectives, buffer-averaged TF-IDF loss derivation, experimental design, and result interpretation.

A Per-Model Training Configuration

The memorization experiments use a shared LoRA training setup across models, with per-model configuration recorded in Table 6 and architecture-specific target modules listed separately.

  • All models use rank r=8, α=32, learning rate 1 × 10−4, block size 256, 1 epoch, and AdamW.
  • Table 6 provides per-model LoRA configurations for the memorization experiments.
  • LoRA target modules are specified separately for each model architecture in Table 7.
Loading 2609.11029v1…