RNN效率媲美Transformer，谷歌新架构两连发：同等规模强于Mamba

去年 12 月，新架构 Mamba 引爆了 AI 圈，向屹立不倒的 Transformer 发起了挑战。如今，谷歌 DeepMind「Hawk 」和「Griffin 」的推出为 AI 圈提供了新的选择。

这一次，谷歌 DeepMind 在基础模型方面又有了新动作。

我们知道，循环神经网络（RNN）在深度学习和自然语言处理研究的早期发挥了核心作用，并在许多应用中取得了实功，包括谷歌第一个端到端机器翻译系统。不过近年来，深度学习和 NLP 都以 Transformer 架构为主，该架构融合了多层感知器（MLP）和多头注意力（MHA）。

Transformer 已经在实践中实现了比 RNN 更好的性能，并且在利用现代硬件方面也非常高效。基于 Transformer 的大语言模型在从网络收集的海量数据集上进行训练，取得了显着的成功。

纵然取得了很大的成功，但 Transformer 架构仍有不足之处，比如由于全局注意力的二次复杂性，Transformer 很难有效地扩展到长序列。此外，键值（KV）缓存随序列长度线性增长，导致 Transformer 在推理过程中变慢。这时，循环语言模型成为一种替代方案，它们可以将整个序列压缩为固定大小的隐藏状态，并迭代更新。但若想取代 Transformer，新的 RNN 模型不仅必须在扩展上表现出相当的性能，而且必须实现类似的硬件效率。

在谷歌DeepMind 近日的一篇论文中，研究者提出了RG-LRU 层，它是一种新颖的门控线性循环层，并围绕它设计了一个新的循环块来取代多查询注意力（MQA）。

他们使用该循环块构建了两个新的模型，一个是混合了MLP 和循环块的模型Hawk，另一个是混合了MLP 与循环块、局部注意力的模型Griffin 。

• 论文标题：Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models
• 论文链接：https://arxiv.org/pdf/2402.19427.pdf

研究者表示，Hawk 和Griffin 在held-out 损失和训练FLOPs 之间表现出了幂律缩放，最高可以达到7B 参数，正如之前在Transformers 中观察到的那样。其中 Griffin 在所有模型规模上实现了比强大 Transformer 基线略低的 held-out 损失。

研究者针对一系列模型规模、在300B tokens 上对Hawk 和Griffin 进行了过度训练，结果显示，Hawk-3B 在下游任务的性能上超越了Mamba-3B，尽管训练的tokens 数量只有后者的一半。 Griffin-7B 和 Griffin-14B 的性能与 Llama-2 相当，尽管训练的 tokens 数量只有后者的 1/7。

此外，Hawk 和 Griffin 在 TPU-v3 上达到了与 Transformers 相当的训练效率。由于对角 RNN 层受内存限制，研究者使用了 RG-LRU 层的内核来实现这一点。

同时在推理过程中，Hawk 和 Griffin 都实现比 MQA Transformer 更高的吞吐量，并在采样长序列时实现更低的延迟。当评估的序列比训练中观察到的更长时，Griffin 的表现比 Transformers 更好，并且可以有效地从训练数据中学习复制和检索任务。不过当在未经微调的情况下在复制和精确检索任务上评估预训练模型时，Hawk 和 Griffin 的表现不如 Transformers。

共同一作、DeepMind 研究科学家 Aleksandar Botev 表示，混合了门控线性循环和局部注意力的模型 Griffin 保留了 RNN 的所有高效优势和 Transformer 的表达能力，最高可以扩展到 14B 参数规模。

                               ^{来源：https://twitter.com/botev_mg/status/1763489634082795780}

Architecture du modèle Griffin

Griffin Tous les modèles contiennent les composants suivants : (i) un bloc résiduel, (ii) un bloc MLP, (iii) un bloc de mélange temporel. (i) et (ii) sont les mêmes pour tous les modèles, mais il existe trois blocs de mélange temporel : l'attention multi-requêtes globale (MQA), le MQA local (fenêtre glissante) et le bloc récurrent proposé dans cet article. Dans le cadre du bloc récurrent, les chercheurs ont utilisé une unité récurrente linéaire vraiment fermée (RG-LRU), une nouvelle couche récurrente inspirée des unités récurrentes linéaires.

Comme le montre la figure 2(a), le bloc résiduel définit la structure globale du modèle Griffin, qui s'inspire du pré-normeTransformer. Après avoir intégré la séquence d'entrée, nous la passons à travers des blocs comme ? (? représente la profondeur du modèle), puis appliquons RMSNorm pour générer les activations finales. Pour calculer les probabilités des jetons, une dernière couche linéaire est appliquée, suivie de softmax. Les poids de cette couche sont partagés avec la couche d’intégration d’entrée.

Modèle récurrent, efficacité de mise à l'échelle comparable à celle de Transformer

La recherche sur la mise à l'échelle fournit des informations importantes sur la manière d'ajuster les hyperparamètres du modèle et son comportement lors de la mise à l'échelle.

Les chercheurs ont défini les modèles évalués dans cette étude, ont fourni des courbes de mise à l'échelle jusqu'aux paramètres 7B et au-delà et ont évalué les performances du modèle sur les tâches en aval.

Ils ont considéré 3 familles de modèles : (1) baseline MQA-Transformer ; (2) Hawk : un modèle RNN pur (3) Griffin : un modèle hybride qui mélange des blocs récurrents avec une attention locale. Les hyperparamètres clés du modèle pour les modèles de différentes tailles sont définis à l'annexe C.

L'architecture Hawk utilise le même motif résiduel et le même bloc MLP que la ligne de base de Transformer, mais les chercheurs ont utilisé un bloc récurrent avec une couche RG-LRU comme bloc de mélange temporel au lieu de MQA. Ils ont élargi la largeur du bloc de boucle d'un facteur d'environ 4/3 (c'est-à-dire ?_??? ≈4?/3) pour correspondre à peu près au nombre de paramètres du bloc MHA lorsque les deux utilisent la même dimension de modèle ?.

