laya-rlcd-training

laya-rlcd-training

A training loop for Laya, the open-weights System One model.

Laya ships the building blocks but no training loop. Its README describes the method (RLCD, Gaussian noise on the logits, a strictly proper scoring rule, REINFORCE with a group-mean baseline), and the repository contains no backward call and no optimizer (checked 2026-09-20). This is that loop, written from the description and tested on a real corpus rather than a toy set.

Everything public about System One models right now rebuilds the inference path: score the allowed answers instead of generating text. That part is well covered. The part everyone leaves out is how you make the model good at your questions. That is what this is for.

What it does

Takes labelled text, trains the encoder to answer several typed questions in one forward pass, and reports accuracy per question against a held-out gold set.

Three question types, matching Laya's own:

Quick start

pip install torch transformers safetensors huggingface_hub datasets

python example/build_example.py --train 4000 --gold 400

python train.py \
    --data example/train.jsonl \
    --gold example/gold.jsonl \
    --questions example/questions.json \
    --batch 16 --stage 200 --save out/

The base model is fetched on first use. Set LAYA_SNAP to a local directory to use a copy you already have, or to continue from your own checkpoint.

The example is built from AG News, so it runs without any private data. It prints the majority-class baseline before training starts: the number your result has to beat before it means anything.

Your own data

One JSON object per line:

{"id": "a1", "text": "...", "topic": "business", "severity": 2, "violent": false}

A field missing from a line is skipped for that record. Questions do not all have to be labelled on the same texts, which lets you combine a large cheaply labelled set for one question with a small carefully labelled set for another.

Question definitions live in a JSON file; see example/questions.json and the docstring in data.py.

Trained model

The four-question model this loop produced on our news corpus is public: InfinimindCreations/laya-news-decisions.

Results on a news corpus

Six questions over 92k news articles, labelled by a larger model, evaluated against gold judgements from three independent LLM annotators (majority vote). Two runs, same loop, same questions, same gold set. The only change is the amount of training data.

A note on the gold set, because it decides what any of these numbers mean: it started at 176 judgements and the extremes were nearly empty: four cases on the top impact level, where a single article moved the score by 25 points. We widened it to 230 by sampling by level rather than by topic and having three LLM annotators judge each case independently, keeping only what at least two agreed on. Levels 3 and 4 went from 20 and 4 cases to 48 and 23. Every number below that says 230 uses the widened set.

All six improved, none traded against another. On the ordinal questions the within-one-level rate reached 0.97–0.99, i.e. the model rarely misses by more than one step.

Both runs: one A100, full fine-tune, batch 16, G=32, lr 1e-5 with cosine decay.

What a seventh question costs

We then added a title-level question to the same model, trained from the base weights so it would not inherit the collapse described below. Measured against a widened gold set (230 judgements for the ordinal questions, 900 for the new one):

There is no free lunch here and we are not going to pretend otherwise. One model is the specialist, the other the generalist; the seventh question cost about two points on average across the other six, worst on urgency. Both were trained with the same budget, so part of that is simply fewer steps per question.

One detail the accuracy column hides: on impact the seven-question model's rank correlation rose to 0.824 from 0.778. It orders the articles better while hitting the exact level about as often, which for that question is the property that matters.

Things that cost us time

• The default threshold of 0.5 is not a decision boundary. On one binary question the raw probability separated the classes perfectly (AUC 1.000) while a 0.5 cut got 15 of 20 right. Fit the operating threshold on real data.
• Narrow questions beat clever ones. "Does this describe violence?" scored F1 0.791. The same question with exclusion clauses scored 0.692, and splitting it into three sub-questions combined in code scored 0.762. This is a 421M encoder, and negations in the instruction hurt.
• Your calibration set needs the real distribution. A set built only from interesting cases reached AUC 0.860 and then failed in the field: a gardening tip scored 0.906. The base rate in the actual stream was 1 in 40.
• Check the majority-class baseline before you believe an accuracy. One of our questions sat at exactly the majority share for fifteen consecutive evaluations: the model had learned to always answer the most common label. With 83/10/7 class shares, always guessing the majority scores 0.83. The same checkpoint read as 0.839 against a held-out set with the natural distribution and 0.333 against a balanced one. Same model, same question, two numbers.
• A collapsed question does not come back by rebalancing the data. When we retrained that question on a balanced set, nothing moved, not in fifteen evaluations. Probing the raw distribution showed why: the model was not hesitating between classes, it answered the majority class with probability 1.0000 on every single item. At saturation the gradient is gone, so there is nothing left for the training signal to pull on. Worse, the base model had been better: it scored 0.4556 on the balanced set before any training, against 0.3333 after. Training on a skewed set destroyed an ability the model came with. The fix is to learn that question again from the base weights on a balanced set, not to fine-tune the collapsed one. Balance the classes the first time; it is much cheaper than noticing later.
• The student can beat the teacher. On one question the trained model scored 0.665 against the labelling model's own 0.631. Where teacher errors scatter randomly (99.4% within one level), training averages them out.
• Keep the best checkpoint, not the last one. REINFORCE oscillates; adjacent evaluations differed by up to 6 points, and one run's final checkpoint landed 4.6 points below its best. But if you pick "best" on the same gold set you report on, the reported number is the top of that oscillation, not a fair estimate. Report the last checkpoint next to it, or hold out a second set for selection.

Resuming

A long run should survive an interruption. --save writes two different files, and conflating them costs you a run:

• checkpoint.pt: the latest state, written at every evaluation, carrying the weights, the AdamW moments and the scheduler position. This is what --resume out/checkpoint.pt reads, so an interruption costs one stage at most.
• model-best.pt: the best state. This is the one you deploy.

Writing the resume file only on improvement looks harmless and is not: a run whose best score falls early would resume hours behind where it stopped.

Weights alone are not enough: without the optimizer moments and the scheduler position the learning rate jumps back to its warmup value and shakes the weights loose. Measured on a 600-step schedule, continuing correctly resumed at lr 4.12e-06, a weights-only warm start would have restarted at 8.82e-06.

Set LAYA_PUSH=org/repo to also push the best checkpoint to the Hub during the run, throttled by --push-every (uploads run on the accelerator's clock).

Memory

A full fine-tune of the 322M-parameter multilingual variant needs roughly 5.2 GB for parameters, gradients and the two AdamW moments, before activations. On a 12 GB consumer card that OOMs at batch 4; use --freeze 12 to freeze the lower encoder layers. On an 80 GB accelerator, leave --freeze at 0. In our runs the unfrozen model passed the frozen one's 5,000-article result after 1,600.

License

Apache 2.0, following Laya itself. The example data comes from AG News.