--- license: cc-by-nc-4.0 library_name: pytorch tags: - quality-estimation - sigmoid-head - machine-translation - confidence-estimation - tower base_model: Unbabel/TowerInstruct-7B-v0.2 --- # Sigmoid Head for TowerInstruct-7B-v0.2 This repo hosts a **sigmoid quality-estimation (QE) head** trained on top of [`Unbabel/TowerInstruct-7B-v0.2`](https://huggingface.co/Unbabel/TowerInstruct-7B-v0.2). 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 (e.g. machine translation). - **Base model:** `Unbabel/TowerInstruct-7B-v0.2` (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:** `[32007, 4096]` - **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 (source-prompt, hypothesis) pair, compute a per-token confidence for the hypothesis. Useful for QE on outputs from any MT system. ```python import torch from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer BASE = "Unbabel/TowerInstruct-7B-v0.2" HEAD = "tuanh23/SigmoidHead-TowerInstruct-7B-v0.2" 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() # Same chat-template format the head was trained on (see prepare_data.py). src_lang, tgt_lang = "English", "German" src = "The cat sat on the mat." hypothesis = "Die Katze saß auf der Matte." user_msg = {"role": "user", "content": f"Translate the following text from {src_lang} into {tgt_lang}.\n{src_lang}: {src}.\n{tgt_lang}: "} asst_msg = {"role": "assistant", "content": " " + hypothesis} # Full conversation -> input_ids for the model input_ids = tokenizer.apply_chat_template( [user_msg, asst_msg], tokenize=True, add_generation_prompt=False, return_tensors="pt" ).to(device) # Same encoding but with the generation prompt added after the user turn -> tells us # where the assistant content begins inside `input_ids`. 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 (hypothesis + closing chat tokens): hyp_conf = conf[0, prompt_len - 1:] hyp_tokens = tokenizer.convert_ids_to_tokens(input_ids[0, prompt_len:].tolist()) print("Hypothesis:", hypothesis) for tok, s in zip(hyp_tokens, hyp_conf.tolist()): print(f" {tok!r:>20s} conf={s:.4f}") print(f"Sentence-level (mean): {hyp_conf.mean().item():.4f}") # Expected output: # Hypothesis: Die Katze saß auf der Matte. # '▁Die' conf=0.9999 # '▁Kat' conf=0.9995 # 'ze' conf=0.9992 # '▁sa' conf=0.9993 # 'ß' conf=1.0000 # '▁auf' conf=1.0000 # '▁der' conf=0.9983 # '▁Mat' conf=0.9992 # 'te' conf=0.9999 # '.' conf=0.9897 # '<|im_end|>' conf=1.0000 # '▁' conf=1.0000 # '<0x0A>' conf=1.0000 # Sentence-level (mean): 0.9988 ``` ### 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. ```python import torch from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer BASE = "Unbabel/TowerInstruct-7B-v0.2" HEAD = "tuanh23/SigmoidHead-TowerInstruct-7B-v0.2" 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() src_lang, tgt_lang = "English", "German" src = "The cat sat on the mat." messages = [{"role": "user", "content": f"Translate the following text from {src_lang} into {tgt_lang}.\n{src_lang}: {src}.\n{tgt_lang}: "}] 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=64, 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] translation = tokenizer.decode(gen_ids[0], skip_special_tokens=True) print("Translation:", translation) 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: # Translation: Die Katze saß auf der Matte. # '▁Die' conf=0.9999 # '▁Kat' conf=0.9994 # 'ze' conf=0.9991 # '▁sa' conf=0.9993 # 'ß' conf=1.0000 # '▁auf' conf=1.0000 # '▁der' conf=0.9983 # '▁Mat' conf=0.9992 # 'te' conf=0.9999 # '.' conf=0.9900 # '<|im_end|>' conf=1.0000 # Sentence-level (mean): 0.9986 ``` ## 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. ## Code Training and evaluation code: .