Probability-Entropy Calibration: An Elastic Indicator for Adaptive Fine-tuning

In plain words

Fine-tuning nudges a language model toward the right answer one token at a time. Many recent methods give each token a weight based on one number: how likely the model already thought the correct token was, or how unsure the model was about the next word. The paper argues that either number alone misleads. A filler word can look important only because the model is unsure, and a word with several acceptable synonyms can look wrong when it is not.

RankTuner combines both signals by comparing two ranks: where the correct token sits in the model's ordered list of candidates, and where a typical guess would land given how spread out the model's predictions are. Tokens the model is behind on, relative to how hard that position is, get a larger weight. Tokens that are uncertain anyway get a smaller one. The authors apply it to math reasoning and code. Their math gains are clear on most benchmarks, but the code results are mixed, and every number on this page is the authors' own.

TL;DR

RankTuner is a token-weighted cross-entropy loss for supervised fine-tuning. Each token's loss is multiplied by a Relative Scale built from the ground-truth probability and a rank-based bound on the entropy of the output distribution §4.5 alg. 1. The authors report better math Pass@1 than SFT, DFT, TALR, OverTone and EAFT on two backbones, best or near-best on most but not all benchmark and metric pairs tab. 2. On Qwen2.5-Coder-3B it does not beat the original model on HumanEval tab. 13, and the RL extension is untested §4.5 app. B.7.

Problem

Token-level reweighting in SFT usually uses one signal. Probability-dominant weights use the ground-truth probability, and entropy-dominant weights use the predictive entropy §1 §3.1. The authors argue that each misreads some tokens. Entropy-dominant weights up-weight high-uncertainty filler such as "umm", and probability-dominant weights over-penalise interchangeable words such as "essentially" §3.2 fig. 1. In their noise-insertion test, both baselines surface injected noise much more often than RankTuner does tab. 1.

Probability-dominant and entropy-dominant weights are also hard to compare directly, because probability and entropy live on different scales §4. The paper's answer is to express both through rank quantities, which share a unit §4.1 §4.3.

Method

Setting. A weighted NLL objective is written as the sum over tokens of w_t times the negative log-probability of the target token §2.1. SFT uses w_t = 1. RankTuner keeps the same loss and multiplies w_t by a per-token scale S_t eq. 11 alg. 1. For math, the paper sets w_t = p_t as the starting weight §4.5 app. B.6. For general and code tasks it uses w_t = 1 app. B.6 app. C.8.

Rank and expected rank. R_t is the rank of the ground-truth token when the vocabulary is sorted by probability §4.1. The expected rank E[R_t] is the expected number of guesses when guessing from the model's distribution, the sum over ranks i of i times the i-th largest probability eq. 1 §2.3.

Relative Rank Indicator. The indicator is I_t = g(f(R_t) - f(E[R_t])), with f(x) = 1 / log2(x + 1) and g(x) = 2^x eq. 2 §4.1. It is 1 when R_t equals E[R_t], falls as R_t grows, and rises with E[R_t] §4.1 fig. 2. The paper says the particular f and g are chosen for stability and a closed form, not as an optimum, and that other monotone pairs give similar results §4.1 app. C.5.

Bounds that replace probability and entropy. Two bounds link rank quantities to p_t and H_t. The first is R_t <= 1 / p_t eq. 4 app. A.1. The second is E[R_t] >= s(H_t), where s(H_t) = (1/4) * 2^H_t + 1 if H_t >= 2, and s(H_t) = 2 - p_max,t otherwise eq. 5 app. A.2. Entropy here is in bits (log base 2) app. A.2 app. B.1. The paper shows empirically that both bounds are tight on Qwen3-8B chain-of-thought tokens §4.3 fig. 2 fig. 3 app. B.5.

Power-law form. Applying the Cauchy mean value theorem to f and log2 gives f(R_t) - f(E[R_t]) = -K(xi_t) * (log2 R_t - log2 E[R_t]) for some xi_t between R_t and E[R_t] eq. 6 §4.4 app. A.3. Here K(xi) = xi / ((xi + 1) * [log2(xi + 1)]^2), which is positive eq. 6. Therefore I_t = (E[R_t] / R_t)^K(xi_t) eq. 7. The paper says K is about 0.5 for typical reasoning tokens, where R and E[R] are small, and that K falls for high-uncertainty positions §4.4 app. A.4.

