FedLEKE: Federated Locate-then-Edit Knowledge Editing for Multi-Client Collaboration

生成日:

FedLEKE: Federated Locate-then-Edit Knowledge Editing for Multi-Client Collaboration

Abstract（日本語訳）

Locate-then-Edit Knowledge Editing（LEKE）は、大規模言語モデル（LLM）を全面的に再学習せずに更新するための主要な技術である。しかし、既存手法は単一ユーザーの設定を仮定しており、病院や金融機関のような分散した組織が、重なりをもつ知識をそれぞれ更新する現実のマルチクライアント環境では効率が低下する。この場合、mediator knowledge vector（MKV）の計算が重複し、プライバシー上の懸念も生じる。これらの課題に対して、本論文は Federated Locate-then-Edit Knowledge Editing（FedLEKE）という新しいタスクを導入する。これは、複数のクライアントがプライバシーを保ち、計算量を抑えながら、協調して LEKE を実行できるようにするものである。そのために、本論文は MKV の選択と再利用を最適化する二段階のフレームワーク FedEdit を提案する。第一段階では、各クライアントがローカルに LEKE を適用し、計算された MKV をアップロードする。第二段階では、サーバ上の MKV 共有だけに依存するのではなく、FedLEKE において各クライアントがコサイン類似度にもとづいて関連する MKV を取得できるようにし、知識の再編集を可能にして、重複計算を減らす。二つのベンチマークデータセットでの実験結果は、FedEdit が非 federated な LEKE の性能の 96% 超を保持しつつ、FedAvg ベースのベースラインを約 2 倍上回ることを示した。さらに、FedLEKE タスクにおいては、本論文の FedEdit フレームワークの下で、MEMIT が PMET より一貫した性能を示すことも分かった。コードは https://github.com/zongkaiz/FedLEKE で公開されている。

論文の面白いところ

LLM の知識編集は、モデル内の特定の事実を置き換える技術として研究されてきた。たとえば「ある人物の国籍」や「ある作品の作者」のような知識を、モデル全体の再学習なしに変更する。従来の Locate-then-Edit Knowledge Editing（LEKE）は、この処理を単一のモデル所有者が行うことを想定する。本論文の着眼点は、同じような編集を複数の組織が別々に繰り返す状況である。病院や金融機関のようにデータを外へ出しにくい組織では、知識更新の必要は共有されても、元データをまとめることは難しい。そこで、編集そのものではなく、編集で得られる中間的なベクトルを共有し、近い編集に再利用するという設計が置かれている。

この発想は、federated learning を知識編集にそのまま持ち込むだけではない。通常の FedAvg は、各クライアントの更新差分を平均してグローバルモデルへ反映する。しかし知識編集では、編集対象の事実が局所的で、単純平均が意味を壊すことがある。本文のケーススタディでも、MEMITAvg は国籍の編集で不自然な出力を生じる一方、FedMEMIT は目的の国名を出している。共有するものを重み差分ではなく MKV にする点が、この論文の中心である。派手な新モデルではないが、知識編集を複数組織の運用に近づける問題設定として読みやすい。

問題設定

LEKE は、LLM の内部パラメータのうち事実知識に関係する部分を探し、その値を書き換える方法である。ROME、MEMIT、PMET などの系譜では、入力プロンプトから対象知識に対応する隠れ状態や key-value 的な表現を求め、Transformer の MLP 層などに編集を加える。本論文が扱う FedLEKE では、複数のクライアントがそれぞれローカルな事実集合を持つ。各クライアントは生データを共有せず、自身のモデルに対して編集を行う。目標は、重複する知識更新を何度も勾配降下で計算しないことと、各クライアントに不要な編集を押し込まないことである。

難しさは二つある。第一に、共有すべき中間表現をどう定めるかである。重み差分をそのまま共有すると、編集対象が異なる場合に平均化の副作用が出やすい。第二に、サーバ上に蓄積された中間表現のうち、どれを各クライアントが使うべきかを選ぶ必要がある。無関係な編集を再利用すると、正しい知識を変えてしまう。逆に選択が保守的すぎると、再利用の利益が小さくなる。FedLEKE は、この選択と再編集を含むマルチクライアントの知識編集タスクとして定義される。