Griffon. Le principal avantage des blocs récurrents par rapport à l'attention globale est qu'ils utilisent une taille d'état fixe pour résumer les séquences, alors que la taille du cache KV de MQA augmente proportionnellement à la longueur de la séquence. L'attention locale a les mêmes propriétés, et le mélange de blocs récurrents avec l'attention locale préserve cet avantage. Les chercheurs ont découvert que cette combinaison était extrêmement efficace, car l’attention locale peut modéliser avec précision le passé récent, tandis que les couches récurrentes peuvent transmettre des informations sur de longues séquences.

Griffin utilise le même modèle résiduel et les mêmes blocs MLP que la ligne de base de Transformer. Mais contrairement à la base de référence MQA Transformer et au modèle Hawk, Griffin utilise un mélange de blocs de boucle et de blocs MQA. Plus précisément, nous adoptons une structure hiérarchique qui alterne deux blocs résiduels avec un bloc récurrent puis un bloc d'attention local (MQA). Sauf indication contraire, la taille de la fenêtre d’attention locale est fixée à 1 024 jetons.

Les principaux résultats de mise à l'échelle sont présentés dans la figure 1 (a). Les trois familles de modèles ont été formées sur des tailles de modèle allant de 100 millions à 7 milliards de paramètres, bien que Griffin propose une version à 14 milliards de paramètres. Les résultats de l'évaluation de

sur les tâches en aval sont présentés dans le tableau 1 :

Hawk et Griffin ont tous les deux très bien joué. Le tableau ci-dessus rend compte de la précision normalisée des fonctionnalités pour MMLU, HellaSwag, PIQA, ARC-E et ARC-C, tout en rapportant la précision absolue et les scores partiels pour WinoGrande. À mesure que la taille du modèle augmente, les performances de Hawk s'améliorent également de manière significative et Hawk-3B fonctionne mieux que Mamba-3B dans les tâches en aval, bien que le nombre de jetons qu'il entraîne ne soit que la moitié de celui de Mamba-3B. Griffin-3B fonctionne nettement mieux que Mamba-3B, et Griffin-7B et Griffin-14B ont des performances comparables à Llama-2, bien qu'ils soient entraînés avec près de 7 fois moins de jetons. Hawk est comparable à la référence MQA Transformer, tandis que Griffin la surpasse.

Entraînez efficacement le modèle de boucle côté appareil

Lors du développement et de l'extension du modèle, les chercheurs ont rencontré deux défis d'ingénierie majeurs. Tout d’abord, comment partager efficacement les modèles de traitement sur plusieurs appareils. Deuxièmement, comment mettre en œuvre efficacement des boucles linéaires pour maximiser l'efficacité de la formation TPU. Cet article aborde ces deux défis et fournit ensuite une comparaison empirique de la vitesse d'entraînement des lignes de base Griffin et MQA.

Les chercheurs ont comparé les vitesses d'entraînement de différentes tailles de modèles et longueurs de séquence pour étudier les avantages informatiques du modèle dans cet article pendant le processus d'entraînement. Le nombre total de jetons par lot reste fixe pour chaque taille de modèle, ce qui signifie qu'à mesure que la longueur de la séquence augmente, le nombre de séquences diminue proportionnellement.

La figure 3 représente la durée d'exécution relative du modèle Griffin par rapport au modèle de base MQA à 2048 longueurs de séquence.

Vitesse d'inférence