Surrogates and the Relative Scale. Replacing R_t by 1/p_t and E[R_t] by s(H_t) gives the surrogate competence (p_t * s(H_t))^K(xi_t) eq. 8. The indicator is at least this surrogate eq. 9 app. A.5. The training weight is the inverse of the indicator, the Relative Scale S_t = I_t^-1 eq. 10 §4.5:

• S_t ≈ (p_t * s(H_t))^(-K(xi_t)), with xi_t = max(R_t, s(H_t)) eq. 10 app. B.1;
• K(xi_t) is simplified to [log2(xi_t + 1)]^-2 by dropping the factor xi / (xi + 1) "for simplicity and training stability" §4.5 app. B.1.

The modified token weight is w~_t = w_t * S_t eq. 11 §4.5.

Training step. Algorithm 1 is a standard SFT loop. For each batch, it runs the model, computes p_t, R_t, H_t, s(H_t), xi_t, K(xi_t), S_t and w~_t for each token, and takes gradient steps on the mean of w~_t times the token loss alg. 1 app. B.1. The algorithm does not define the token loss, so it is taken as the NLL from §2.1. The paper says the extra cost is O(|V|) per token, the same order as the cross-entropy itself, with O(1) extra scalars per token app. B.2.

# RankTuner-weighted SFT loss for one batch  [alg. 1] [app. B.1]
# logits: (N, V) at N supervised positions; targets: (N,)
# base_weights: p_t for math, ones for code  [§4.5] [app. C.8]
# Output: scalar loss

def ranktuner_loss(logits, targets, base_weights):
    # Token NLL l_t = -log p_t  [§2.1]
    nll = cross_entropy(logits, targets)                 # (N,)
    # Ground-truth probability p_t  [§2.1]
    gt_prob = exp(-nll)                                  # (N,)
    # Rank R_t: count of logits >= the target's logit, target included [§4.1]
    gt_rank = count(logits >= logits[targets])           # (N,)
    # Output distribution over the full vocabulary  [§2.2]
    probs = softmax(logits)                              # (N, V)
    # Top probability p_max,t  [§4.5] [eq. 5]
    max_prob = max(probs, axis=-1)                       # (N,)
    # Entropy H_t in bits over the full vocabulary  [app. B.1]
    entropy = -sum(probs * log2(probs), axis=-1)         # (N,)
    # Lower bound s(H_t) on the expected rank E[R_t]  [eq. 5]
    bound = where(entropy >= 2, 0.25 * 2 ** entropy + 1, 2 - max_prob)
    # Clamped argument xi_t = max(R_t, s(H_t))  [app. B.1]
    xi = maximum(gt_rank, bound)                         # (N,)
    # Exponent K(xi_t) = [log2(xi_t + 1)]^-2, simplified form  [§4.5]
    k_xi = 1 / log2(xi + 1) ** 2                         # (N,)
    # Relative Scale S_t = (p_t * s(H_t))^-K(xi_t)  [eq. 10]
    scale = (gt_prob * bound) ** (-k_xi)                 # (N,)
    # Token weight w~_t = w_t * S_t  [eq. 11]
    weights = base_weights * scale                       # (N,)
    # Weighted token loss, averaged over the N tokens  [app. B.1]
    return mean(weights * nll)

Pseudocode notes. The loop runs over batches alg. 1. The code treats the scale as a constant with respect to the parameters; see Paper vs code. The pseudocode above follows the paper, not the code's exact arithmetic.

Hyperparameters

Results

Author-reported. Setup: NuminaMath-CoT, first 10k training instances; evaluated on MATH-OAI (called Math500 in §5.1), Minerva Math, OlympiadBench, AIME24 and AMC23 §5.1 tab. 2. Values are percentages tab. 2. The paper reports no seeds, repeats or variance §5 app. C.

Main math results tab. 2:

Source: tab. 2. The paper's Delta rows compare against the Original model and against the best non-RankTuner baseline in each block tab. 2.

Points the table shows and the text does not stress:

• On Qwen2.5-Math-7B, MATH-OAI Pass@1 is 68.60 for RankTuner and 69.15 for DFT, so DFT is higher there tab. 2.
• On Qwen2.5-Math-7B, Minerva Pass@1 is 33.30 for RankTuner and 35.68 for TALR tab. 2.
• The AIME24 Pass@16 of RankTuner equals the Original model on both backbones (23.33 and 26.67) tab. 2 §5.2.
• The AIME24 and AMC23 test sets have 30 and 40 problems in the released data, so single-point gaps there are a few problems [raw data: math_evaluation/data/aime24/test.jsonl, math_evaluation/data/amc23/test.jsonl].

