DecisionFormer

Fine-tune any Hugging Face model into a calibrated, one-pass decision model, with full fine-tuning, LoRA or QLoRA.

Our Lumma-Fev models score the options you give them in a single forward pass, instead of generating a label for you to parse. Until now, the code that trains models like that lived in our own repositories. Today we are open-sourcing it as DecisionFormer.

Think of it as the Transformers library for decision models. You choose a backbone, a fine-tuning mode (full, LoRA or QLoRA), an attention mode (causal, bidirectional or prefix-bidirectional) and a head. You write them in one YAML file, run one command, and get a checkpoint you can load in three lines of Python or serve over HTTP.

The library is Apache 2.0 and lives at FrontiersMindAI/decisionformer. To give you something to train on, we are also releasing a dataset mix, deciformer-15k.

This post is the documentation: what a decision model is, how to run one, and how to train your own with every option the library offers. Each example below uses the real commands and config keys, and the last section covers a few things to know before you start.

State in, probabilities out

A decision model takes a state (text, a JSON object, or an array) and a dict of typed questions. For every question it returns a probability for each option, computed in one forward pass with no text generation. Nothing is sampled, so nothing needs parsing, and the same input always gives the same answer.

• choice: pick one of your criteria (up to 255 options)
• noul: a yes/no gate, answered as the probability of yes
• score: pick one of 2 to 10 ordered levels

The API shape is the /v1/systemone contract that TypeSafe Jev, vLLM and the Decision Index already use, so a model you train with DecisionFormer drops in wherever those clients are used.

Input, output and why a classifier is enough

What goes in

Every question is rendered into one prompt with the same layout: the state, the question, the options each tagged with a short code (A, B, C, and so on), and a closing instruction. A decoder gets this wrapped in its own chat template, with a system message that tells it to treat the state as data and reply with only a code. The position after the template’s assistant header is where the answer gets decided.

State:
The order arrived damaged yesterday. The customer has a receipt and asks for a replacement today.

Question:
Which team should handle this request?

Options:
A: returns: Refunds, replacements and damaged deliveries
B: billing: Payments, invoices and charges
C: technical: Product setup and faults

Reply with only the code of the best option.

A noul question uses two options, no/false as A and yes/true as B. A score question lists its levels in order. A state that is not a string is written as JSON, and an empty state is shown as (empty). Encoders such as ModernBERT do not use a chat template: the state and the question block go in as a sentence pair and the first token is read.

What comes out

At the decision position the head produces one logit for each option code. DecisionFormer divides them by a calibrated temperature and applies a softmax, so every question gets a probability for every option, and the probabilities sum to 1. That list is turned into the answer for the question’s type:

Why turning a generator into a classifier works

A language model that answers with a single letter has already made its decision before it writes anything. The decision is the probability distribution over the next token at the answer position, and generating text only samples one token from it. A decision model skips the sampling and reads the distribution directly.

• The answer lives in one place. With one-token option codes, the whole answer is a choice among a few tokens at one position, so one forward pass contains everything a generation loop would have produced.
• Only valid answers are possible. The head looks at the code tokens only, so the model cannot answer with a typo, an explanation or a code that does not exist. There is nothing to parse and no retry.
• It starts from what the model already knows. The code_readout head is initialized from the language-model head’s rows for the code tokens, so an untrained model already reads options the way the pretrained one does. (lm_verbalizer keeps those rows frozen, which gives a zero-shot baseline.) Fine-tuning then sharpens it with a cross-entropy loss over the options.
• You get probabilities, not a sampled guess. The same input always returns the same distribution, and the temperature fitted on held-out data aligns those probabilities with how often the model is right.

How confidence is computed

confidence is not the top probability. It measures how decisively the distribution points at one option, from 0 for a flat distribution to 1 for a certain one. With n options and probabilities p after the temperature:

p = softmax(z / T)     z: option logits, T: temperature fitted per question type

d3 (default)
  confidence = 1 − H(p) / ln n,     H(p) = −Σ pi ln pi

autojev, choice
  confidence = (pmax − 1/n) / (1 − 1/n)

autojev, score
  confidence = 1 − Σ pi·|i − best| / b,     b = (1/n) Σ |i − (n−1)/2|

