Sigmoid Head for OLMo-2-0425-1B-SFT

This repo hosts a sigmoid quality-estimation (QE) head trained on top of allenai/OLMo-2-0425-1B-SFT.

It is the model from the paper Sigmoid Head for Quality Estimation under Language Ambiguity. Unlike the usual softmax LM head, this head uses a sigmoid activation, so multiple equally-valid tokens can simultaneously receive high scores. This produces a more reliable per-token quality / confidence score in settings with language ambiguity.

Base model: allenai/OLMo-2-0425-1B-SFT (frozen during training)
Head type: new unembedding head — a torch.nn.Embedding(vocab_size, hidden_size) applied to the last hidden state
Activation: sigmoid (per-token, not normalized over vocab)
Shape: [100352, 2048]
Trained with: ambiguity-aware negative sampling

Files

model.safetensors — the trained head weights (single tensor weight).
config.json — SigmoidHeadConfig (vocab/hidden sizes + auto_map).
sigmoid_head.py — SigmoidHead(PreTrainedModel) definition; auto-loaded by transformers via trust_remote_code=True.

Usage

The head is loaded with transformers.AutoModel. Pass trust_remote_code=True so transformers downloads sigmoid_head.py from this repo automatically.

1. Score an existing output (teacher forcing)

Given a (prompt, completion) pair, compute a per-token confidence for the completion. Useful for QE on outputs from any generator.

import torch
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

BASE = "allenai/OLMo-2-0425-1B-SFT"
HEAD = "tuanh23/SigmoidHead-OLMo-2-0425-1B-SFT"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(BASE)
base_model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16).to(device).eval()
head = AutoModel.from_pretrained(HEAD, trust_remote_code=True).to(device).eval()

# Use the same chat-template wrapping the head was trained with.
user_msg = {"role": "user",      "content": "What is the capital of France?"}
asst_msg = {"role": "assistant", "content": "Paris."}

input_ids = tokenizer.apply_chat_template(
    [user_msg, asst_msg], tokenize=True, add_generation_prompt=False, return_tensors="pt"
).to(device)
prompt_len = tokenizer.apply_chat_template(
    [user_msg], tokenize=True, add_generation_prompt=True, return_tensors="pt"
).shape[1]

with torch.no_grad():
    out = base_model(input_ids, output_hidden_states=True)
    last_hidden = out.hidden_states[-1].float()              # [1, T, hidden]
    conf_full = head.score(last_hidden)                      # [1, T, vocab]   in (0, 1)

# Per-token confidence for the actual next token at each position (shifted by 1)
target_ids = input_ids[:, 1:]
conf = conf_full[:, :-1, :].gather(-1, target_ids.unsqueeze(-1)).squeeze(-1)  # [1, T-1]

# Confidence over just the assistant span (completion + closing chat tokens):
comp_conf = conf[0, prompt_len - 1:]
comp_tokens = tokenizer.convert_ids_to_tokens(input_ids[0, prompt_len:].tolist())

print("Completion:", asst_msg["content"])
for tok, s in zip(comp_tokens, comp_conf.tolist()):
    print(f"  {tok!r:>20s}  conf={s:.4f}")
print(f"Sentence-level (mean): {comp_conf.mean().item():.4f}")

# Expected output:
# Completion: Paris.
#                'Paris'  conf=0.9940
#                    '.'  conf=0.9603
#        '<|endoftext|>'  conf=0.9998
# Sentence-level (mean): 0.9847

2. Generate and score

The sigmoid head only needs the last-layer hidden states, which transformers.generate already returns when you ask for them. So you can generate with the base LM and score with the sigmoid head in one forward pass — no re-decoding.

import torch
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

BASE = "allenai/OLMo-2-0425-1B-SFT"
HEAD = "tuanh23/SigmoidHead-OLMo-2-0425-1B-SFT"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(BASE)
base_model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16).to(device).eval()
head = AutoModel.from_pretrained(HEAD, trust_remote_code=True).to(device).eval()

messages = [{"role": "user", "content": "What is the capital of France?"}]
input_ids = tokenizer.apply_chat_template(
    messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
).to(device)

with torch.no_grad():
    gen = base_model.generate(
        input_ids=input_ids,
        max_new_tokens=32,
        do_sample=False,                # greedy
        output_hidden_states=True,
        return_dict_in_generate=True,
    )
    # Stitch together per-step last-layer hidden states into [B, gen_len, hidden].
    # Step 0 returns hidden states for the whole prompt — keep only the last position.
    last_hidden = [step[-1] for step in gen.hidden_states]
    last_hidden[0] = last_hidden[0][:, -1:, :]
    last_hidden = torch.cat(last_hidden, dim=1).float()    # [B, gen_len, hidden]

    gen_ids = gen.sequences[:, input_ids.shape[1]:]        # [B, gen_len]
    conf_full = head.score(last_hidden)                    # [B, gen_len, vocab] in (0, 1)
    conf = conf_full.gather(-1, gen_ids.unsqueeze(-1)).squeeze(-1)  # [B, gen_len]

answer = tokenizer.decode(gen_ids[0], skip_special_tokens=True)
print("Answer:", answer)
for tok, s in zip(tokenizer.convert_ids_to_tokens(gen_ids[0].tolist()), conf[0].tolist()):
    print(f"  {tok!r:>20s}  conf={s:.4f}")
print(f"Sentence-level (mean): {conf[0].mean().item():.4f}")

# Expected output:
# Answer: The capital of France is Paris.
#                  'The'  conf=0.9675
#             'Ġcapital'  conf=0.9999
#                  'Ġof'  conf=0.9989
#              'ĠFrance'  conf=0.9993
#                  'Ġis'  conf=0.9987
#               'ĠParis'  conf=0.9989
#                    '.'  conf=0.9997
#        '<|endoftext|>'  conf=0.9998
# Sentence-level (mean): 0.9953

Why sigmoid?

A standard softmax head forces the probability mass to sum to 1 across the vocab, so when several outputs are equally valid the mass is split and valid tokens might look low-confidence. The sigmoid head decouples tokens, so all valid options can score high simultaneously — a better proxy for quality.

Citation

@article{dinh2026sigmoid,
  title   = {Sigmoid Head for Quality Estimation under Language Ambiguity},
  author  = {Dinh, Tu Anh and Niehues, Jan},
  journal = {arXiv preprint arXiv:2601.00680},
  year    = {2026}
}

Accepted to ACL 2026 (Main); proceedings not yet released.