Apple lässt große Modelle lernen, faul zu sein: den ersten Token schneller ausspucken und die Genauigkeit beibehalten

偷懒才能更好地工作。

Llama 3.1 刚刚发布，你是否已经尝试了呢？就算你的个人计算机是最近的顶尖配置，运行其中最小的 8B 版本可能也依然会有明显延迟。为了提升模型的推理效率，研究者想出了多种多样的方法，但其中很多都会让模型牺牲一些准确度。

近日，苹果和 Meta AI 的一个研究团队提出了一种新方法，可在保证准确度不明显下降的同时，将 Llama 2 预填充阶段的推理速度提升到原来的 2 倍以上，这或许能为 Llama 3.1 的加速提供一些启发。他们把这种方法称为 LazyLLM，即懒惰大型语言模型。

• 论文标题：LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

• 论文地址：https://arxiv.org/abs/2407.14057

那么他们是怎么让 LLM 偷懒的呢？要理解他们的方法，我们首先需要知道标准的基于 prompt 的 LLM 推理过程是怎样的。简单来说，该过程分为两个阶段：预填充和解码，如图 1 所示。

在预填充阶段，模型计算和保存 prompt 中每个 token 的 KV 缓存，并预测首个 token。我们将预填充阶段所耗费的时间称为「首个 token 时间（TTFT）」。

预填充阶段之后是解码阶段。在这个阶段，模型再次使用缓存的 KV 来迭代式地解码下一个 token，直到满足停止标准。

在预填充阶段，所有 Transformer 层都会使用 prompt 中的所有 token。当 prompt 较长时，TTFT 可能很慢，因为当前最佳的基于 Transformer 的 LLM 既深又宽，并且计算注意力的成本会随 prompt 中 token 数量而呈二次增长。举个例子，Llama 2（7B 版本）堆叠了 32 层 Transformer，模型维度为 4096。在这种情况下，TTFT 需要的 walltime 是每个后续解码步骤的 21 倍，在 LongBench 基准上这些时间大约占用了总生成时间的 23%。

因此，要让 LLM 推理高效进行，优化 TTFT 是非常关键的步骤。

尽管 LLM 推理优化方面是一个活跃的研究领域，但很多方法关注的重心都是提升解码阶段的推理速度。研究者很少关注 TTFT 的改进。一些基于压缩的研究成果可通过减少 LLM 的大小隐式地提升 TTFT。

另一个研究方向是在静态的 Transformer 架构下实现对 TTFT 的改进。对于这个研究方向，很自然会引出一个问题：在生成首个 token 时，所有 prompt token 都必不可少吗？

图 2 给出了在 LongBench 基准上的 LLM 分析结果。

可以看到，对于首个生成的 token，输入 token 的注意力分数非常稀疏，这说明输入 prompt 中的许多 token 是多余的，就算移除也不会影响到下一 token 预测。这一观察正是该团队提出 LazyLLM 的基础。

LazyLLM 的优势包括适用范围广、无需训练、效果好。图 3 对比了标准 LLM 与 LazyLLM。

LazyLLM

图 4 展示了 LazyLLM 的整体框架。

从完整上下文开始，LazyLLM 会逐渐对 token 进行剪枝，从而逐渐减少得到最终模型所使用的计算数量。请注意，LazyLLM 允许模型在不同的生成步骤选取不同的 token 子集，即便它们中的一些可能在之前的步骤中被剪枝了。相比于静态剪枝（一次性对所有 token 进行剪枝），动态剪枝会在每个生成步骤对下一 token 预测进行优化，这有助于维持模型的性能表现。

渐进式 token 剪枝

之前也有一些研究成功使用过 token 剪枝来优化 LLM 推理。但是，这些方法需要积累预测前几个 token 的完整注意力图，以便在剪枝开始之前分析 prompt token 的重要性。也因此，它们不适合用于降低 TTFT，因为它们在预填充阶段仍需要计算所有 KV 缓存。

Im Vergleich dazu ist LazyLLM „sehr faul“ und berechnet ab der ersten Iteration der Inferenz (Vorfüllschritt) nur Token, die für die Vorhersage des nächsten Tokens wichtig sind.

In der ersten Iterationsrunde besteht ein Hauptproblem darin, die Bedeutung jedes Tokens zu bestimmen. Inspiriert durch frühere Untersuchungen, die zeigen, dass sich verborgene Token-Zustände entwickeln, wenn sie Transformer-Schichten durchlaufen, besteht die Lösung des Teams darin, bei jedem Generierungsschritt eine schichtweise Token-Bereinigung durchzuführen. Insbesondere verwenden sie die Aufmerksamkeitskarte jeder Schicht, um die Bedeutung des Eingabe-Tokens für das vorherzusagende Token zu bestimmen.

Nach der Berechnung des Konfidenzwerts des Tokens besteht ein weiteres schwieriges Problem darin, den Schwellenwert für die Bereinigung des Tokens zu bestimmen.

