MoME: miscela di embedding di memoria per la ricerca sparsa contestuale

Abstract degli autori

La necessità di scalare in modo efficiente i modelli linguistici di grandi dimensioni ha portato allo sviluppo di meccanismi di capacità sparsa, come Mixture-of-Experts e, più di recente, la memoria condizionale: tabelle di embedding indicizzate per token che ampliano il modello di base con consultazioni parametriche a basso costo. I metodi esistenti basati su embedding di memoria effettuano il recupero tramite una funzione deterministica della forma superficiale, accorpando sensi contestuali diversi dello stesso token (per esempio, Python come linguaggio di programmazione e come animale) in un’unica voce fissa. Presentiamo Mixture of Memory Embeddings (MoME), un meccanismo di memoria sensibile al contesto che sostituisce la singola riga di memoria di ogni token con una miscela di M slot e usa un gate appreso sullo stato nascosto per scegliere quali slot leggere in ciascuna posizione. Negli esperimenti controllati di preaddestramento condotti con le architetture di base nanochat, Llama-3/MobileLLM e Qwen3, MoME supera le baseline Value Embedding, Bigram e STEM a parità di numero di parametri e di FLOP di addestramento, mostra una tendenza più promettente all’aumento della dimensione della memoria su scala inferiore al miliardo e mantiene l’efficienza sia nell’addestramento sia nell’inferenza. Le analisi qualitative del routing di token polisemici suggeriscono inoltre che la miscela appresa presenta un certo grado di interpretabilità semantica, indirizzando lo stesso token superficiale verso slot di memoria distinti a seconda del senso.