Clasificación Escalable en Contexto con Modelos Generativos

Scalable In-context Ranking with Generative Models

October 6, 2025

Autores: Nilesh Gupta, Chong You, Srinadh Bhojanapalli, Sanjiv Kumar, Inderjit Dhillon, Felix Yu

cs.AI

Resumen

La Clasificación en Contexto (ICR, por sus siglas en inglés) es un paradigma emergente en la Recuperación de Información (IR), que aprovecha la comprensión contextual de los Modelos de Lenguaje de Gran Escala (LLMs) al incorporar directamente la descripción de la tarea, los documentos candidatos y la consulta en el mensaje de entrada del modelo, asignando al LLM la tarea de identificar los documentos relevantes. Aunque es efectivo, la eficiencia es un desafío significativo en este paradigma, especialmente a medida que la lista de candidatos crece debido al escalamiento cuadrático/super-lineal de la operación de atención con la longitud del contexto. Con este fin, este artículo identifica primero estructuras inherentes y explotables en la atención de los LLMs ajustados para ICR: (1) esparcidad de bloques inter-documentos: la atención es densa dentro de cada bloque de documentos pero dispersa entre diferentes documentos en el contexto; y (2) relevancia de bloques consulta-documento: las puntuaciones de atención de ciertos tokens de la consulta a un bloque de documentos en las capas intermedias se correlacionan fuertemente con la relevancia real de ese documento. Motivados por estas observaciones, introducimos BlockRank (Clasificación en Contexto por Bloques), un método novedoso que adapta la operación de atención en un LLM mediante (a) la imposición arquitectónica de la esparcidad de bloques inter-documentos observada, reduciendo la complejidad de la atención de cuadrática a lineal sin pérdida de rendimiento, y (b) la optimización de la relevancia de bloques consulta-documento para documentos verdaderamente relevantes durante el ajuste fino utilizando un objetivo de entrenamiento contrastivo auxiliar, mejorando la recuperación en la atención. Los experimentos en BEIR, MSMarco y NQ con Mistral-7B demuestran que FLARE Mistral iguala o supera a los clasificadores listwise de última generación existentes y a la línea base ajustada finamente controlada, siendo significativamente más eficiente en la inferencia (4.7x para 100 documentos de MSMarco en contexto) y escalando de manera elegante a listas cortas de contexto largo, alrededor de 500 documentos en contexto (aproximadamente 100K de longitud de contexto) en un segundo, presentando una solución escalable y efectiva para ICR.

English

In-context Ranking (ICR) is an emerging paradigm for Information Retrieval (IR), which leverages contextual understanding of LLMs by directly incorporating the task description, candidate documents, and the query into the model's input prompt and tasking the LLM to identify relevant document(s). While it is effective, efficiency is a significant challenge in this paradigm, especially as the candidate list grows due to quadratic/super-linear scaling of attention operation with context length. To this end, this paper first identifies inherent and exploitable structures in the attention of LLMs finetuned for ICR: (1) inter-document block sparsity: attention is dense within each document block but sparse across different documents in the context; and (2) query-document block relevance: the attention scores from certain query tokens to a document block in middle layers strongly correlate with that document's actual relevance. Motivated by these observations, we introduce BlockRank (Blockwise In-context Ranking), a novel method that adapts the attention operation in an LLM by (a) architecturally enforcing the observed inter-document block sparsity, reducing attention complexity from quadratic to linear without loss in performance, and (b) optimizing query-document block relevance for true relevant documents during fine-tuning using an auxiliary contrastive training objective, improving retrieval in attention. Experiments on BEIR, MSMarco and NQ with Mistral-7B demonstrate that FLARE Mistral matches or outperforms existing SOTA listwise rankers and controlled fine-tuned baseline while being significantly more efficient at inference (4.7x for 100 MSMarco documents in context) and scaling gracefully to long-context shortlists, around 500 documents in-context (approximately 100K context length) within a second, presenting a scalable and effective solution for ICR.