Reti Delta con cancelli: Migliorare Mamba2 con la regola Delta

Gated Delta Networks: Improving Mamba2 with Delta Rule

December 9, 2024

Autori: Songlin Yang, Jan Kautz, Ali Hatamizadeh

cs.AI

Abstract

I trasformatori lineari hanno attirato l'attenzione come alternative efficienti ai trasformatori standard, ma le loro prestazioni nei compiti di recupero e contesti lunghi sono state limitate. Per affrontare tali limitazioni, lavori recenti hanno esplorato due meccanismi distinti: il gating per il controllo adattivo della memoria e la regola di aggiornamento delta per modifiche precise della memoria. Osserviamo che questi meccanismi sono complementari: il gating consente una rapida cancellazione della memoria mentre la regola delta facilita gli aggiornamenti mirati. Sfruttando questa intuizione, introduciamo la regola delta con gating e sviluppiamo un algoritmo di addestramento parallelo ottimizzato per hardware moderno. La nostra architettura proposta, Gated DeltaNet, supera costantemente modelli esistenti come Mamba2 e DeltaNet su diversi benchmark, inclusi modellizzazione del linguaggio, ragionamento di senso comune, recupero in contesto, estrapolazione della lunghezza e comprensione di contesti lunghi. Miglioriamo ulteriormente le prestazioni sviluppando architetture ibride che combinano strati di Gated DeltaNet con attenzione a finestra scorrevole o strati di Mamba2, ottenendo sia un'efficienza di addestramento migliorata che prestazioni superiori nei compiti.

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.