Resumo Técnico: Estudando a "Soupabilidade" de Documentos em Modelos de Espaço de Estados

Declaração do Problema

Muitas tarefas de PLN do mundo real, como questionamento multi-documento (QA), sumarização científica e análise jurídica, exigem o raciocínio sobre corpora inteiros em vez de documentos individuais longos. Os modelos atuais baseados em Transformers lutam contra isso devido aos custos proibitivos de atenção $O(L^2)$ e à inflexibilidade das estratégias de codificação monolíticas. Nas abordagens padrão, todos os documentos são concatenados em uma única sequência antes do processamento; consequentemente, qualquer modificação no corpus (por exemplo, adicionar ou remover um único documento) exige a recodificação de todo o input do zero. Isso impede o reúso modular de representações de documentos e cria um gargalo para sistemas de recuperação escaláveis e dinâmicos.

Embora os Modelos de Espaço de Estados Estruturados (SSMs) como Mamba e Mamba2 ofereçam o processamento de sequências em tempo linear e comprimam entradas em estados ocultos de comprimento fixo, eles são tipicamente implantados usando a mesma estratégia de concatenação monolítica, herdando as mesmas limitações de recodificação.

Metodologia: Document Souping (Sopa de Documentos)

Os autores propõem o Document Souping, uma estratégia inspirada no model souping (média de pesos de checkpoints ajustados) mas aplicada a estados ocultos. A hipótese central é que os mecanismos de recorrência linear dos SSMs permitem que as representações de documentos codificadas independentemente sejam fundidas post hoc via operações lineares simples (por exemplo, média) para formar um estado composto adequado para o raciocínio subsequente.

A Arquitetura

1. Codificação Independente: Cada documento $d_i$ em um corpus é codificado independentemente por um codificador SSM compartilhado ($SSM_\theta$), produzindo um conjunto de estados ocultos por camada $\{h^{(l)}_i\}$.
2. Agrupamento (Pooling): Esses estados por documento são agregados em cada camada usando um operador de agrupamento comutativo (por exemplo, média elemento a elemento, soma ou max) para criar um estado "sopado" único $h^{(l)}_{soup}$.
3. Decodificação: Os estados agrupados $\{h^{(1)}_{soup}, \dots, h^{(L)}_{soup}\}$ são injetados no decodificador junto com a consulta $q$ para gerar a resposta $\hat{y}$.

Treinamento e Otimização

• Ajuste Fino (Finetuning): O artigo demonstra que os SSMs pré-treinados não suportam nativamente o souping; é necessário o ajuste fino completo do codificador-decodificador para ensinar o modelo a interpretar estados agrupados de forma eficaz.
• Eficiência de Memória: Para lidar com a sobrecarga de memória de retropropagação através de múltiplos codificadores de documentos independentes, os autores utilizam checkpointing de ativação em nível de documento. Isso recomputa as ativações do forward pass durante o backward pass, permitindo que o uso de memória permaneça constante independentemente do número de documentos ($k$).
• Estratégia de Agrupamento: Experimentos extensivos identificam a média elemento a elemento sem normalização como o método de agregação mais robusto e estável, superando a soma (que sofre com crescimento ilimitado de magnitude) e o max pooling.

Principais Contribuições

1. Estudo Empírico Sistemático: A primeira investigação abrangente sobre a "soupabilidade" — a capacidade dos SSMs de fundir representações de documentos codificadas independentemente para raciocínio multi-documento.
2. Validação de Desempenho: Demonstração de que modelos Mamba2 finetunados com representações "sopadas" alcançam desempenho competitivo ou superior em QA de múltiplos saltos (HotpotQA), QA de documentos longos (RACE, QuALITY) e recuperação esparsa (NIAH) comparado às bases de comparação de concatenação padrão.
3. Escalabilidade e Generalização: Evidência de que a abordagem escala para centenas de documentos (testada até 256) e generaliza para diversos tamanhos de contexto, desde que o modelo seja finetunado em um número suficiente de segmentos.
4. Especificidade Arquitetural: Prova de que esta abordagem modular é unicamente adequada para SSMs. Um experimento análogo de "cache souping" em Transformers (LLaMA 3) falhou significativamente, destacando que o estado recorrente de tamanho fixo dos SSMs é um habilitador crítico para esta técnica.
5. Eficiência de Inferência: Um novo fluxo de trabalho de inferência onde os documentos são codificados uma vez, armazenados em cache e reassemblados dinamicamente, oferecendo economias substanciais de latência (até 14,6x de aceleração para 32k tokens) em comparação com a recodificação de contextos concatenados.

Resultados Experimentais

Os autores avaliaram modelos Mamba2-8B e Mamba2-2.7B através de vários benchmarks:

• QA de Múltiplos Saltos (HotpotQA): Modelos com ajuste fino completo de codificador-decodificador e média de agrupamento atingiram 47.8 EM em 5 documentos, superando o ajuste fino de apenas decodificador e igualando ou excedendo as bases de comparação de concatenação. O método permaneceu robusto conforme o número de distratores aumentava.
• QA de Documentos Longos (RACE/QuALITY): O souping superou substancialmente as abordagens tradicionais de concatenação em benchmarks de documentos longos, particularmente quando a configuração de treinamento se alinhava ao tamanho do soup em tempo de teste.
• Recuperação Esparsa (NIAH): Modelos treinados em um maior número de segmentos (ex: 8 segmentos) mantiveram alta precisão (71.45% EM) ao serem testados em 16 segmentos, enquanto modelos treinados em menos segmentos falharam em generalizar.
• Escalabilidade: Embora um modelo treinado em 64 documentos tenha falhado em extrapolar para 256, o ajuste fino contínuo em 128 documentos permitiu que o modelo alcançasse 37.25 EM em 256 documentos, provando que a arquitetura pode se adaptar a contextos vastos.
• Latência: Para um contexto de 32k tokens, o processamento de uma consulta usando estados "sopados" em cache levou 65.8 ms, comparado a 963.2 ms para o método padrão de concatenação.

Significância e Alegações

O artigo afirma que o document souping desbloqueia um novo paradigma para o raciocínio de corpus modular e escalável. Ao desacoplar a codificação do documento do processamento da consulta, a abordagem permite:

• Atualizações Eficientes: Documentos podem ser codificados uma vez e armazenados em cache, permitindo atualizações dinâmicas de corpus e recuperação sem reprocessar o input completo.
• RAG Escalável: O método serve como um substituto direto para a etapa de "concatenar passagens recuperadas" em sistemas de Geração Aumentada por Recuperação (RAG), suportando números variáveis, mas limitados, de documentos de forma eficiente.
• Adequação Arquitetural: O trabalho estabelece que as propriedades de recorrência linear dos SSMs os tornam unicamente capazes de suportar a composição de nível de estado, uma capacidade que os Transformers padrão não possuem devido à sua dependência de caches KV de comprimento variável.

Os autores concluem que, embora a extrapolação para tamanhos de soup não vistos possa ser desafiadora, estratégias de ajuste fino apropriadas permitem que os SSMs lidem efetivamente com conjuntos de documentos de grande escala e dinâmicos, oferecendo uma solução leve e eficaz para o raciocínio de contexto longo.