Memorización de Contexto para la Generación Eficiente de Contextos Largos

Resumen

Las aplicaciones modernas de modelos de lenguaje grande (LLM) dependen cada vez más de prefijos de condicionamiento largos para controlar el comportamiento del modelo en tiempo de inferencia. Si bien la inferencia aumentada por prefijo es efectiva, presenta dos limitaciones estructurales: i) la influencia del prefijo se desvanece a medida que avanza la generación, y ii) el cálculo de atención sobre el prefijo escala linealmente con su longitud. Los enfoques existentes mantienen el prefijo en la atención mientras lo comprimen, o lo internalizan en los parámetros del modelo mediante entrenamiento basado en gradientes. El primero aún atiende al prefijo durante la inferencia, mientras que el segundo requiere mucho entrenamiento y es inadecuado para actualizaciones del prefijo. Para abordar estos problemas, proponemos la memoria de estados de atención, un enfoque sin entrenamiento que externaliza el prefijo en una memoria ligera basada en búsqueda de estados de atención precomputados entre los tokens del prefijo y los de la consulta. En ManyICLBench con LLaMA-3.1-8B, nuestro método mejora la precisión sobre el aprendizaje en contexto con presupuestos de memoria de 1K a 8K, al tiempo que reduce la latencia de atención en 1.36x con 8K, y supera el rendimiento de RAG con atención completa en el benchmark NBA utilizando solo el 20% de su huella de memoria.

English

Modern large language model (LLM) applications increasingly rely on long conditioning prefixes to control model behavior at inference time. While prefix-augmented inference is effective, it incurs two structural limitations: i) the prefix's influence fades as generation proceeds, and ii) attention computation over the prefix scales linearly with its length. Existing approaches either keep the prefix in attention while compressing it, or internalize it into model parameters through gradient-based training. The former still attends to the prefix at inference, while the latter is training-intensive and ill-suited to prefix updates. To address these issues, we propose attention-state memory, a training-free approach that externalizes the prefix into a lightweight, lookup-based memory of precomputed attention states between prefix and query tokens. On ManyICLBench with LLaMA-3.1-8B, our method improves accuracy over in-context learning at 1K-8K memory budgets while reducing attention latency by 1.36x at 8K, and surpasses full-attention RAG performance on NBA benchmark using only 20% of its memory footprint.