feat(common): parallel option layout for order-invariant decisions

[sharath-sms]

feat(common): parallel option layout for order-invariant decisions

What

An opt-in parallel option layout that makes a checkpoint's decisions invariant to option order
by construction, plus the training switch to produce such a checkpoint.

• Every option starts at the same position id. The instruction's closing [SEP] and the state
  continue after the longest option.
• An option cannot attend to another option's tokens. Everything else (instruction, state, [CLS])
  attends as before.
• ModernBERT's local layers keep their sliding window, but it is measured in position ids
  instead of sequence index, so every slot sees the same neighbourhood.
• The decision head already has no positional encoding, so with order-blind marker embeddings the
  logits simply permute with the options. The act head reads [CLS] and order-free statistics of
  the probabilities, so it is unchanged under reordering too.

The layout is a property of the weights, so it lives in the checkpoint config:
"option_layout": "parallel" in rl_agent_config.json. Every published checkpoint stays on the
default "sequential" layout, and nothing changes for them. Producing a parallel checkpoint:

from laya.train import TrainConfig, finetune

finetune("train.jsonl", "./laya_base", "./laya_parallel", TrainConfig(option_layout="parallel"))

Why

Where an option is listed changes the answer. The repo already measures this from several angles:

• tests/test_option_order.py: five options with identical text spread by 3.59 logits on the
  English checkpoint, by position alone.
• BENCHMARKS.md: option-order flip rates of 28% (CS) and 10.6% (voice).
• laya-multilingual never selects the first-listed score option — in English too (0/290), while laya does #131: laya-multilingual never selects the first-listed score option.

The current tools reduce the effect without removing it. shuffle_options teaches the model to
ignore order approximately, and option_order rotation averaging (README) costs k rows per question
and cut flips from 16.3% to 6.1%. A structural fix removes it: with the parallel layout, reordering
can only reorder the logits.

The technique is the one Parallel Context Windows (Ratner et al., ACL 2023)
uses for long contexts (shared position ids plus a block-diagonal mask), applied here to the option
block of a bidirectional encoder. PINE (Wang et al., 2024) removes
position bias between documents in decoder LMs along similar lines.

Results

Base checkpoint convaiinnovations/laya (English). All three arms were fine-tuned with
laya.train.finetune on typed-decisions train, using the notebook recipe (4 epochs, effective
batch 64, RLCD, seed 0) on one Colab T4. They were evaluated on typed-decisions test
(400 cases, 2,000 decisions). Each question was asked twice: in its listed order, and under a random
non-identity option_order.

arm	accuracy (listed order)	accuracy (shuffled)	ECE	same answer when shuffled	max prob. change
sequential (current)	0.779	0.767	0.153	91.7%	0.267
sequential + shuffle_options (all types)	0.757	0.750	0.136	95.0%	0.257
parallel (this PR)	0.764	0.764	0.143	100%	0.003 *

* fp16 autocast rounding. In fp32 the change is ≤1e-4, which is the 4-decimal rounding of the returned
probabilities.

Per question type, accuracy in listed order (shuffled order in brackets):

arm	choice	score	noul
sequential	0.757 (0.750)	0.735 (0.715)	0.860 (0.852)
sequential + shuffle_options	0.737 (0.732)	0.704 (0.694)	0.848 (0.843)
parallel	0.735	0.726	0.843

How to read it:

• Exact invariance. The parallel arm gives the same answer under reordering for 100% of
  decisions. The sequential arm changes 8.3% of them (12.6% on score), and shuffle_options
  reduces that to 5.0%.
• No measurable accuracy cost under arbitrary order. The sequential arm's 0.779 is partly a
  learned prior on the dataset's listed order: shuffle the options and it drops to 0.767. The
  parallel arm is 0.764 either way.
• Compared with shuffle_options, today's mitigation: the parallel arm has higher accuracy and
  full order invariance. shuffle_options keeps the best ECE.
• Where it trails. The parallel arm is slightly behind the sequential arm on choice (0.735 vs
  0.750 shuffled).
• Variance. This is one seed per arm. The same sequential recipe gave 0.7715 on Apple MPS and
  0.779 on the T4, so differences under about 0.01 should not be read as real.

Zero-shot (no fine-tuning) is not a substitute. Switching the shipped weights to the parallel
layout is already exactly invariant: on MASSIVE (20 options, research/eval/metamorphic.py) the
max probability drift goes from 0.999 to 7.8e-6. But MASSIVE accuracy falls from 0.795 to 0.54,
because the weights were trained on the other layout. A checkpoint has to be trained on it, which is
why the layout lives in the config and defaults to sequential.

Scope and limitations

• Requires transformers ≥ 5. The masks go through ModernBERT's per-layer-type attention_mask
  mapping, which transformers 4.x does not have. uses_parallel_layout raises a clear error on 4.x.
  The new test skips there.
• Eager backend only. The ONNX runtime refuses a parallel checkpoint at load. The TileLang fast
  path, compile=True and other non-eager backends refuse it at inference with a message pointing to
  set_backend('eager'). backend="auto" resolves to eager for such a checkpoint.
• Label text still matters. The layout removes the effect of position, not of option text.
  presentation_checks.py permutes score levels while the level N: prefix stays with its slot,
  so it mixes the two. There the slot-0 margins drop (score 0.66 → −0.13, choice 0.90 → 0.11 at
  zero-shot) but cannot reach 0.
