Reproduce “Set-LLM: A Permutation-Invariant LLM” on OpenResearch

About this paper

Beni Egressy, Jan Stühmer

Set-LLM, developed by researchers at HITS and KIT, introduces architectural modifications to Large Language Models (LLMs), enabling them to process mixed set-text inputs without sensitivity to element order. The framework achieves robust, permutation-invariant performance, consistently matching or outperforming traditional baselines and computationally expensive methods while maintaining the efficiency of a single LLM inference.

Key insights

• The core attention mechanism of transformers is inherently permutation-invariant; order sensitivity arises from how positional encodings and causal masks imbue sequential information.
• Permutation invariance for sets can be achieved by selectively encoding intra-sequence order while explicitly discarding inter-set element order at both the positional encoding and attention mask levels.
• A formal proof confirms that an attention layer with SetPE and SetMask is set permutation equivariant, guaranteeing that the stacked network will produce truly invariant outputs.

Summary by alphaXiv.