Taint-aware, entropy-based set block decoding for language model-based token prediction

Abstract

Certain aspects of the disclosure provide techniques for language model token prediction. A method generally includes processing an input token sequence to generate, for each respective token position of a plurality of token positions in an output token sequence: an entropy score based on a probability distribution for a set of candidate output tokens associated with the respective token position; and a taint score indicative of an extent to which the probability distribution is influenced by a untrusted subset of tokens of the input token sequence; for each respective token position, determining whether the entropy score satisfies an adjusted entropy threshold that is based on at least the taint score associated with the respective token position; based on the determination for each respective token position: performing a set block decoding operation; performing a single-token decoding operation; or generating a message in accordance with a refusal policy.