Réseaux Delta à Portes: Amélioration de Mamba2 avec la Règle Delta

Gated Delta Networks: Improving Mamba2 with Delta Rule

December 9, 2024

Auteurs: Songlin Yang, Jan Kautz, Ali Hatamizadeh

cs.AI

Résumé

Les transformateurs linéaires ont attiré l'attention en tant qu'alternatives efficaces aux transformateurs standard, mais leurs performances dans les tâches de recherche et de contexte étendu ont été limitées. Pour remédier à ces limitations, des travaux récents ont exploré deux mécanismes distincts : le contrôle adaptatif de la mémoire par des mécanismes de filtrage et la règle de mise à jour delta pour des modifications précises de la mémoire. Nous observons que ces mécanismes sont complémentaires : le filtrage permet une effacement rapide de la mémoire tandis que la règle delta facilite les mises à jour ciblées. En nous appuyant sur cette observation, nous introduisons la règle delta filtrée et développons un algorithme d'entraînement parallèle optimisé pour le matériel moderne. Notre architecture proposée, Gated DeltaNet, surpasse de manière constante les modèles existants tels que Mamba2 et DeltaNet sur plusieurs bancs d'essai, notamment la modélisation de langage, le raisonnement de bon sens, la récupération en contexte, l'extrapolation de longueur et la compréhension de contexte étendu. Nous améliorons encore les performances en développant des architectures hybrides qui combinent des couches Gated DeltaNet avec une attention à fenêtre glissante ou des couches Mamba2, obtenant à la fois une efficacité d'entraînement améliorée et des performances de tâche supérieures.

English

Linear Transformers have gained attention as efficient alternatives to standard Transformers, but their performance in retrieval and long-context tasks has been limited. To address these limitations, recent work has explored two distinct mechanisms: gating for adaptive memory control and the delta update rule for precise memory modifications. We observe that these mechanisms are complementary: gating enables rapid memory erasure while the delta rule facilitates targeted updates. Building on this insight, we introduce the gated delta rule and develop a parallel training algorithm optimized for modern hardware. Our proposed architecture, Gated DeltaNet, consistently surpasses existing models like Mamba2 and DeltaNet across multiple benchmarks, including language modeling, common-sense reasoning, in-context retrieval, length extrapolation, and long-context understanding. We further enhance performance by developing hybrid architectures that combine Gated DeltaNet layers with sliding window attention or Mamba2 layers, achieving both improved training efficiency and superior task performance.