Attentionスコアの分布類似性を用いた大規模言語モデルの動作効率化および省メモリ化

現在、GPTやLlama,MistralといったLLM(LargeLanguageModel)が開発され、一部はOSSで誰でも利用・改良できるようになっている。これらのモデルは活性化関数やレイヤ正規化の手法、学習データ数の差などはあれど基本的な構造はTransformer-デコーダと変わらず、Attentionヘッドと呼ばれる処理機構を並列に持つTransformerブロックが複数個直列に繋がれたモデルである。基本的なTransformer-デコーダモデルのAttentionヘッド１つあたりのPrefill処理時(初回のプロンプト入力時)のAttention処理の計算は、以下のような式で表される。Output=Softmax푄퐾⊤+Mask√푉(1)푑ここでQ,K,Vはそれぞれクエリ、キー、バリュー行列であり、dはQueryベクトルの次元数、Maskは値が1、大きさが入力長と等しい下三角行列である。KVキャッシュを利用したDecode手法では式1は以下のように置き換えられる。퐾cache←[퐾cache,푘],푉cache←[푉cache,푣](2)Output=Softmax푞퐾⊤cache√푉cache(3)푑これによりDecode時にはq,k,vベクトルのみを計算すればよいため生成速度の向上が期待できる。反面、式2のキャッシュサイズは入力長に応じて線形に増加するため、プロンプトが長大な場合はメモリ(GPUを利用する場合はVRAM)を圧迫する課題が存在している。またQ,K,Vの作成やFeedForward処理の計算量はO(n)であるのに対して、式1,3の計算量は푄퐾푇∈ℝ푁×푁のためO(n2)となる。そのため既存のTransformerモデルでは、入力長に応じてメモリ使用量が線形に増加し、また計算量が入力長の二乗に比例して大きくなる課題が存在している。近年ではLLMの活用方法として、本や映画の脚本の要約・複数書類を参照対象とした質問タスクなどが期待されていて、その要求に対して長大なコンテキスト長で学習したLLMモデルも発表されている。しかし上記のような問題により、長文タスクを解く場合はメモリ量・計算力ともに強大な計算機を所持する必要があり、LLMの長文タスクへの応用の妨げとなっている。