DevAaron2404/gain-predictor

Gain predictor

Predicts from the source text alone how much a heavy LLM improves a summary over a light one, BERTScore(heavy) - BERTScore(light) for Qwen3-4B against Qwen3-0.6B, before either model runs. A router uses it to decide when the heavy model is worth its extra energy.

Non-commercial use only (CC BY-NC-SA 4.0), because the DialogSum training data is.

Outputs

model.onnx takes the tokens of one text from tokenizer.json (cut at 256 tokens):

Output	Meaning
gain	Expected BERTScore gain of the heavy model
gain_low, gain_high	90% interval of the real gain
real_gain_probability	Chance that the gain is above the noise threshold (0.055)
familiarity	Cosine similarity to the closest training text

The graph holds MiniLM, a ridge regression for the gain, a second one for the size of its error, and normalised conformal prediction for the interval and probability. Do not trust the gain when familiarity is below the familiarity_threshold in the metadata, or when the text has more than max_words words.

Training and results

1,100 English texts: 500 CNN/DailyMail, and 150 each from XSum, DialogSum, BillSum and Reddit TL;DR. Each label averages 5 sampled summaries per model. Evaluated with 5-fold cross-validation repeated 10 times:

Model	RMSE	Skill over the mean [95% CI]
Always the average gain	0.058	0
Average gain per dataset	0.058	0.024 [-0.003, 0.051]
This model	0.057	0.036 [0.018, 0.054]

The 90% interval held 90% of real gains (87% to 93% per dataset).

Limitations

• The gain is hard to predict: the model beats the average only slightly, and mostly learns the kind of text.
• On a kind of text it never saw it has no skill, so check familiarity.
• English only, and MiniLM reads the first 256 tokens.