Rebalancing Token Importance in Language Models with TF-IDF Weighted Cross-Entropy Loss
Zhijian Li, Stefan Larson, Kevin Leach
cs.CL · cs.LG
Quick verdict
Claims TF-IDF-weighted CE cuts LLM memorization (14% LoRA, 58% full-weight TinyLLaMA) at no utility cost; results credible, but baselines are missing.
Summary
The authors replace uniform token weighting in language-model training with a TF-IDF-weighted cross-entropy loss: term frequencies and a smoothed inverse document frequency are estimated over a rolling buffer of 16 mini-batches, weights are normalized to unit mean, and all tokens keep supervision. Five decoder-only models (TinyLLaMA 1.1B, Pythia 1.4B, GPT-J 6B, LLaMA-2 7B/13B) are fine-tuned on a controlled corpus where 100 WikiText-2 sequences are injected and repeated 10 times into 20,000 Pile sequences, then probed with Longest Memorized Substring (LMS), prefix-match, and ROUGE-L. Under LoRA, LMS falls for four of five models (e.g., LLaMA-2 13B 3.90 to 3.17, GPT-J 4.46 to 3.55) with bootstrap confidence intervals excluding zero; TinyLLaMA shows no significant change. Under full-weight fine-tuning of TinyLLaMA, average LMS drops 57.8% (8.30 to 3.50, CI [34.9%, 71.1%]). Perplexity, CNN/DailyMail summarization, and SQuAD QA stay within roughly noise of the CE baseline. Main limitations: no head-to-head comparison with competitor losses (goldfish, MiLe), a single synthetic 1-epoch injection setup, an unverified mechanism story, and an asserted-but-unsubstantiated novelty claim.
Review complete
Assessment at a glance
4
risk flags identified
● high 0
● medium 2
● low 2
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Replacing uniform token weighting with buffer-averaged TF-IDF weights in the cross-entropy loss consistently reduces memorized substring length across five decoder-only LLMs (1.1B–13B).
- The magnitude of memorization reduction scales with the baseline's memorization capacity: 14% average LMS reduction under LoRA, 58% under full-weight fine-tuning of TinyLLaMA 1.1B.
- The loss preserves utility: perplexity improves or stays stable in five of six configurations, and summarization/QA scores remain within roughly 0.5 ROUGE points and 1.5% EM of the CE baseline.
- The method is architecture-agnostic and adds less than 3% computational overhead (256 ms vs 262 ms per step on TinyLLaMA 1.1B, batch size 8, full-weight).
- Down-weighting ubiquitous tokens removes the incentive to memorize exact surface forms while retaining supervision on every token.
Strengths
Consistent memorization reductions across five architectures with proper uncertainty
Table 1 shows Avg LMS drops for all LoRA models except TinyLLaMA (GPT-J 4.46→3.55, LLaMA-2 7B 3.48→2.98, LLaMA-2 13B 3.90→3.17, Pythia 2.44→2.07), with percentile bootstrap CIs (n=10,000, 95%) excluding zero for Pythia [6.9%, 22.6%], LLaMA-2 7B [5.6%, 22.4%], and LLaMA-2 13B [10.4%, 25.9%]. The full-weight TinyLLaMA result (8.30→3.50, CI [34.9%, 71.1%]) is substantial, and the null LoRA-TinyLLaMA result is honestly reported with a CI spanning zero.
Honest effect-size framing tied to baseline memorization capacity
Section 5 explicitly attributes the small LoRA-TinyLLaMA effect (3.23→3.22) to LoRA already suppressing memorization, and validates the claim via the full-weight control (CE reaches Avg Prefix 5.65, Max 128) rather than hiding the null. Figure 2 additionally shows the gap is a stable offset across training checkpoints, not a delayed onset.
Controlled, like-for-like experimental design
Variance-reduction observation, not just mean reduction
Candid limitations section
Weaknesses
Uniform CE is the only baseline
Single synthetic memorization setup
Mechanism claim is asserted, not measured
Subword and buffer interactions unexamined
Novelty claim and 'preserves' framing are loose
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
The paper proposes replacing the uniform token weighting of standard cross-entropy (CE) with a TF-IDF-weighted CE objective for LLM continued pretraining and fine-tuning: per-token weights `w_i = tf_i · idf(x_i)`, normalized to unit mean over non-padding tokens (Eqs. 1–6), with TF counted within-sequence and smoothed IDF estimated over a rolling buffer of K=16 mini-batches rather than the full corpus. Three claims are at stake: (1) the loss reduces verbatim memorization consistently across five decoder-only models (1.1B–13B); (2) it preserves perplexity and downstream summarization/QA performance; (3) it is architecture-agnostic with <3% overhead. All three are supported by the presented data, though claim (3) rests on a single timing measurement and the design-space superiority over competing losses is asserted rather than tested.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 4 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is TF-IDF weighted cross-entropy loss for language models?
It is a modification of the standard training loss in which each token's contribution is multiplied by a TF-IDF-style weight: tokens that are frequent within a sequence but rare across recent training batches (distinctive content words) get larger weight, while ubiquitous tokens like 'the' get smaller weight. Weights are computed from a rolling buffer of 16 mini-batches and normalized so the average weight is 1, keeping standard learning rates usable.
How much does TF-IDF weighting reduce memorization in LLMs?
In this paper, LoRA fine-tuning across five models (1.1B–13B) yields about a 14% average reduction in longest memorized substring length, though TinyLLaMA under LoRA shows no significant change. Full-weight fine-tuning of TinyLLaMA 1.1B shows a 58% reduction (average 8.30 to 3.50 tokens, a statistically significant effect). The effect scales with how much the cross-entropy baseline memorizes.
Does TF-IDF weighted loss hurt model quality or downstream tasks?
The paper reports no meaningful cost: perplexity on WikiText-2 improves for five of six model configurations, CNN/DailyMail summarization stays within about 0.5 ROUGE points, and SQuAD QA stays within 1.5 points exact match and 0.6 F1 of the standard cross-entropy baseline. Exact match drifts slightly downward in all QA comparisons, which the authors describe as within noise.
How expensive is it to add TF-IDF weighting to LLM training?
The authors report less than 3% computational overhead, measured as 256 ms versus 262 ms per training step on TinyLLaMA 1.1B with batch size 8 under full-weight fine-tuning. Because weights come from corpus statistics in a rolling buffer rather than a reference model, it avoids the cost of auxiliary models or global corpus pre-computation.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- The GitHub repository URL is a hyperlink in the PDF and does not appear in the extracted text; its absence from this review's text is an extraction artifact, not proof the paper lacks it.
- Appendix A reports only the LoRA target-module and quantization tables; full-weight TinyLLaMA training details beyond batch size 8, block 256, and lr 5e-5 (e.g., schedule, weight decay) are not in the provided text.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.