KV 預測以改善首次標記的時間

使用變壓器語言模型的推理過程

使用變壓器（transformer）語言模型進行推理時，首先需要處理提示（prompt）。在這個步驟中，模型會生成第一個輸出標記，並儲存未來生成步驟所需的KV快取（KV cache）。這個提示處理的過程可能會很耗費計算資源，尤其是當提示的長度或批次大小增加時，對於擁有十億參數的模型，在邊緣設備上可能需要花費十幾秒甚至更長的時間。這會影響使用者的體驗，因為模型的輸出會有明顯的延遲。

為了減少生成預訓練模型的第一個輸出所需的時間（稱為“第一次標記時間”，或TTFT），我們提出了一種新方法，稱為KV預測（KV Prediction）。在這個方法中，我們使用一個小型輔助模型來處理提示，並生成基礎模型所需的KV快取的近似值。這個近似的KV快取隨後會與基礎模型一起使用，進行自回歸生成，而不需要再次查詢輔助模型。

我們的研究顯示，這種方法在效率和準確性之間達到了最佳的平衡，與基準相比，我們在TriviaQA上展示了15%到50%的相對準確性提升，這是在不同的TTFT FLOPs預算下的結果。我們還在固定的TTFT FLOPs預算下，展示了在HumanEval的Python代碼補全上，準確性提升高達30%。此外，我們在蘋果M2 Pro CPU（Apple M2 Pro CPU）上進行了基準測試，證明我們在FLOPs上的改進也轉化為硬體上的TTFT加速。我們在這裡釋出我們的代碼。