L'inférence de LLM se compose de deux étapes. La phase de « pré-remplissage » consiste à recevoir et à traiter les invites. Cette étape effectue en fait une passe avant sur le modèle. Étant donné que les invites peuvent être traitées en parallèle tout au long de la séquence, la plupart des opérations de modèle à ce stade sont liées au calcul. Par conséquent, nous nous attendons à ce que la vitesse relative des transformateurs et des modèles de boucle dans l'étape de pré-remplissage soit la même que celles évoquées précédemment. pendant l’entraînement étaient similaires.

Après le pré-remplissage, vient l'étape de décodage, au cours de laquelle le chercheur extrait de manière autorégressive les jetons du modèle. Comme indiqué ci-dessous, en particulier pour les séquences plus longues, où le cache clé-valeur (KV) utilisé pour l'attention devient volumineux, le modèle récurrent a une latence plus faible et un débit plus élevé dans l'étape de décodage.

Il y a deux mesures principales à prendre en compte lors de l'évaluation de la vitesse d'inférence. Le premier est la latence, qui mesure le temps nécessaire pour générer un nombre spécifié de jetons pour une taille de lot spécifique. Le second est le débit, qui mesure le nombre maximum de jetons pouvant être générés par seconde lors de l'échantillonnage d'un nombre spécifié de jetons sur un seul appareil. Étant donné que le débit est calculé comme le nombre de jetons échantillonnés multiplié par la taille du lot divisé par la latence, vous pouvez augmenter le débit en réduisant la latence ou en réduisant l'utilisation de la mémoire pour utiliser une taille de lot plus grande sur l'appareil. La prise en compte de la latence est utile pour les applications temps réel qui nécessitent des temps de réponse rapides. Le débit mérite également d'être pris en compte car il nous indique le nombre maximum de jetons pouvant être échantillonnés à partir d'un modèle particulier dans un temps donné. Cette propriété est intéressante lorsque l'on considère d'autres applications linguistiques, telles que l'apprentissage par renforcement basé sur la rétroaction humaine (RLHF) ou la sortie d'un modèle de langage de notation (comme cela se fait dans AlphaCode), car être capable de générer un grand nombre de jetons dans un temps donné est une option intéressante. fonctionnalité.

Ici, les chercheurs ont étudié les résultats d'inférence du modèle avec le paramètre 1B. En termes de lignes de base, ils sont comparés au transformateur MQA, qui est nettement plus rapide lors de l'inférence que le transformateur MHA standard couramment utilisé dans la littérature. Les modèles comparés par les chercheurs sont : i) le convertisseur MQA, ii) Hawk et iii) Griffin. Pour comparer différents modèles, nous rapportons la latence et le débit.

Comme le montre la figure 4, les chercheurs ont comparé la latence du modèle avec une taille de lot de 16, un pré-remplissage vide et un pré-remplissage de 4096 jetons.

La figure 1(b) compare le débit maximum (jetons/seconde) des mêmes modèles lors de l'échantillonnage de 512, 1024, 2048 et 4196 jetons respectivement après des indices vides.

Modélisation de contextes longs

Cet article explore également l'efficacité de Hawk et Griffin en utilisant des contextes plus longs pour améliorer les prédictions des prochains jetons et étudie leur capacité à extrapoler pendant l'inférence. Les performances de Griffin sur des tâches nécessitant des capacités de copie et de récupération sont également explorées, à la fois dans des modèles formés à de telles tâches et lorsque ces capacités sont testées à l'aide de modèles de langage pré-entraînés.

À partir du graphique sur le côté gauche de la figure 5, on peut observer que dans une certaine plage de longueur maximale, Hawk et Griffin peuvent tous deux améliorer la capacité de prédiction du prochain jeton dans un contexte plus long, et ils sont globalement capables pour déduire des séquences plus longues (au moins 4 fois) que lors de l'entraînement. Griffin, en particulier, fonctionne très bien en raisonnement même en utilisant RoPE dans la couche d'attention locale.

Comme le montre la figure 6, dans la tâche de copie sélective, les 3 modèles peuvent parfaitement accomplir la tâche. En comparant la vitesse d'apprentissage sur cette tâche, Hawk est nettement plus lent que Transformer, ce qui est similaire aux observations de Jelassi et al (2024) qui ont constaté que Mamba apprenait beaucoup plus lentement sur une tâche similaire. Il est intéressant de noter que même si Griffin n'utilise qu'une couche d'attention locale, sa vitesse d'apprentissage est à peine ralentie et est comparable à celle de Transformer.

Pour plus de détails, veuillez lire l'article original.

以上是RNN效率媲美Transformer，谷歌新架构两连发：同等规模强于Mamba的详细内容。更多信息请关注PHP中文网其他相关文章！