Set-LLM: A Permutation-Invariant LLM

TL;DR

Set-LLM addresses the critical vulnerability of input-order sensitivity in decoder-only LLMs by embedding permutation invariance into the architecture. It achieves this through four changes: removing sequential position encoding, removing the causal mask, and introducing permutation-invariant SetPE and SetMask, with formal proofs of invariance. Empirically, Set-LLM matches or exceeds baseline performance across four MCQA datasets while eliminating order-based accuracy degradation and avoiding runtime penalties, even on adversarial input orders. The approach generalizes across multiple base LLMs and shows strong out-of-distribution robustness, offering a more reliable framework for MCQA, model evaluation, and mixed set-text tasks. This work introduces the first permutation-invariant decoder-only LLM and opens avenues for robust evaluation and reasoning pipelines in real-world AI systems.

Abstract

While large language models (LLMs) demonstrate impressive capabilities across numerous applications, their robustness remains a critical concern. This paper is motivated by a specific vulnerability: the order sensitivity of LLMs. This vulnerability manifests itself as the order bias observed when LLMs decide between possible options (for example, a preference for the first option) and the tendency of LLMs to provide different answers when options are reordered. The use cases for this scenario extend beyond the classical case of multiple-choice question answering to the use of LLMs as automated evaluators in AI pipelines, comparing output generated by different models. We introduce Set-LLM, a novel architectural adaptation for pretrained LLMs that enables the processing of mixed set-text inputs with permutation invariance guarantees. The adaptations involve a new attention mask and new positional encodings specifically designed for sets. We provide a theoretical proof of invariance and demonstrate through experiments that Set-LLM can be trained effectively, achieving comparable or improved performance and maintaining the runtime of the original model, while eliminating order sensitivity.

Figures (8)

• Figure 1: An example of the vulnerability of LLMs to choice permutations. The LLM's response changes simply due to a reordering of the answer options. (Example for illustrative purposes only.) Set-LLM eliminates this vulnerability by building invariance directly into the model architecture.
• Figure 2: Three types of attention masks and their corresponding directed graphs. The colored squares on the left indicate attention scores that are not masked. For example, in a causal mask, the $4^{\text{th}}$ token attends to the first $4$ tokens, and the remaining tokens are masked. The circles on the right represent the tokens as nodes of an attention graph. Red, orange, and yellow tokens correspond to the prompt, and blue tokens correspond to the response. Orange and yellow tokens corresponds to elements of a set. Causal masks are standard in decoder-only LMs, whereas prefix masks are used in bidirectional encoder-decoder LMs. SetMask is introduced in this work. Some edges are grayed out and self-loops are omitted to improve clarity. In addition, the figure shows three types of token positions, standard consecutive positions (PE), no positional encoding (NoPE), and set position encoding (SetPE). These are indicated by the numbers inside the token nodes on the right.
• Figure 3: An example of a multiple-choice question with set positional encoding (SetPE) positions.
• Figure 4: A failure case for an LLM with set position encoding (SetPE) but without set attention mask (SetMask). SetPE positions are shown for the tokens that are part of the set (of facts). Without the SetMask, the model is unable to distinguish the two inputs. The model can not "know" which position $8$ token belongs to which position $10$ token.
• Figure 5: Causal Mask+PE$^{\text{Ultra}}$+Majority Vote vote count versus accuracy (left). Causal Mask+PE$^{\text{Ultra}}$+Majority Vote vote count versus agreement rate with Set-LLM (right) on PIQA.

Theorems & Definitions (4)

• Claim 1
• Claim 2
• Claim 3
• Claim 4