Locas: Seus Modelos são Inicializadores Fundamentados de Memórias Paramétricas com Suporte Local

Resumo

Neste artigo, visamos estabelecer uma ponte entre o treinamento durante o teste e um novo tipo de memória paramétrica que pode ser flexivelmente descarregada ou fundida com os parâmetros do modelo. Apresentamos o Locas, uma memória paramétrica de Suporte Local que partilha a estrutura dos blocos FFN dos transformadores modernos, permitindo que seja flexivelmente permanentizada nos parâmetros do modelo, ao mesmo tempo que suporta uma aprendizagem contínua eficiente. Discutimos duas variantes principais do Locas: uma com um design convencional de MLP de duas camadas, que possui uma garantia teórica mais clara; e outra que partilha a mesma estrutura GLU-FFN com os LLMs de última geração, podendo ser facilmente acoplada a modelos existentes para uma aprendizagem contínua eficiente em termos de parâmetros e de computação. Crucialmente, demonstramos que uma inicialização adequada destas memórias de baixo rank, no estilo FFN lateral — realizada de forma fundamentada através da reutilização de parâmetros, ativações e/ou gradientes do modelo — é essencial para uma convergência rápida, uma generalização melhorada e a prevenção do esquecimento catastrófico. Validamos o mecanismo de memória proposto nas tarefas de modelagem de linguagem de livros inteiros PG-19 e de resposta a perguntas em diálogos de contexto longo LoCoMo. Com apenas 0,02% de parâmetros adicionais no caso mais reduzido, o Locas-GLU é capaz de armazenar informações do contexto anterior, mantendo uma janela de contexto muito menor. Adicionalmente, também testamos a perda de capacidade geral do modelo após memorizar um livro inteiro com o Locas, através de uma avaliação comparativa MMLU. Os resultados mostram a capacidade promissora do Locas em permanentizar o contexto passado em conhecimento paramétrico, com uma minimização do esquecimento catastrófico do conhecimento interno existente do modelo.

English

In this paper, we aim to bridge test-time-training with a new type of parametric memory that can be flexibly offloaded from or merged into model parameters. We present Locas, a Locally-Supported parametric memory that shares the design of FFN blocks in modern transformers, allowing it to be flexibly permanentized into the model parameters while supporting efficient continual learning. We discuss two major variants of Locas: one with a conventional two-layer MLP design that has a clearer theoretical guarantee; the other one shares the same GLU-FFN structure with SOTA LLMs, and can be easily attached to existing models for both parameter-efficient and computation-efficient continual learning. Crucially, we show that proper initialization of such low-rank sideway-FFN-style memories -- performed in a principled way by reusing model parameters, activations and/or gradients -- is essential for fast convergence, improved generalization, and catastrophic forgetting prevention. We validate the proposed memory mechanism on the PG-19 whole-book language modeling and LoCoMo long-context dialogue question answering tasks. With only 0.02\% additional parameters in the lowest case, Locas-GLU is capable of storing the information from past context while maintaining a much smaller context window. In addition, we also test the model's general capability loss after memorizing the whole book with Locas, through comparative MMLU evaluation. Results show the promising ability of Locas to permanentize past context into parametric knowledge with minimized catastrophic forgetting of the model's existing internal knowledge.