Insbesondere für verschiedene Ebenen und verschiedene Aufgaben kann sich dieser Schwellenwert ändern, wenn sich der Aufmerksamkeitswert ändert. Die Lösung des Teams besteht darin, die Top-K-Perzentil-Auswahlstrategie zu verwenden. Insbesondere wenn der Konfidenzwert eines Tokens kleiner als das k-te Perzentil der Eingabe-Tokens ist, wird es bereinigt. Sobald ein Token bereinigt wird, nimmt es nicht mehr an der Berechnung aller nachfolgenden Schichten teil.

Mit anderen Worten: Die von nachfolgenden Schichten verwendeten Token sind eine Teilmenge der von vorherigen Schichten verwendeten Token.

Spätere Experimente zeigen, dass sich auch die Leistung ändert, wenn die Position der Beschneidungsschicht und die Anzahl der beschnittenen Token unterschiedlich sind. Insbesondere für dieselbe Transformer-Schicht nimmt die Leistung des Modells allmählich ab, da durch Beschneiden immer mehr Token entfernt werden.

Sie fanden außerdem heraus, dass beim Beschneiden in späteren Schichten im Vergleich zum Beschneiden in frühen Schichten eine bessere Leistung erzielt wird, was zeigt, dass spätere Schichten weniger empfindlich auf einen Token-Beschnitt reagieren. Um Geschwindigkeit und Genauigkeit besser auszubalancieren, verwendete das Team eine progressive Bereinigung, wie in Abbildung 4 dargestellt, wobei mehr Token in frühen Schichten beibehalten wurden und dann die Anzahl der Token schrittweise reduziert wurde, während sie in spätere Schichten fließen.

Aux Cache (Hilfscache)

In der Vorfüllphase gibt es keinen KV-Cache und jedes Token wird in einem verborgenen Zustand ausgedrückt. Daher kann eine progressive Token-Bereinigung erreicht werden, indem der verborgene Zustand der bereinigten Token entfernt wird. Die Ausweitung des progressiven Token-Prunings auf nachfolgende Decodierungsschritte ist jedoch nicht einfach. Der Grund dafür ist, dass jeder Decodierungsschritt die Aufmerksamkeit mithilfe des KV-Cache berechnet, der in der Vorfüllphase berechnet wurde. Da LazyLLM in der Pre-Population-Phase eine progressive Token-Bereinigung durchführt, wird der KV eines Tokens, der auf einer bestimmten Ebene bereinigt wird, nicht im KV-Cache der nächsten Ebene angezeigt.

Zur Erinnerung: Das LazyLLM-Framework ermöglicht es jedem Generierungsschritt, bei jedem Schritt eine andere Teilmenge von Token aus der vollständigen Eingabe-Token-Sequenz auszuwählen, unabhängig davon, ob sie in vorherigen Schritten bereinigt wurden. Beispielsweise können im nachfolgenden Decodierungsschritt bereinigte Token, die nicht im KV-Cache vorhanden sind, für die Aufmerksamkeitsberechnung erneut ausgewählt werden. In diesem Fall kann das Modell den KV-Cache für diese Token nicht abrufen.

Eine intuitive Lösung besteht darin, diese Token durch den Startpunkt des Transformers zu leiten. Dies führt jedoch zu einer Doppelzählung desselben Tokens und verlangsamt letztendlich die Gesamtgenerierungsgeschwindigkeit.

Um dieses Problem zu lösen, hat das Team zusätzlich zum ursprünglichen KV-Cache einen weiteren Cache eingeführt: Aux Cache (Hilfscache).

Wenn der KV der bereinigten Token (T4 und T7 in Abbildung 4) nicht im KV-Cache nachfolgender Schichten erscheint, werden ihre verborgenen Zustände vom Aux-Cache gespeichert, um sie in nachfolgenden Iterationen abzurufen.

Wie in Abbildung 4 dargestellt, ruft jede Transformer-Schicht bei jedem Decodierungsschritt zunächst den KV-Cache vergangener Token ab (falls vorhanden). Für die Token, die sich nicht im KV-Cache befinden, wird ihr verborgener Zustand direkt aus dem Aux-Cache der vorherigen Ebene abgerufen, ohne dass die vorherige Ebene erneut durchlaufen werden muss. Aux Cache stellt sicher, dass jedes Token höchstens einmal in jeder Transformer-Schicht berechnet wird, und stellt außerdem sicher, dass LazyLLM schneller ist als Standard-LLM in seiner langsamsten Form.

Experimente

Das Team testete diesen „faulen“ neuen Ansatz an zwei großen Sprachmodellen: Llama 2 7B und XGen 7B. Das zum Vergleich verwendete Standard-LLM ist dasselbe öffentlich veröffentlichte, vorab trainierte Checkpoint-Modell ohne zusätzliche Schulung.