Both are clamped to the range 0 to 1, and a question with a single option has confidence 1. The d3 form is one minus the normalized entropy. For p = (0.80, 0.15, 0.05) the entropy is 0.613 and ln 3 is 1.099, so the confidence is 0.44, even though the top probability is 0.80. The autojev choice form measures the top probability above chance, which gives 0.70 for the same distribution. For ordered levels it uses distance from the chosen level, so probability spread over neighbouring levels costs less than spread across the whole scale. Which form you get follows the prompt format of the model (data.prompt.style).

Install and decide

You need Python 3.10 or newer, PyTorch 2.6+ and transformers 5.20+.

pip install decisionformer
pip install "decisionformer[serve]"            # adds the API server

# or from source
git clone https://github.com/FrontiersMindAI/decisionformer.git
cd decisionformer && pip install -e ".[dev]"

Extras combine, for example pip install "decisionformer[serve,qlora]".

Decide in Python

Any checkpoint loads the same way, whether it is one of ours, one you trained, local or on the Hub.

from decisionformer import AutoDecisionModel

model = AutoDecisionModel.from_pretrained("vllm-sr/d3-edge")   # or your own fine-tune

response = model.decide(
    state="The order arrived damaged yesterday. The customer has a receipt and asks for a replacement today.",
    questions={
        "route": {
            "type": "choice",
            "instructions": "Which team should handle this request?",
            "criteria": {
                "returns": "Refunds, replacements and damaged deliveries",
                "billing": "Payments, invoices and charges",
                "technical": "Product setup and faults",
            },
        },
        "receipt": {"type": "noul", "instructions": "Does the customer have a receipt?"},
        "urgency": {
            "type": "score",
            "instructions": "How urgent is this request?",
            "criteria": ["Routine", "Soon", "Today"],
        },
    },
)

Output:

{
  "model": "vllm-sr/d3-edge",
  "answers": {
    "route": {
      "type": "choice",
      "choice": "returns",
      "probabilities": {"returns": 0.9927, "billing": 0.0001, "technical": 0.0073},
      "confidence": 0.9602
    },
    "receipt": {"type": "noul", "noul": 0.9995},
    "urgency": {
      "type": "score",
      "score": 0.5777,
      "probabilities": {"0": 0.5393, "1": 0.3438, "2": 0.1169},
      "confidence": 0.1343,
      "legend": {"0": "Routine", "1": "Soon", "2": "Today"}
    }
  },
  "usage": {"input_tokens": 345, "output_tokens": 0}
}

output_tokens is always 0: nothing is generated. score is the expected level index, so with three levels it runs from 0 to 2, and decide_many([...]) answers several requests in shared batches.

Serve it as an API

decisionformer serve vllm-sr/d3-edge --port 8000

curl localhost:8000/v1/systemone -H 'content-type: application/json' -d '{
  "state": "I was charged twice for my subscription this month.",
  "questions": {
    "angry": {"type": "noul", "instructions": "Is the customer angry?"}
  }
}'

The server returns the same JSON as decide, plus id, object and created. Over-long requests are refused with HTTP 422 and never silently truncated. Set --api-key-env MY_KEY_VAR to require Authorization: Bearer <key>.

One config, any fine-tuning mode

Every setting lives in a YAML file. Recipes inherit shared defaults with base:, and any value can be overridden on the command line with dotted keys. decisionformer config FILE prints the fully resolved configuration, and a bad value fails when the config loads, not an hour into training.

model:
  backbone: Qwen/Qwen3.5-0.8B        # any Hub id or local path
  attention: causal                  # causal | bidirectional | prefix_bidirectional
  head: {type: code_readout}
finetune:
  mode: lora                         # full | lora | qlora | head_only
  lora: {r: 16, lora_alpha: 32, target_modules: all-linear}
data:
  train: data/train.jsonl
  calibration: data/calibration.jsonl
  eval: {dev: data/dev.jsonl}
train:
  output_dir: runs/my-decider
  lr: 2.0e-4
  push_to_hub: my-org/my-decider     # optional: uploads best/ when training ends

decisionformer train configs/recipes/my_recipe.yaml
decisionformer train configs/recipes/my_recipe.yaml train.lr=1e-4 finetune.lora.r=32

A run writes runs/<name>/best and last: finished checkpoints in the standard Hugging Face layout, ready for from_pretrained, serve and push. It also writes metrics.jsonl, per-example predictions at every evaluation, resumable state and summary.json.

Full fine-tuning