Diagnostic on noise tab. 1 (the paper marks these as lower-is-better; the token columns are precision and recall of the top 10% of tokens, and the sequence column is the hit rate at 10% app. B.3):

Out-of-distribution tab. 3. Qwen2.5-Math-7B, Pass@1, temperature 0.8, 16 samples, max 3072 tokens §5.3:

Other backbones tab. 9. The paper reports Qwen2.5-Math-1.5B, Qwen3-4B and Llama-3.1-8B with the same baselines. RankTuner is not best everywhere. For example, on Qwen2.5-Math-1.5B MATH-OAI Pass@1 it reports 62.00 against 62.94 for TALR tab. 9. The paper notes a few small regressions app. C.4.

Code generation tab. 13. HumanEval and HumanEval+, Pass@1 and Pass@10, with general-task weights w_t = 1 app. C.8:

On the 3B model, the original model is the best row on every column tab. 13. The paper reads this as limited capacity app. C.8.

Ablations fig. 4 (values only in the figure; not transcribed here). On AIME24 Pass@k for Qwen2.5-Math-7B, removing the probability term keeps Pass@1 about the same but weakens gains at larger k, and removing the entropy term degrades all k. The paper says both terms matter §5.4 fig. 4. Its dynamic K(xi_t) beats fixed exponents K = 0.5, 0.9 and 1.0 on five math benchmarks; the average Pass@1 is 37.28 for the dynamic version, against 34.66, 27.40 and 23.67 tab. 12 app. C.7.

Inference-time entropy §5.5 fig. 4. On AIME24 with Qwen2.5-Math-7B, the paper reports OverTone with the highest inference entropy, DFT with the largest collapse, and RankTuner closest to the original model. These are the paper's qualitative readings of its figure.

Paper vs code

Official repository: LvAoAo/Ranktuner_VERL, linked from the abstract. The owner name is not one of the paper's authors. The card uses it on the basis of that link. The clone is at commit 111c55de1ed5cc9a3bedee74f383eddc2a651ffe. It contains a vendored copy of verl (Apache-2.0) and math_evaluation/ (MIT); the root has no LICENSE file.

The loss is in verl/verl/trainer/compute_loss.py, and the trainer selects it when trainer.loss_type is ranktuner fsdp_general_trainer.py.

Matches the paper:

• Bounds on s(H_t): threshold of 2 bits, 0.25 * 2^H + 1 above it, and 1 + (1 - p_max) below it compute_loss.py eq. 5.
• xi = max(R, s(H)) and K = [log2(xi + 1)]^-2 compute_loss.py §4.5.
• Scale (p * s)^-K multiplied by p_t compute_loss.py eq. 10.
• Learning rate 5e-5, batch 256, max length 2048 and one epoch in the example command README.md.

Differs from the paper:

1. Training data size. The paper uses the first 10k NuminaMath-CoT examples §5.1. The README's data step passes --train_end 100000 README.md, and the default of 0 means no limit numina_cot.py. The evaluated checkpoint is global_step_390 README.md. At batch 256, 390 steps is about 99.8k sequences, so the example run is about 100k examples. The paper's 10k is not confirmed by this code.
2. Base weight is fixed to p_t. compute_loss.py multiplies every token by gt_prob compute_loss.py, with no option for w_t = 1. The paper uses w_t = 1 for code app. C.8. The repository has no switch for it, so this loss alone does not reproduce the code-generation setting.
3. The scale is a constant for the gradient. The scale is computed under torch.no_grad() compute_loss.py compute_loss.py. Gradients flow only through the per-token cross-entropy. The paper does not state whether S_t is differentiated.
4. Normalisation. The paper writes the loss as 1/T times the sum over tokens alg. 1. The code computes a mean over valid tokens per micro-batch compute_loss.py and then divides by the number of micro-batches and multiplies by the data-parallel size fsdp_general_trainer.py fsdp_general_trainer.py. The num_items_in_batch argument, which would give a global token mean, is commented out fsdp_general_trainer.py. Micro-batches with different token counts are weighted differently from the paper's formula.
5. Floors not in the paper. The product p_t * s(H_t) is floored at 1e-6 compute_loss.py. xi is floored at 1 compute_loss.py; this has no effect, because both R and s are at least 1. Inside the entropy log, 1.19e-7 (fp32 epsilon) is added to the log-probabilities compute_loss.py.
6. Entropy base. The main text does not give a log base §2.2. The appendix uses bits app. B.1. The code computes entropy in nats and divides by ln 2 compute_loss.py, so it matches the appendix.
7. Ties in rank. The code counts all tokens whose logit is at least the target's logit, so ties count against the target compute_loss.py. The paper does not say how ties are ranked.
8. Name collision. The variable expected_rank in the code holds s(H_t), the lower bound, not E[R_t] compute_loss.py. The code never computes E[R_t].
9. Optimizer defaults. The paper does not state betas, weight decay or gradient clipping. The code defaults are betas (0.9, 0.95), weight decay 0.01 and clip 1.0 sft_trainer.yaml. These are defaults, not reported settings.
10. Memory. The paper says memory equals SFT app. B.2. The code converts logits to float32 compute_loss.py and creates several tokens-by-vocabulary temporaries for the rank and entropy compute_loss.py compute_loss.py compute_loss.py. This is more than the paper's O(1) per-token extra.
11. Baselines. The paper says its baselines are re-implementations in its unified weighted-NLL framework, not the original code. OverTone is approximated as a discrete reweighting app. C.2. Compare baseline numbers with that in mind.