Der experimentelle Benchmark ist LongBench, ein Multitasking-Benchmark für das Verständnis langer Inhalte. Der LongBench-Benchmark enthält 16 Datensätze, die 6 Aufgaben abdecken, darunter Fragen und Antworten für einzelne Dokumente, Fragen und Antworten für mehrere Dokumente, Zusammenfassung, Lernen mit wenigen Schüssen, Syntheseaufgaben und Code-Vervollständigung.

Die Bewertungsmetrik ist die Wirksamkeit und Effizienz jeder Methode im Hinblick auf den Kompromiss zwischen TTFT-Beschleunigung und Genauigkeit.

Ergebnisse

Tabelle 1 enthält die TTFT-Beschleunigungs- und Genauigkeitsergebnisse für LazyLLM, Standard-LLM und andere Basismethoden.

In dieser Tabelle bezieht sich die Baseline auf die standardmäßige LLM-Inferenz. Zufälliger Token-Drop bezieht sich auf die zufällige Bereinigung von Token. Statische Token-Bereinigung bezieht sich auf die Durchführung einer einmaligen Bereinigung des Eingabe-Tokens basierend auf der Aufmerksamkeitsmethode der vorherigen Transformer-Schichten in der Vorfüllphase. Bei der Prompt-Komprimierung handelt es sich um die Methode der Prompt-Komprimierung, die LLM verwendet, um Redundanz im Eingabekontext zu entfernen.

Wie aus Tabelle 1 ersichtlich ist, ist LazyLLM bei der TTFT-Beschleunigung umfassend überlegen, während der Rückgang der Genauigkeit grundsätzlich vernachlässigbar ist. Es ist zu beachten, dass die Verwendung von LLM zum Komprimieren von Eingabeaufforderungen rechenintensiv ist. Daher ist die tatsächliche TTFT länger als bei Standard-LLM, obwohl die Prompt-Komprimierung die Inferenz schneller macht.

Auswirkungen auf die Gesamtgenerierungsgeschwindigkeit

Um die Auswirkungen der neuen Methode auf die Gesamtgenerierungsgeschwindigkeit zu bewerten, analysierte das Team den Prozentsatz der in der Berechnung verwendeten Prompt-Tokens und die Generierungsbeschleunigung, siehe Tabelle 2.

Man erkennt, dass der Anteil der in LazyLLM-Berechnungen verwendeten Token immer unter 100 % liegt, was zeigt, dass LazyLLM am Ende der Generierung nicht alle Token im Prompt aufgebraucht hat, sondern theoretisch das Modell kann alle Token verwenden. Dies kann den gesamten Generierungsprozess für verschiedene Aufgaben zusätzlich beschleunigen.

Abwurfrate verschiedener Schichten

Das Team analysierte auch die Auswirkungen der Position der Beschneidungsschicht und der Anzahl der beschnittenen Token. Die Ergebnisse sind in Abbildung 6 dargestellt.

Es ist ersichtlich, dass bei der Bereinigung auf derselben Transformer-Ebene die Leistung des Modells umso schlechter ist, je weniger Token übrig sind. Dies steht auch im Einklang mit unserem intuitiven Verständnis. Darüber hinaus führt die Bereinigung in späteren Schichten im Vergleich zur Durchführung der Bereinigung in früheren Transformer-Schichten zu einer besseren Leistung, was zeigt, dass spätere Schichten weniger empfindlich auf die Token-Bereinigung reagieren.

Basierend auf diesen Beobachtungen kann man sagen, dass die Wirkung des progressiven Token Pruning nachgewiesen ist.

Progressives KV-Wachstum

Schließlich versuchte das Team auch, die Interna des Modells mithilfe der Token-Pruning-Logik zu verstehen. Insbesondere möchten sie den kumulierten Anteil der verwendeten Prompt-Tokens und den entsprechenden Anteil der nicht verwendeten kennen. Diese „kumulative Token-Nutzung“ kann äquivalent als die KV-Cache-Größe bei jedem Schritt definiert werden. Abbildung 7 zeigt diese kumulative Prompt-Token-Nutzung für jede Phase von LazyLLM.

Dieses Ergebnis stützt die Hypothese, dass das Modell niemals viele Token auswählen wird (auch wenn das Modell theoretisch alle Token in der Eingabeaufforderung verwenden kann).

Angesichts der Tatsache, dass das Modell die Genauigkeit der Ausführung der Aufgabe weiterhin aufrechterhalten kann, es kann Schlussfolgerung: Das Modell kann Token effektiv verwerfen, ohne die Ausgabequalität zu beeinträchtigen

.

Das obige ist der detaillierte Inhalt vonApple lässt große Modelle lernen, faul zu sein: den ersten Token schneller ausspucken und die Genauigkeit beibehalten. Für weitere Informationen folgen Sie bitte anderen verwandten Artikeln auf der PHP chinesischen Website!