Every weight trains. This is the right choice for small decoders and for encoders.

finetune:
  mode: full
  gradient_checkpointing: true
  freeze_embeddings: true
  freeze_layers: "0-7"           # full fine-tuning only
  layerwise_lr_decay: 0.9

LoRA

The lora block is passed to peft with peft’s own field names and is checked when the config loads.

finetune:
  mode: lora
  gradient_checkpointing: true
  lora:
    r: 16
    lora_alpha: 32
    lora_dropout: 0.05
    target_modules: all-linear   # list | regex | all-linear | attention | mlp | attention_mlp
    use_dora: false
    use_rslora: false
    init_lora_weights: true      # true | gaussian | pissa | olora | loftq | ...
    rank_pattern: {q_proj: 32}
  export: merged                 # merged (plain weights, no peft needed) | adapter

QLoRA

The base model is quantized to 4 or 8 bits and only the adapters train. It needs pip install "decisionformer[qlora]" and a CUDA GPU.

finetune:
  mode: qlora
  gradient_checkpointing: true
  quantization: {bits: 4, quant_type: nf4, double_quant: true, compute_dtype: bfloat16}
  lora: {r: 16, lora_alpha: 32, lora_dropout: 0.05, target_modules: all-linear}

Head only

mode: head_only trains just the decision head on a frozen backbone. With a lm_verbalizer head there is nothing to train at all, so the run is a zero-shot baseline that still gets calibrated.

In every mode the decision head trains fully, with its own learning rate (train.head_lr). The code readout starts from rows of the language-model head, so a small head learning rate such as 2e-5 works well.

Precision

train:
  mixed_precision: bf16      # bf16 (default) | fp16 (with loss scaling) | no (float32)
finetune:
  model_dtype: auto          # weights while training: float32 | bfloat16 | float16
  export_dtype: auto         # saved checkpoints: the base checkpoint's dtype by default
  quantization:
    compute_dtype: auto      # QLoRA compute: follows mixed_precision

With model_dtype: auto, LoRA keeps the frozen base in bf16, QLoRA in 4-bit, and full fine-tuning keeps float32 master weights while the compute runs in bf16, as PyTorch AMP does. GPUs without bf16 (Turing, V100) fall back to fp16 with loss scaling and log a warning. Exports use the base checkpoint’s dtype, so a fine-tuned model is the same size as the one it started from.

Which mode for which model

Approximate sizes per 80 GB GPU, from the repository’s documentation:

Training data

Training data is the request together with the gold answers, one JSON object per line.

{
  "id": "t-17",
  "family": "ticket-17",
  "state": {"ticket": 17, "text": "I was charged twice this month."},
  "questions": {
    "route": {"type": "choice", "criteria": {"billing": "Payments and invoices", "returns": null}},
    "angry": {"type": "noul", "instructions": "Is the customer angry?"}
  },
  "gold": {"route": "billing", "angry": false},
  "targets": {"route": {"billing": 0.9, "returns": 0.1}}
}

• gold may also be called expected or answers.
• targets (soft labels such as vote shares) and teachers (probabilities to distill from) are optional, per question.
• A record with one question and a label also works, and so do flat rows with one question each: state, question, question_type, options, option_descriptions, label and soft_target.
• family keeps related records together in splits, data checks and bootstrap intervals.
• Files can be .jsonl, .jsonl.gz, or a Hub dataset written as hf:<name>[:<config>]:<split>.

decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k
decisionformer data check dev=dev.jsonl calibration=calibration.jsonl train=train.jsonl

data split holds out 10% for calibration and 10% for dev by default, keeps whole families together and gives each task an equal share of every split. data check finds leakage between splits, and exits non-zero when it finds some, so you can run it in CI.

In flat rows, a yes/no question’s options are true and false, so labels such as yes and no are rejected when the file loads.

Calibration comes with training

The calibration file is never trained on. After training, DecisionFormer fits one temperature per question type on it and stores them in the checkpoint, which is why the probabilities you get back are calibrated. Evaluation reports accuracy, ECE, Brier score and NLL, raw and calibrated.

Causal, bidirectional or prefix-bidirectional

Decoders read left to right by default. For a decision, the whole prompt is known up front, so you can let tokens look in both directions.

Bidirectional mode uses transformers’ own switches (is_causal=False on the text decoder only), so SDPA, flash and sliding-window kernels all respect it, and a vision tower keeps its own attention.

