💬 NLP

Weighting What Matters: Boosting Sample Efficiency in Medical Report Generation via Token Reweighting

This paper proposes a token reweighting strategy that prioritizes semantically salient tokens in the loss function, enabling vision-language models to generate high-quality medical reports with up to ten times less training data compared to standard cross-entropy approaches.

Original authors: Alexander Weers, Daniel Rueckert, Martin J. Menten

Published 2026-04-24

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Alexander Weers, Daniel Rueckert, Martin J. Menten

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a student to write a medical report based on an eye scan. You have a limited amount of time and only a few practice exams to work with.

In a traditional classroom (standard AI training), the teacher grades every single word in the student's report with the same level of strictness. If the student writes "The eye looks okay" instead of "The eye looks normal," they get a small deduction. If they write "The eye has no damage" instead of "The eye has severe damage," they also get a small deduction. The teacher treats both mistakes as roughly equal errors.

The Problem:
In the real world, these two mistakes are not equal. The first one is just a minor wording issue. The second one is a life-altering error that could lead to a patient losing their sight. Because the teacher treats them the same, the student doesn't learn to prioritize the critical words.

The Solution: "Weighting What Matters"
This paper proposes a new way to teach the AI, which we can call "The Highlighter Method."

Instead of grading every word equally, the teacher uses a special red highlighter to mark specific, life-critical words before the student even starts writing. These words are things like "drusen" (a specific eye condition), "fluid," "severe," or "healthy."

Here is how the new method works:

The Spotlight: When the AI makes a mistake on a regular word (like changing "scan" to "image"), it gets a gentle tap on the wrist.
The Siren: When the AI makes a mistake on a highlighted, critical word (like confusing "no drusen" with "multiple drusen"), the penalty is massive. It's like hitting a siren. The AI feels this mistake much more acutely.
The Result: Because the AI feels the pain of missing the important words so much more, it learns to focus its energy on getting those right, even if it has fewer practice exams to study.

The Analogy: The Chef and the Salt
Think of training an AI like teaching a chef to make a soup.

Standard Training: The chef tastes the soup and gets criticized equally for adding too much water, using the wrong brand of pot, or forgetting the salt.
Reweighting: The teacher says, "The brand of pot doesn't matter. The amount of water is fine. But if you forget the salt, the whole dish is ruined."
The Outcome: The chef learns to be perfect with the salt (the critical clinical terms) very quickly, even if they haven't cooked 1,000 soups yet. They might still make small mistakes with the garnish, but the soup is safe and delicious.

What the Paper Found
The researchers tested this on eye disease reports. They found that by using this "Highlighter Method":

The AI learned 10 times faster. It could achieve the same high quality of reports with only 10% of the data that a standard AI needed.
It was better at spotting dangerous conditions (like Age-Related Macular Degeneration) because it was forced to pay extra attention to the words that describe those conditions.

In a Nutshell
Medical data is expensive and hard to get. This paper shows that you don't necessarily need more data to build a better AI; you just need to teach it to care more about the words that actually matter. By "weighting" the important words, we can train smarter, safer medical AI with less data.

1. Problem Statement

The generation of medical reports using Vision-Language Models (VLMs) is frequently constrained by the scarcity of high-quality, paired image-report datasets. This issue is particularly acute in specialized medical subdisciplines like ophthalmology, where data curation is expensive and time-consuming.

Standard training for these models typically employs a per-token cross-entropy loss, which treats every prediction error equally. The authors argue that this approach is suboptimal for medical contexts because it fails to distinguish between:

Inconsequential errors: e.g., predicting "OCT image" instead of "OCT scan."
Clinically critical errors: e.g., confusing "multiple drusen" with "no drusen," which fundamentally alters the diagnostic meaning.

The core problem is how to maximize the utility of limited data by ensuring the model prioritizes the learning of semantically and clinically significant tokens over generic ones.

2. Methodology

The proposed solution is a Token Reweighting strategy that modifies the standard training objective to penalize errors on clinically important terms more heavily.

A. Keyword Definition

The authors define three specific sets of keywords relevant to ophthalmology:

Quantitative Keywords ( $\kappa_Q$ ): 34 terms describing severity or extent (e.g., thick, several, increased).
Diagnostic Keywords ( $\kappa_D$ ): 22 terms corresponding to biomarkers and disease concepts (e.g., healthy, late, fluid, drusen).
Combined Set ( $\kappa_C$ ): The union of the above ( $n=56$ ).

B. Loss Function Modification

For a target token sequence $x = (x_1, \dots, x_T)$ , the authors identify the indices $I_\kappa(x)$ corresponding to the selected keywords using the model's tokenizer. They replace the standard cross-entropy loss with a normalized weighted objective:

$L_{tw}(x, \lambda) = -\frac{1}{\Lambda} \sum_{i=1}^{T} \lambda_i \log p_\theta(x_i | x_{<i})$

Where:

$\lambda_i = \gamma$ if token $i$ is in the keyword set $I_\kappa(x)$ , and $\lambda_i = 1$ otherwise.
$\gamma > 1$ is the upweighting factor for clinically important tokens.
$\Lambda = \sum_{i=1}^{T} \lambda_i$ is a normalization factor ensuring the overall loss scale remains consistent across reports with varying numbers of keywords.

C. Experimental Setup

Model Architecture: A VLM comprising a pretrained image encoder, a Llama3-3B language model, and a connecting projector layer.
Training Strategy: Fine-tuning via LoRA (Low-Rank Adaptation) on the projector and language model, with the image encoder frozen.
Data: Experiments conducted across varying dataset sizes (1%, 3%, 10%, 30%, and 100% of the available data).
Evaluation: Mean $F1_{macro}$ scores for Age-Related Macular Degeneration (AMD) staging and biomarker identification.

3. Key Contributions

Semantic-Aware Loss Function: Introduction of a simple yet effective token reweighting mechanism that aligns the training objective with clinical relevance rather than treating all tokens equally.
Systematic Analysis: A comprehensive evaluation of how different keyword sets (quantitative vs. diagnostic vs. combined) impact model performance across various data scales.
Data Efficiency Demonstration: Empirical proof that reweighting can achieve performance levels comparable to standard training with significantly less data.

4. Results

The experiments yielded several significant findings:

Consistent Performance Gains: Token weighting outperformed standard cross-entropy across nearly all data regimes and metrics.
Superior Sample Efficiency:
- Models trained with a weighted loss on 10% of the dataset outperformed unweighted models trained on the full (100%) dataset for AMD staging.
- In many cases, the benefit of using a weighted loss was equivalent to or greater than tripling the dataset size.
Keyword Set Effectiveness: All keyword sets improved performance, but the Diagnostic Keyword Set ( $\kappa_D$ ) yielded the most significant gains. This highlights that prioritizing biomarkers and disease-specific concepts is the most critical factor for report quality.
Quantitative Improvements:
- At 10% data size, the weighted model achieved an AMD F1 score of 0.490 compared to 0.422 for the standard model.
- At 100% data size, the weighted model reached 0.573 vs. 0.481 for the standard model.

5. Significance

This work demonstrates that data efficiency in medical AI can be significantly enhanced not just by collecting more data, but by optimizing how existing data is learned.

Practical Impact: For medical subfields with limited annotated data (like ophthalmology), this method allows for the training of high-quality VLMs with up to 10 times less data than standard approaches.
Clinical Safety: By explicitly prioritizing the correct prediction of critical diagnostic terms, the method reduces the risk of generating reports with dangerous clinical ambiguities.
Generalizability: The approach offers a simple, plug-and-play tool for improving VLMs in any data-constrained medical domain where specific terminology carries high clinical weight.