提案手法

提案フレームワーク FedEdit は、編集と再編集の二段階からなる。第一段階では、各クライアントがローカルに既存の LEKE 手法を実行する。本論文の実験では、主に MEMIT と PMET をこの基礎編集器として用いる。編集の過程で、対象知識に対応する最適化済み隠れ状態 z と、層ごとの key 表現 k が得られる。著者らは、この組を mediator knowledge vector（MKV）として扱う。zsRE 上の分析では、テキスト中の知識と z ベクトルのコサイン類似度に強い正の相関があり、Pearson 係数は 0.74 と報告されている。

第二段階では、各クライアントが MKV をサーバへアップロードし、所定の time slot にサーバ上の MKV を参照する。クライアントは、自身の MKV 群とサーバ上の候補 MKV 群のコサイン類似度を計算する。候補のうち、類似度が閾値 α を超えるものが十分に多い場合、その候補は再編集に使われる。条件を満たした MKV だけを用いて、クライアントはすでに編集済みのモデルをさらに更新する。この仕組みは、同じ領域の近い知識を持つクライアント同士では再利用を許し、遠い知識の混入を抑えるためのものである。中心にあるのは、全クライアントの知識を一つに集めることではなく、関連する編集の痕跡を限定的に使い回すことである。

結果

実験は、COUNTERFACT と zsRE の二つの知識編集ベンチマークを、FedLEKE の設定に組み替えて行われた。バックボーンには GPT-J 6B と GPT-NeoX 20B が使われ、クライアント数は主設定で 8、編集数は約 10,000、time slot は 10 とされている。評価は、編集した事実が正しく出るかを測る Efficacy、言い換えプロンプトでも成功するかを測る Generalization または Paraphrase、無関係な知識が保たれるかを測る Specificity を中心に行われた。COUNTERFACT では Fluency と Consistency も加えられている。

COUNTERFACT の GPT-J 6B では、FedMEMIT が非 federated な MEMIT の Score の 99.2%、FedPMET が PMET の 97.0% を保持した。単純な FedAvg 型の MEMITAvg と PMETAvg は、それぞれ 73.3% と 41.6% にとどまる。GPT-NeoX 20B でも、FedMEMIT と FedPMET は MEMITAvg、PMETAvg を大きく上回った。zsRE では FedMEMIT が MEMIT の 99.6% の Score を示し、Efficacy と Generalization でも近い値を保った。一方、FedPMET は zsRE の 10K 編集では Score が 82.4% まで落ち、MEMIT の方が安定していた。アブレーションでは、再編集条件を外すと各規模で Score が下がり、関連する MKV だけを選ぶ処理が性能に寄与していることが示された。

具体例

ある病院 A のローカルモデルが、「患者 Tom の治療終了日は 1 月 12 日である」と答える状態にあるとする。新しい記録では、正しい治療終了日が 2 月 15 日に変わっている。病院 A は生のカルテを外へ出さず、自分の環境で LEKE を実行し、「Tom finished his treatment on ...」というプロンプトに対して「February 15th」を出すようにモデルを編集する。このとき得られた z と key の組が MKV としてサーバに置かれる。

別の病院 B でも、同じ治療イベントに近い形式の更新が発生したとする。病院 B はサーバ上の MKV を無差別に使うのではなく、自分の編集で得た MKV と候補 MKV のコサイン類似度を比べる。十分に近いものだけが選ばれ、再編集に使われる。期待される出力は、言い換えられた質問、たとえば「When did Tom complete treatment?」に対しても 2 月 15 日を返すことである。間違えやすい点は、似ていない医療記録の更新まで取り込むと、別の患者や別のイベントの日付を壊すことである。FedEdit の再編集条件は、この混入を避けながら、近い知識更新の計算を繰り返さないために置かれている。