What each mode lets a token see

Attention decides which other tokens a token can read while it builds its representation. Think of the prompt as three parts: the state, the question with its options, and the short answer slot at the end, where the code is read.

The same eight-token prompt under each attention mode. In every mode the answer is read from the last token (the outlined row).

• Causal. A token sees itself and what came before it. The state is encoded before the model has seen the question, so state tokens cannot know what you are asking, and only the last token sees everything. This is how decoders are pre-trained.
• Bidirectional. Every token sees every other token, in both directions. State tokens are encoded knowing the question and the options, and the answer token also sees the whole prompt. Recurrent and linear-attention layers, such as Qwen3.5’s Gated DeltaNet, stay causal.
• Prefix-bidirectional. The prompt (state, question and options) is bidirectional, and the answer slot stays causal: it sees the whole prompt, and the prompt does not see it. The mask is built explicitly, so it runs on sdpa or eager attention.

In all three the answer is read from the last token, so the output format does not change.

Which one to use

A short rule: start with causal, train a bidirectional version of the same recipe, and keep whichever scores better on your dev set. Bidirectional attention is what Decision 3.0, pplx-decider and GLiDE use, but a decoder that was pre-trained causal may need the MNTP adaptation below, or enough fine-tuning, to benefit.

Set it in the config:

model:
  backbone: Qwen/Qwen3.5-0.8B
  attention: bidirectional
  head: {type: code_readout}

Compare against a causal baseline before you adopt it. A decoder pre-trained as causal does not automatically work well when you open its attention. In our MNLI tests, SmolLM2-360M made bidirectional without adaptation stayed near chance where the causal model reached 0.80. Qwen3.5-0.8B with bidirectional attention and QLoRA did reach 0.84, so try it per model rather than assume. The numbers are in docs/benchmarks.md in the repository.

If you want bidirectional attention on a decoder that needs it, decisionformer adapt runs masked next-token prediction (as in LLM2Vec) first, and you then train from the adapted checkpoint. This stage is experimental: a short run did not help in our tests, so budget generously and use in-domain text.

decisionformer adapt configs/recipes/adapt_mntp.yaml
decisionformer train my_recipe.yaml model.backbone=runs/smollm2-mntp/adapted model.attention=bidirectional

Bring the model you already use

Backbones are the architectures built into Hugging Face transformers, loaded with its Auto classes and saved with save_pretrained. DecisionFormer has no model code of its own: it builds on the architectures transformers already ships. Decoders read the answer as a code token (A, B, …) at the end of the prompt, which works from step 0. Encoders score each option.

Hub models that need their own modeling code (trust_remote_code) are not supported, and encoder-decoder models without a standalone text encoder in transformers, such as BART, are refused with a clear error.

An encoder example

model:
  backbone: answerdotai/ModernBERT-base
  pooler: {type: cls}
  head: {type: fixed_softmax, labels: [entailment, neutral, contradiction]}
finetune:
  mode: full
  layerwise_lr_decay: 0.9
train:
  lr: 5.0e-5
  head_lr: 5.0e-4

How the answer is read

Use model.head for one task, or model.heads to give every task its own head. A record picks its head with the task field. A head can also set its own pooler: decision_token, last_token, cls, mean, attention or option_spans.

One model, several tasks

model:
  backbone: Qwen/Qwen3-1.7B
  heads:
    nli: {type: code_readout}
    sentiment: {type: code_readout}
    boolq: {type: binary}
finetune:
  mode: qlora
  quantization: {bits: 4, quant_type: nf4, double_quant: true, compute_dtype: bfloat16}
  lora: {r: 16, lora_alpha: 32, target_modules: all-linear}
data:
  train: [data/mnli/train.jsonl, data/sst2/train.jsonl, data/boolq/train.jsonl]
  calibration: [data/mnli/calibration.jsonl, data/sst2/calibration.jsonl, data/boolq/calibration.jsonl]
  eval:
    mnli: data/mnli/dev.jsonl
    sst2: data/sst2/dev.jsonl
    boolq: data/boolq/dev.jsonl
eval:
  selection_set: mnli
  constraints: ["ece <= 0.05"]

selection_set and constraints choose the checkpoint that becomes best/: here, MNLI accuracy picks the checkpoint, subject to a calibration error of at most 0.05.

