Afgeschermde Delta-netwerken: Verbetering van Mamba2 met de Delta-regel

Gated Delta Networks: Improving Mamba2 with Delta Rule

December 9, 2024

Auteurs: Songlin Yang, Jan Kautz, Ali Hatamizadeh

cs.AI

Samenvatting

Lineaire Transformers hebben aandacht gekregen als efficiënte alternatieven voor standaard Transformers, maar hun prestaties bij ophalen en taken met een lange context zijn beperkt gebleven. Om deze beperkingen aan te pakken, heeft recent onderzoek twee verschillende mechanismen verkend: gating voor adaptieve geheugenregeling en de delta-update regel voor nauwkeurige geheugenwijzigingen. We merken op dat deze mechanismen complementair zijn: gating maakt snelle geheugenwissing mogelijk terwijl de delta-regel gerichte updates vergemakkelijkt. Voortbouwend op deze inzichten introduceren we de gated delta-regel en ontwikkelen we een parallel trainingsalgoritme geoptimaliseerd voor moderne hardware. Onze voorgestelde architectuur, Gated DeltaNet, overtreft consequent bestaande modellen zoals Mamba2 en DeltaNet op meerdere benchmarks, waaronder taalmodellering, gezond verstand redeneren, in-context ophalen, lengte-extrapolatie en begrip van lange context. We verbeteren de prestaties verder door hybride architecturen te ontwikkelen die Gated DeltaNet-lagen combineren met aandacht met schuifvenster of Mamba2-lagen, waardoor zowel de trainings-efficiëntie als de taakprestaties worden verbeterd.

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.