Textual Self-attention Network: Test-Time Preference Optimization through Textual Gradient-based Attention

TL;DR

This paper introduces TSAN, a parameter-free, text-based self-attention mechanism that optimizes language model responses at test time by analyzing, weighing, and synthesizing multiple candidates to better align with human preferences.

Contribution

TSAN is a novel, parameter-free test-time preference optimization method that emulates self-attention in natural language, enabling effective synthesis of multiple candidate responses without fine-tuning.

Findings

TSAN outperforms supervised models with just three test-time iterations.

TSAN surpasses existing test-time alignment methods.

The method effectively leverages multiple candidate responses for better alignment.

Abstract

Large Language Models (LLMs) have demonstrated remarkable generalization capabilities, but aligning their outputs with human preferences typically requires expensive supervised fine-tuning. Recent test-time methods leverage textual feedback to overcome this, but they often critique and revise a single candidate response, lacking a principled mechanism to systematically analyze, weigh, and synthesize the strengths of multiple promising candidates. Such a mechanism is crucial because different responses may excel in distinct aspects (e.g., clarity, factual accuracy, or tone), and combining their best elements may produce a far superior outcome. This paper proposes the Textual Self-Attention Network (TSAN), a new paradigm for test-time preference optimization that requires no parameter updates. TSAN emulates self-attention entirely in natural language to overcome this gap: it analyzes multiple candidates by formatting them into textual keys and values, weighs their relevance using an LLM-based attention module, and synthesizes their strengths into a new, preference-aligned response under the guidance of the learned textual attention. This entire process operates in a textual gradient space, enabling iterative and interpretable optimization. Empirical evaluations demonstrate that with just three test-time iterations on a base SFT model, TSAN outperforms supervised models like Llama-3.1-70B-Instruct and surpasses the current state-of-the-art test-time alignment method by effectively leveraging multiple candidate solutions.