Losses and distillation

train:
  loss:
    type: cross_entropy        # or focal
    label_smoothing: 0.05
    distill_alpha: 0.5         # mix in the teachers' probabilities
    distill_temperature: 2.0

Prompts

data.prompt.style picks the prompt: d3 (the default, Decision 3.0’s byte-exact prompt, servable by vLLM) or autojev. Decoders use the tokenizer’s own chat template, and you can add your own format with decisionformer.register_format. Training examples longer than data.max_length have the end of their state cut, while loaded and served models never truncate and refuse what does not fit.

From a laptop to many GPUs

Multi-GPU

Launch with torchrun. NCCL is used on GPUs and gloo on CPU.

torchrun --nproc-per-node 8 -m decisionformer train configs/recipes/deciformer15k_qwen35.yaml

Every rank computes the same step plan with balanced token counts. Ranks with fewer micro-batches run zero-weight padding batches so synchronization never deadlocks, and the loss is scaled so that DDP averaging gives the global-batch mean. train.shard_optimizer_states: true adds ZeRO-1.

Resume exactly

train:
  save_every: 500
  resume: auto        # rerun the same command to continue

Resuming restores the optimizer, data order and random state, so a run that you pause and continue ends with the same weights as one that never stopped.

On a laptop

The device is picked automatically: a CUDA GPU, then Apple silicon (MPS), then the CPU. Set DECISIONFORMER_DEVICE=cpu to force the CPU. This recipe fully fine-tunes SmolLM2-135M-Instruct for 300 steps:

pip install "decisionformer[data]"
decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k
decisionformer train configs/recipes/local_smollm2.yaml
decisionformer serve runs/local-smollm2-135m/best

On MPS and CPU, training runs in float32. The recipe uses micro-batches of 4, at most 4,096 tokens per micro-batch and gradient checkpointing. On a 6-core CPU it peaked at 3.7 GB of memory at about 15 s per step. If memory is tight, lower them: train.micro_batch_size=2 train.token_budget=2048.

Evaluate, compare and ship

decisionformer eval runs/my-decider/best data/dev.jsonl --predictions dev.predictions.jsonl
decisionformer predict runs/my-decider/best requests.jsonl --output responses.jsonl
decisionformer compare a.predictions.jsonl b.predictions.jsonl     # paired bootstrap over families
decisionformer push runs/my-decider/best my-org/my-decider          # private unless --public

• eval scores any checkpoint on labeled requests, and --calibration FILE refits the temperatures on a file you choose first.
• compare tells you whether the difference between two models is larger than noise, by resampling whole families.
• A checkpoint is a plain Hugging Face model repo: the backbone’s config.json and safetensors, the tokenizer, decision_config.json, readout.safetensors (or heads.safetensors) and a model card.
• from_pretrained also loads other checkpoints of this standard, such as Decision 3.0 (vllm-sr/d3*) and AutoJev exports.

Serve with vLLM

Code-readout checkpoints on a Qwen3.5 backbone, the Decision 3.0 layout, can be served by vLLM behind the same API. Other checkpoints fail with a clear error that points to the default backend.

pip install "decisionformer[vllm]"
decisionformer serve vllm-sr/d3-edge --backend vllm

Run the Decision Index

pip install "decision-index @ git+https://github.com/apolinario/decision-index"
python -m decision_index pipeline \
    --engine decisionformer.integrations.decision_index:DecisionFormerEngine \
    --option model=my-org/my-decider --out runs/index-my-decider

Before you start

• Bidirectional attention on a decoder needs checking per model, as described in the attention section.
• The vLLM backend serves only code-readout checkpoints with a Qwen3.5 or Qwen3.8 backbone. Everything else runs on the default transformers backend.
• trust_remote_code models are not supported.
• Checkpoints are large and every export writes full weights, so keep an eye on disk space when you run many experiments.
• In a container with a CPU quota, set OMP_NUM_THREADS to a small number. PyTorch starts one thread per host core, which can slow training badly until you limit it.

What we are building next

The roadmap in the repository is FSDP2 for full fine-tuning beyond the DDP limits, encoding a request’s state once for all of its questions, image and video inputs, and more heads such as multi-label and regression.

If you train a decision model with DecisionFormer, we would like to hear what worked and what did not. Open an issue on GitHub or write to support@frontiersmind.ai.