No further discrepancies were found in the code paths for the Relative Scale.

Pitfalls

• Entropy must be in bits. Eq. 5 uses 2^H, so H has to be in log base 2 eq. 5 app. A.2. Section 2.2 gives no base, and the appendix and code use bits app. B.1 compute_loss.py.
• K(1) differs between analysis and implementation. The main analysis says K is about 0.5 for typical tokens §4.4 app. A.4. The implemented K = [log2(xi + 1)]^-2 is 1.0 at xi = 1 eq. 10 alg. 1. Anyone implementing the Cauchy-based K from Section 4.4 gets different weights from the paper's training code.
• The xi choice and the dropped factor. The paper compares alternative means for xi only app. C.6. Dropping the factor xi / (xi + 1) is stated as a choice for simplicity and stability §4.5; no ablation of it is reported.
• Numerical floors. The 1e-6 floor on p_t * s(H_t) and the eps in the entropy are not in the paper compute_loss.py compute_loss.py. For tokens with very small p_t, the floor sets the weight.
• Large weights for rare tokens. The scale (p s)^-K grows as p_t falls. The code caps it by flooring p_t * s at 1e-6 compute_loss.py. The paper states no cap §4.5.
• Normalisation across micro-batches. The paper's 1/T is not tied to a micro-batch. The code averages per micro-batch alg. 1 fsdp_general_trainer.py. The paper does not say how it normalised across micro-batches §5.1.
• The RL extension is not validated. The paper describes an illustrative form for PPO and GRPO token ratios but leaves it untested §4.5 app. B.7.
• Math-only evidence for the strongest claims. The OOD results are on two benchmarks with one backbone tab. 3 §5.3. Code results are mixed tab. 13.
• No variance reported. Runs, seeds and confidence intervals are not stated §5 app. C. AIME24 and AMC23 differences are a few problems.
• Training data size is unclear. The paper says 10k examples; the example command and checkpoint imply about 100k §5.1 README.md README.md. Confirm which one a reproduction used.
• Code-task weights. The paper's code results use w_t = 1 app. C.8. This repository's only loss uses p_t compute_loss.py. Set the base weight yourself when reproducing code results.
• Stop-gradient choice. The code detaches the scale compute_loss.py. Whether the paper intends this is not stated.
• Tie handling. Equal logits count against the target compute_loss.py. This affects R_t and therefore xi.
• Memory. Plan for several tokens-by-vocabulary float32 temporaries in the reference implementation compute_loss.py compute_loss.py.

Lineage

• Builds on: the guessing-problem view of expected number of guesses §2.3, and the rank-based discounting f in the Relative Rank Indicator, which the paper attributes to ranking metrics §4.1. The Cauchy mean value theorem step is standard analysis §4.4 app. A.3.
• Compared against: OverTone, DFT, TALR and EAFT, all reimplemented in the paper's framework app. C.2. EAFT is arXiv:2601.02151, as cited in the paper's reference list. Also Alpha Power and Entropy Regularisation loss shapes §5.4.
• Related RL methods: PPO and GRPO are named as the RL-style setting the paper does not validate §4.5.
• Superseded by: not stated in this card's sources.