Sigmoid Head for Quality Estimation under Language Ambiguity

TL;DR

Ambiguity in natural language causes LM probabilities to underreport output quality. The authors propose the Sigmoid Head, an additional unembedding head with sigmoid activation that computes $p_i(v)=\sigma(z^{qe}_{i}[v])$ where $z^{qe}_{i}=W_{qe} h_i + b_{qe}$, trained on the same data as the original LM, and paired with ambiguity-informed negative sampling that excludes dominant tokens. This method yields stronger quality signals than the standard softmax head, requires no human-annotated quality data, and demonstrates robustness to domain shift, outperforming unsupervised baselines and rivaling supervised QE on out-of-domain MT, paraphrase, and QA tasks. The approach is computationally efficient at training and inference since only a small negative-sample subset updates the unembedding, and it generalizes across multiple model families and evaluation settings.

Abstract

Language model (LM) probability is not a reliable quality estimator, as natural language is ambiguous. When multiple output options are valid, the model's probability distribution is spread across them, which can misleadingly indicate low output quality. This issue is caused by two reasons: (1) LMs' final output activation is softmax, which does not allow multiple correct options to receive high probabilities simultaneuously and (2) LMs' training data is single, one-hot encoded references, indicating that there is only one correct option at each output step. We propose training a module for Quality Estimation on top of pre-trained LMs to address these limitations. The module, called Sigmoid Head, is an extra unembedding head with sigmoid activation to tackle the first limitation. To tackle the second limitation, during the negative sampling process to train the Sigmoid Head, we use a heuristic to avoid selecting potentially alternative correct tokens. Our Sigmoid Head is computationally efficient during training and inference. The probability from Sigmoid Head is notably better quality signal compared to the original softmax head. As the Sigmoid Head does not rely on human-annotated quality data, it is more robust to out-of-domain settings compared to supervised QE.

Figures (2)

• Figure 1: Extended language model (LM) architecture with our proposed Sigmoid Head. The weights of the original components from the LM are kept unchanged. We initialize the Sigmoid Head from the original softmax head, and train it to predict independent scores for each token in the vocabulary.
• Figure 2: Ground-truth versus predicted quality scores.