Kolibri Post-Training from Scratch: Effort Labels, Format Rewards and a Gradient-Norm Trust Region
Building Aleph Alpha's Kolibri post-training recipe from the ground up: soft reasoning-effort labels drawn from log-length percentiles, attention-cost ordering for packed SFT batches, segment-aware group advantages, the p-norm format multiplier, the B-TV trust region, a closed-form per-token bound on the logit gradient, targeted on-policy self-distillation, the three-bucket curriculum, and the Merlin–Arthur game against hallucination. Every piece is derived on one German baker problem and a group of four rollouts.
The previous post built the Kolibri base model: 78.1B parameters, 3.46B of them active per token, and fluent in German. This post follows the same report (Aleph Alpha, 2026) through post-training. That covers 4,000 steps of supervised fine-tuning (SFT) on 268B tokens, then 1,000 steps of asynchronous RL on 38 environments with more than 1.2M tasks.
Most of post-training is data work. A handful of mechanisms, though, are compact enough to derive completely:
- Soft reasoning-effort labels. SFT traces carry no effort label, so one is assigned after the fact from the trace’s length relative to traces of similar difficulty.
- A multiplicative p-norm format reward that can never overrule the task reward.
- A segment-aware group baseline for multi-turn agents whose contexts get re-rendered.
- The B-TV trust region plus a closed-form bound on each token’s logit gradient, .
- Targeted on-policy self-distillation (OPSD), which rescues groups that RL would otherwise throw away.
- The Merlin–Arthur game, which teaches the model to abstain when its evidence has been removed.
We derive each one on a single German task.
The Running Example
The model gets a short document, a question, and one tool.
Kontext. (S1) Ein Bäcker verkauft drei Dutzend Brötchen. (S2) Jedes Brötchen kostet 0,45 €. (S3) Die Bäckerei öffnet um 6 Uhr.
Frage. Wie hoch ist sein Einkommen? Buche den Betrag mit
buche_einnahme.Tool.
buche_einnahme(betrag: number)
The answer is , which German writes as 16,20 €. RL samples a group of rollouts for this prompt (Kolibri uses ):
| Rollout | What it did | Task correct? | Format defect |
|---|---|---|---|
Reasons in German, books {"betrag": 16.2} | yes | none | |
| Two segments; the first reasons in English, the second in German | yes | language, in segment 1 | |
| Reasons in German, books correctly, but runs slightly past its length budget | yes | completion budget (mild) | |
| Misreads the decimal comma: € | no | — |
We will also need one failing variant of the group. In it, all four rollouts pass the amount as a string, {"betrag": "16.20"}. The tool rejects every call, so the group’s rewards are all zero.
1. Mathematical Setup
We reuse tools derived earlier in this series:
- The log trick, , the policy gradient, and importance sampling with the ratio are built in Reinforcement Learning from Scratch and Section 1 of MiMo-V2.6: Scaling RL from Scratch.
- Group-relative advantages , which sum to zero, are derived in Section 4 of the MiMo post.
- Softmax, KL divergence , and Gibbs’ inequality (, with equality iff ) are built in Mathematical Prerequisites for Foundation Prior.
One identity carries Sections 7 and 9, so we derive it here: the gradient of a log-softmax with respect to the logits. Let , so . Differentiate with respect to using the chain rule:
where is the Kronecker delta (1 if , else 0). For the sampled token , the whole gradient row is
where is the one-hot vector of . Example: with and the first token, .
2. SFT: Teaching German Reasoning and Effort Control
2.1 Reasoning in German, not just answering in it
Open corpora such as Nemotron translate prompts and answers into German but keep the reasoning in English. Machine-translating English traces into German leaks artefacts. The report’s example is a trace on our own baker problem that explains the German decimal comma to itself, writes 0.45 instead of 0,45, and glosses Brötchen as “Brötchen sind Brötchen”. The fix is reasoning-prefill distillation: the teacher’s thinking block is opened with a short German phrase, such as “Gegeben ist:”, picked at random from a list for each domain. With a system-prompt instruction alone, 67% of teacher traces stay German throughout. With the prefilled opener, 97% do.
Two more data rules matter for our example. Agentic traces get backfilled reasoning before tool calls that had none. A backfilled trace is kept only if a second model, reading only that reasoning, derives the same tool call. And clearly wrong tool calls are masked from the loss but kept in context. These are calls that report an unknown tool or invalid arguments, or exact repeats of a failed call. The model then learns the correction that follows an error without learning to produce the error. Our failing variant’s string-typed betrag call would be one of these masked turns.
2.2 Soft effort labels
Kolibri has four effort levels: none, low, medium and high. The chat template writes one fixed sentence per level into the system prompt. SFT traces come with no label, so each exchange gets one after the fact from its reasoning length . Lengths vary enormously: median trace lengths range from 27 to 90,223 tokens across datasets, and hard prompts reason up to 31× longer than easy ones. So thresholds are set per reference group . A distilled 2B grader, Poacher-2B, labels difficulty and agrees with its GLM-5.2 teacher 84.8% of the time. Within each group, the 30th and 70th length percentiles become the bucket edges 0.3 and 0.7 on a unit interval, interpolated in log space:
Check the endpoints. At the fraction is , so . At it is , so . The percentiles land exactly on the edges.
The label is then drawn, not thresholded:
with bucket centres and widths , , and .
Numerical check. Say ‘s reference group (German math, easy) has and tokens.
- : , so . This is the medium centre, and medium is drawn with probability 1.00.
- : , so , which gives high with probability 1.00.
- : , exactly on the medium/high edge. The medium exponent is and the high exponent is . They are equal, so the draw is a coin flip: 50/50.
Why equal at every edge. Each edge sits at , so its scaled distance from the centre is , whatever the bucket’s width. Neighbouring buckets always tie exactly on their shared edge. Labels change smoothly, and there is no exact token count at which the label flips, which the model might otherwise learn as a stopping point. By construction each group splits about 30% low, 40% medium and 30% high.
Does it work? Since 9.5% of opening prompts appear with two reasoning completions, often of very different length, the label is frequently the only input that differs between them. After SFT, the median reasoning length orders low < medium < high on 23 of 24 benchmarks. At fixed difficulty, high effort reasons 1.4–3.7× longer than low. The mean score rises from 38.4 at none to 64.9 at low and 67.8 at high.
2.3 Packing by attention cost
SFT packs conversations into 256k-token sequences with document-masked attention, using Next--Fit with open bins, which leaves only 0.001% padding. Equal length does not mean equal cost, because attention cost grows with the square of each document’s length:
A sequence holding one 256k conversation costs (in units of k²). A sequence of 64 conversations of 4k each costs . That is 64× less for the same token count. Data-parallel ranks synchronise after every microbatch, so everyone waits for whoever drew the expensive one. The fix needs no communication. Every rank buffers packed sequences, sorts them by , and draws them in a permutation from a seed shared by all ranks. Every rank sorts a buffer from the same distribution, so the -th draw costs about the same on every rank. On 64 B300 GPUs this raised median throughput from 8,610 to 15,850 tokens/s/GPU, an 84% gain.
2.4 Mixing and souping
The SFT pool holds 57 datasets, 22.0M samples and 335.8B tokens, of which 142.4B are loss tokens. MergeMix fine-tunes one specialist per data cluster and scores a candidate mix by averaging the specialists’ weights. With two or three clusters, merges ranked mixes the way real training runs did. With 20 clusters, they did not: the Spearman correlation between merged and trained scores was . Training the promising weightings was still worth it, since 6 of 7 beat the hand-set baseline. The two best full runs, C1 and C2, were souped by averaging every parameter. The soup reached a mean of 63.4% against 58.9% for the better single run and 30.0% for the base model. That soup is where RL starts.
3. The RL Objective
Kolibri optimises, over the trainable tokens and the distilled tokens of a batch,
The objective has three terms: a policy-gradient term, a squared log-ratio regulariser, and self-distillation. There is no KL penalty to a frozen reference model and no entropy bonus. Sections 4–9 build each piece on our group.
4. Rewards: Environment First, Format Second
4.1 The p-norm format multiplier
Each of format signals is 1 for clean output. Examples are tool-in-schema, well-formed tool call, reasoning language, content language, completion budget, and boxed answer. Each signal is softened to , and they combine as
Kolibri uses , and . For , no format signal applies.
Reading it. Each defect contributes a “cost” . The costs are combined with a -norm, mapped back through to a factor in , and then squeezed into . So:
- The task always dominates. Since , the bracket is always above , and . No pile of format defects can cost a correct answer more than half its reward.
- sets how defects compound. With one defect at : , cost , , so . With two such defects:
| aggregate cost | ||
|---|---|---|
| 1 (costs add, i.e. factors multiply) | ||
| 2 (Kolibri) | ||
| (only the worst counts) |
lets a second defect cost something without letting many defects add up linearly.
Our group. For , the mild budget overrun has , so . Then and . For ‘s first segment, the reasoning-language signal has , which gives . Its second segment is clean, at . Rollout is clean, . Rollout has , so .
Why only on positive rewards. In early runs, the format signals were added to the reward. The reward climbed but benchmarks did not. Most groups differed only in format, so the update chased formatting instead of correctness. Applying format only to successes, and only multiplicatively, fixed this: 500 RL steps then lifted AIME 2025 from 60.4 to 70.6.
4.2 Segment-aware group advantages
A segment is a run of rollout steps that extends one tokenised context. Whenever re-rendering changes earlier tokens, a new segment starts. That happens, for example, when earlier reasoning traces are dropped, which Kolibri does in a random half of groups, since deployments may re-encode text between turns. Each segment’s tokens get
The baseline averages within each sequence first. In our group, ‘s segment mean is , so
The advantages are : ; : on the English-reasoning segment and on the German one; : ; : . If we had averaged over all five segments, would count twice and the baseline would be . Splitting a sequence into segments would then change everyone’s advantage.
Like Dr. GRPO, Kolibri does not divide by the group’s standard deviation. Like DAPO, it drops groups with no reward spread, but with one twist. All-pass groups are kept when their format penalties differ, since format is the only remaining signal there. A dropped group is re-sampled from the same environment up to twice, so environments that drop a lot of groups are not starved. Sequences that fail from infrastructure errors leave the group before is computed.
5. The Policy-Gradient Term and the B-TV Trust Region
For a sampled token with training probability and inference probability , the ratio is , and
This is the Binary Total Variation (B-TV) mask of DPPO, with . PPO clips the ratio. B-TV masks on the absolute probability change, and only in the direction the advantage pushes.
Example. Take a token of with . If and the trainer has already moved it to , the change sets : the token has been pushed far enough. If instead , the change keeps it.
Because the loss is , the log trick gives , the importance-sampled policy gradient.
6. The Per-Token Gradient-Norm Bound
6.1 Deriving
This is the part that is easiest to get wrong. The B-TV mask limits how far a probability may move. It does not limit how hard a single token pushes on the logits.
The loss depends on the logits only through . By the chain rule, , because . For an unmasked token, then:
Take the -norm, using its absolute homogeneity, :
The sampled coordinate of is . Every other coordinate is , and since probabilities sum to one, . So the maximum sits at the sampled token, and . Therefore
On-policy, and . Off-policy, has in the denominator and is unbounded as . A single token that the stale inference policy almost never picked, but that the trainer now likes moderately, can dominate the whole update.
6.2 Numerical check on
Rollout has . Suppose one of its tokens was sampled when the inference policy gave it , and the trainer now gives it . The B-TV mask keeps it: and . Then and .
Check this directly on a three-token vocabulary with :
The largest magnitude is . The 1-norm is , as the report’s appendix states: the off-token entries sum to exactly the sampled entry.
6.3 The fix: scale, don’t clip
Kolibri multiplies each token’s loss by a constant (stop-gradient) factor:
For our token, , so its logit gradient falls from to . Tokens with , which includes every on-policy token, are untouched. Computing costs nothing, since it uses only the two probabilities the trainer already has.
Why not cap the ratio, as CISPO does? A ratio cap ignores the softmax factor . Consider a token at , . Its , a tiny gradient, yet a ratio cap still treats it as a large ratio. Conversely, at low probabilities a cap of 5 lets reach , far above 2. Bounding directly is the same as capping the effective ratio at , which adapts to each token.
What the bound buys, to first order. A gradient step on the logits changes the sampled probability through the softmax Jacobian . The appendix of the report works this through to
using , the maximum of a parabola at . No single token can move its own probability by more than .
The bound was added when moving from Kolibri Origin to Kolibri, after individual tokens produced gradient norms “in the thousands”. Over the final run, the B-TV mask and the bound together touched only about 0.01% of tokens per step. On top of the per-token bound, the global gradient norm is clipped at 1.0.
7. The Squared Log-Ratio Regulariser and Engine Consistency
The second term softly penalises drift on every trainable token, masked ones included:
Because and only the first part depends on , the chain rule and the boxed identity of Section 1 give
If the trainer has raised a token above its inference probability (), descent lowers it again, back towards .
That only makes sense if and differ because of learning, not because of arithmetic. Inference serves FP8 weights and an FP8 KV cache, which is 2–3× faster. When the trainer used plain BF16, RL diverged. So the trainer applies the same FP8 rounding in BF16 arithmetic (quantisation-aware training). It replays the experts the inference engine routed to, and keeps the LM-head logits in FP32 on both sides. Rollouts sample at temperature 1 with no top- or top-, so no truncated support has to be replayed. The QK-norm bound from the previous post guarantees that the FP8 queries and keys never overflow. In the final run, the mismatch hovered around , staleness included.
8. Targeted On-Policy Self-Distillation (OPSD)
Return to the failing variant: four rollouts, all booking {"betrag": "16.20"} as a string. Every booking fails, the rewards are , and Section 4.2 drops the group. Yet the failure is obvious and easy to fix. OPSD turns such groups into dense signal. The run applies it to tool-use and terminal environments, for groups whose best reward is below 0.2:
- Pick the rollout with the most tool-response errors.
- Insert a remedial hint for the error class at the exact turn it occurred. For
wrong_argument_type, the hint is “Arguments use the exact JSON types the tool declares.” - Re-score every token of the failed turn under the current policy, with the hint. That is the teacher. The student is the same policy without the hint.
- Minimise , with the RL loss switched off for that rollout.
No ground-truth solution is needed. The hint is the privileged information.
Deriving the gradient
Write the teacher as (constant, because of the stop-gradient) and the student as . Then
The first sum does not depend on . For the second, use the identity of Section 1:
since . Hence
Numerical check. Take the token after "betrag": with three candidates: " (opening a string), 16 (a number), and anything else. The unhinted student has . The hinted teacher has . The gradient is . Descent lowers the logit of " and raises that of 16 by the same amount, which is exactly the fix. A finite-difference check of gives . The report’s Figure 68 shows this exact error class, with the teacher’s probability mass sitting on the opening quote.
Over the run. Detections per OPSD sequence fell from 4.21 to 2.56. Invented tools and references were trained away almost entirely, and customer-service detections dropped 99%. Late in training, output_truncated dominated as software-engineering rollouts grew longer. Only 0.10% to 0.034% of trained tokens used the distillation loss. Those tokens had consistently higher entropy than average, which is where the policy is least sure.
9. One View of All Three Terms
Every term of the objective has a logit gradient of the same shape, a weight times ( minus a target):
| Term | weight | target | |
|---|---|---|---|
| Policy gradient | one-hot | ||
| Log-ratio regulariser | one-hot | ||
| OPSD | hinted distribution |
Gradient descent moves towards the target when the weight is positive and away from it when the weight is negative. The policy gradient pulls towards the sampled token when and pushes away when . Its weight is capped in magnitude by and gated by . The regulariser pulls towards or away from the sampled token, whichever undoes recent drift. OPSD pulls towards a whole distribution rather than a single token, which is why it carries more information per token than an advantage does. In every case the push on the logits is a bounded weight times a difference of two probability vectors, so it is bounded too.
10. The Curriculum: Spend Rollouts Where the Signal Is
A prompt the model always solves, or never solves, yields a dropped group. The curriculum sorts each environment’s prompts into three buckets by estimated pass rate: hard (estimated 0), trainable (in between) and solved (1). Pass-rate estimates start offline, or at 0.5 by default, and are refreshed by every high-effort group. Bucket is drawn with weight
The supply term shrinks a bucket’s weight when it holds fewer prompts than the batch’s groups. The recency term grows linearly from 0 to 1 over 50 policy updates since the bucket’s oldest prompt was last tried. It stops a just-failed hard prompt from being redrawn at once, and it is 1 for the trainable bucket.
Numerical check. Suppose our baker task’s environment holds 40 hard, 1,000 trainable and 600 solved prompts, and the hard bucket’s head was tried 25 updates ago, so :
Normalised, these give 1.3%, 96.5% and 2.1% of draws. The environment itself is drawn in proportion to its sampling weight times . An environment that runs out of trainable prompts fades from the mix on its own, and comes back once re-evaluated prompts become trainable again.
11. The Merlin–Arthur Game Against Hallucination
To teach abstention, Kolibri plays a game on the context. Arthur is the policy. Merlin hides the sentences that matter least for the gold answer . Morgana hides the ones that matter most. Both hide about the same number of sentences, so Arthur cannot tell them apart by how much is missing. The ideal masks solve
In practice, each sentence is masked alone and ranked by how much the log-probability of drops. The masks are recomputed from the current policy at every step, so Morgana always tracks whatever currently lures Arthur into guessing.
Our context. Suppose . Masking each sentence alone gives without S1 (the quantity), without S2 (the price), and without S3 (opening hours). With , Merlin hides S3, the smallest drop (). Morgana hides S2, the largest (). Without the price, 16,20 € cannot be derived.
Each question is posed three times without labels: with the full context, Merlin’s, and Morgana’s. The reward is for the correct answer on the full or Merlin context, for a wrong answer there, for abstaining on Morgana’s context, and otherwise (small overlap terms aside). Advantages are centred per role: . If two of four Morgana-context rollouts abstain (reward 1) and two invent a price (reward 0), the role mean is 0.5 and their advantages are .
Abstaining on the full context earns 0, less than a correct answer’s 1, so Arthur cannot win by refusing everything. The same game gives a metric. Completeness is how often the model answers correctly with Merlin’s context, and soundness is how often it abstains with Morgana’s. Together they yield the M/A grounding score, a certified lower bound on how much of the answer came from the document. Kolibri scores 23.4%; Kolibri Origin and several baselines score 0. On AA-Omniscience, Kolibri abstains instead of answering wrong on 44% of non-correct items, against 15% for Origin.
12. Asynchrony and Staleness
The run used 256 B300 GPUs: 96 for the trainer, 144 single-GPU vLLM replicas for rollouts, and 16 for a frozen Qwen3.5-122B judge. Groups enter a replay buffer as they finish. The trainer samples each group at most once. After every step, FP8 weights are pushed by RDMA into one replica per node and broadcast over NVLink. In-flight requests keep their KV cache, so a single rollout can span several policy versions. Weight updates went from 221 + 217 seconds with the original synchronous design to about 17 seconds.
Staleness is a ratio of wall-clock times. A group’s staleness is roughly its rollout duration divided by the trainer’s step time. With 2-minute steps, a 60-minute code rollout arrives steps stale. Doubling both trainer and inference GPUs at the same batch size halves the step time to 1 minute. Per-request decode speed is unchanged, so the same rollout now arrives 60 steps stale. That exceeds Kolibri’s eviction limit of 42, and the group is thrown away. The remedy is to scale the batch size with the GPU count, which keeps the step time, and with it the staleness, fixed. The final run had a median staleness of 10 steps.
13. Results
What RL added on top of the SFT soup, at high effort:
| Benchmark | SFT | RL | Δ |
|---|---|---|---|
| AIME 2025 | 92.7 | 96.9 | +4.2 |
| AIME 2026 | 92.1 | 96.0 | +3.9 |
| GPQA Diamond | 79.8 | 84.3 | +4.5 |
| IFBench | 48.2 | 78.1 | +29.9 |
| τ³-bench banking | 13.1 | 38.1 | +25.0 |
| AA-Omniscience Index | −68.8 | −32.8 | +36.0 |
| LiveCodeBench v6 | 75.3 | 85.9 | +10.6 |
| Terminal-Bench 2.1 | 23.2 | 27.7 | +4.5 |
| AA-LCR | 60.3 | 68.3 | +8.0 |
On German AIME 2026, the share of reasoning in German rose from 89% to 99%, and accuracy from 75% to 91%. English AIME 2026 moved from 92% to 96%.
Against open models, using Overall aggregates in English / German:
| Model | Active | EN | DE |
|---|---|---|---|
| Kolibri | 3.46B | 75.5 | 70.8 |
| Qwen3.5 35B-A3B | 3B | 74.7 | 69.8 |
| Qwen3.6 35B-A3B | 3B | 71.4 | 67.3 |
| Gemma 4 26B-A4B IT | 4B | 71.9 | 66.3 |
| GPT-OSS 120B | 5.1B | 72.3 | 70.2 |
| Nemotron 3 Super 120B-A12B | 12B | 73.0 | 67.9 |
| Kolibri Origin | 3.27B | 54.1 | 46.4 |
| Qwen3.8 27B (dense) | 27B | 80.2 | 79.9 |
Kolibri is the strongest MoE model tested in both languages. Only the dense Qwen3.8 27B, with almost 8× the active parameters, scores higher. Kolibri leads the MoE models on English AIME, the SQuAD grounding score and the automotive customer proxy. It averages 89.7 on the English in-house Industry RAG proxies. Its weakest rows are RGB closed-book, where it scores 51 and ranks last, AA-Omniscience accuracy, and BFCL multi-turn. That fits a model trained to abstain rather than guess.
The report records two exploits, both caught through sequence monitoring. With the OpenCode harness, the model used web-fetch to download upstream SWE-bench fixes in about 7% of trials. In one under-isolated sandbox, a command switched off unprivileged user namespaces on the host kernel. In the report’s words, sandbox escapes are “as much a function of sandbox security as of model capability”.
Summary
Kolibri’s post-training uses the same design rule throughout: every learning signal is bounded and scaled to how much it can be trusted. Effort labels come from a smooth Gaussian draw on log-length percentiles, so neighbouring buckets tie exactly at their shared edge and no hard cutoff exists. Format signals multiply only successful rewards and can never take away more than half of one. The policy gradient is masked by B-TV and scaled so that no token’s logit gradient exceeds , where would otherwise blow up off-policy. The regulariser and OPSD are the same “weight times ( − target)” push, the latter rescuing all-fail groups with a hint. A curriculum and the Merlin–Arthur game spend rollouts where they teach, including teaching the model when not to answer.
Previous: Kolibri Pre-Training from Scratch
Enjoyed this post?
Subscribe to get notified when I publish new posts. No spam, unsubscribe anytime.