Classificação Escalável em Contexto com 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

Resumo

A Classificação em Contexto (In-context Ranking - ICR) é um paradigma emergente para Recuperação de Informação (Information Retrieval - IR), que aproveita a compreensão contextual de LLMs (Large Language Models) ao incorporar diretamente a descrição da tarefa, documentos candidatos e a consulta no prompt de entrada do modelo, atribuindo ao LLM a tarefa de identificar os documentos relevantes. Embora seja eficaz, a eficiência é um desafio significativo nesse paradigma, especialmente à medida que a lista de candidatos cresce devido à escala quadrática/super-linear da operação de atenção com o comprimento do contexto. Para isso, este artigo primeiro identifica estruturas inerentes e exploráveis na atenção de LLMs ajustados para ICR: (1) esparsidade de bloco interdocumento: a atenção é densa dentro de cada bloco de documento, mas esparsa entre diferentes documentos no contexto; e (2) relevância de bloco consulta-documento: as pontuações de atenção de certos tokens da consulta para um bloco de documento nas camadas intermediárias correlacionam-se fortemente com a relevância real desse documento. Motivados por essas observações, introduzimos o BlockRank (Classificação em Blocos em Contexto), um método novo que adapta a operação de atenção em um LLM ao (a) impor arquitetonicamente a esparsidade de bloco interdocumento observada, reduzindo a complexidade da atenção de quadrática para linear sem perda de desempenho, e (b) otimizar a relevância de bloco consulta-documento para documentos verdadeiramente relevantes durante o ajuste fino usando um objetivo de treinamento contrastivo auxiliar, melhorando a recuperação na atenção. Experimentos no BEIR, MSMarco e NQ com o Mistral-7B demonstram que o FLARE Mistral iguala ou supera os classificadores listwise SOTA existentes e a linha de base de ajuste fino controlado, enquanto é significativamente mais eficiente na inferência (4,7x para 100 documentos do MSMarco em contexto) e escala graciosamente para listas curtas de contexto longo, cerca de 500 documentos em contexto (aproximadamente 100K de comprimento de contexto) em menos de um segundo, apresentando uma solução escalável e eficaz 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.