Token Constraint Decoding Improves Robustness on Question Answering for Large Language Models

Overview

Token Constraint Decoding (TCD) is an inference-time decoding method designed to keep multiple-choice QA behavior stable when prompts contain noise or adversarial perturbations.

Instead of changing model weights, TCD modifies the decoding process itself, making it attractive for teams that need stronger robustness without retraining costs.

Method Summary

The method applies a soft penalty when local token decisions drift beyond a tolerance range, encouraging more consistent token trajectories during answer generation.

This acts like a lightweight regularizer at inference time and can be combined with prompt engineering strategies to further stabilize output quality.

Experimental Results

Evaluations on CommonsenseQA, MMLU, and MMLU-Pro show meaningful robustness improvements under perturbed input conditions, with particularly notable gains in smaller model regimes.

The analysis also surfaces model-family differences in optimal penalty scheduling, revealing a practical robustness-versus-overconfidence trade-off.

Practical Impact

Because TCD is model-agnostic and inference-only, it can be integrated into existing QA systems with lower operational risk, making it useful for reliability-sensitive applications.