• No new published checkpoint. This PR adds the capability and the evidence. Training and
  publishing a parallel checkpoint (English or multilingual, for laya-multilingual never selects the first-listed score option — in English too (0/290), while laya does #131) is the maintainer's call.

Changes

file	change
laya/common.py	parallel_layout, parallel_option_masks, uses_parallel_layout, OPTION_LAYOUTS; build_sequence(return_layout=); collate_items pads the layout and refuses a mixed batch; DecisionModel.forward(position_ids=, option_ids=)
laya/agent.py	reads option_layout at load; _encode_state carries the layout; _infer passes it and refuses non-eager forwards
laya/train.py	TrainConfig.option_layout, threaded through encode_item, train_model, calibration_records, finetune; written to the saved config only when set
laya/onnx_agent.py	refuses a non-sequential checkpoint at load
laya/backends/__init__.py	auto resolves to eager for a parallel checkpoint
research/eval/laya_eval.py	score_cases honours the layout, so presentation_checks / metamorphic measure parallel checkpoints correctly
tests/test_parallel_options.py	new assert-based suite, registered in scripts/test_suites.py
README.md	"Order-invariant checkpoints" under Option order

The default path is byte-for-byte unchanged. Without option_layout nothing new is computed, and
build_sequence keeps its return shape unless return_layout=True.

How it was verified

ruff check laya/ --select=E9,F63,F7,F82,F401,F811 --line-length=120   # All checks passed!
python -m compileall -q laya/ tests/
python tests/test_parallel_options.py   # all parallel-option tests passed (21 checks)
python tests/test_router.py             # all routing tests passed
python tests/test_criteria.py           # 198 passed, 0 failed
python tests/test_hooks.py              # all hook tests passed
python tests/test_hooks_api.py          # all hook API tests passed
python tests/test_option_order.py       # all option-order tests passed
python tests/test_train.py              # OK
python tests/test_backends.py           # 78 passed, 0 failed
python tests/test_packaging.py          # all packaging tests passed

Every tests/test_*.py was also run as a script on this branch and on main. The suites that fail
do so identically on both: they need CUDA/tilelang, the onnx extra, local checkpoints or network
services.

tests/test_parallel_options.py runs a tiny randomly initialised ModernBERT inside the real
DecisionModel, with no download. It checks:

• the layout itself (shared option positions, the state after the longest option, a max_len clamp
  that does not widen the last surviving option);
• that with no options marked, the hand-built masks reproduce the encoder's own padding and
  sliding-window masks;
• that across all 24 orders of four options, the parallel logits equal the canonical ones
  permuted (max diff < 1e-5), and the act head output is unchanged;
• that the sequential layout on the same weights is not invariant, so the check can fail;
• the collate_items and config plumbing.

Real-checkpoint numbers above come from research/eval/metamorphic.py, research/eval/presentation_checks.py
and an evaluation script over typed-decisions test, which is attached and can be added under
research/ if wanted.

Checklist

• Focused on one change
• Rebased on the latest main (8a6e132)
• Tests pass locally
• Docs updated (README Option order)
• Public API: additive only (new keyword arguments and config key); tests/test_hooks_api.py passes unchanged

Refs #131, #244, #602

[NandhaKishorM]

Merged, with the default left exactly where you put it.

This is a structural answer to a problem the repository has been measuring from four angles and mitigating approximately: tests/test_option_order.py spreading five identical options by 3.59 logits, the 28% and 10.6% flip rates in BENCHMARKS.md, #131, and rotation averaging costing k rows a question to get flips from 16.3% to 6.1%. 100% identical answers under reordering, with a max probability change of 0.003 attributable to fp16 rounding and 1e-4 in fp32, is a different kind of claim from any of those.

What made it mergeable is the discipline around the default, so I want to name the specific parts:

• The layout lives in the checkpoint config, not in a call argument, because it is a property of the weights. That is the right place and it is why nothing can turn this on by accident.
• build_sequence keeps its return shape unless return_layout=True, and with no option_layout nothing new is computed.
• The ONNX runtime refuses a parallel checkpoint at load, the non-eager backends refuse at inference with a message naming set_backend('eager'), and backend="auto" resolves to eager. Refusing beats silently computing the sequential answer on a parallel checkpoint.
• collate_items refusing a mixed batch rather than padding across layouts.
• Reporting the zero-shot result, MASSIVE 0.795 down to 0.54 while drift goes from 0.999 to 7.8e-6, instead of leaving it out. That number is the argument for why the config key exists, and omitting it would have made the PR look better and be worth less.
• Saying plainly that the parallel arm trails on choice (0.735 against 0.750 shuffled), that it is one seed per arm, and that shuffle_options keeps the best ECE.

I verified the default path myself rather than taking the byte-for-byte claim on trust, since this edits DecisionModel.forward, _encode_state and collate_items. On transformers 5.17.0, so tests/test_parallel_options.py actually ran rather than skipping: all 81 suites pass, and 74 real-weights decisions over choice, score and noul across English, Spanish, Hindi, French, Chinese and German, including budgeted calls, predict_long and usage accounting, come back with 0 flips and a max probability delta of 0.000000 against main.

Merge note. This conflicted with #933 in laya/train.py in two places and both sides were real: your parallel=parallel on calibration_records, and #933's calibration report and warnings. The merged call is calibration_records(..., parallel=parallel) feeding fit_temperature_map and then calibration_report, with the warning loop kept. The import block takes both OPTION_LAYOUTS and QTYPE_NAMES. Worth a glance at 6ab95ea and c062088 to confirm the threading is what you intended.

On the last point in your scope section: training and publishing a parallel checkpoint is mine to do, and the evidence here is what I would want before spending the compute. Leave #131 open against it.