MOSAIC: Composable Safety Alignment with Modular Control Tokens

本論文は、凍結された基盤モデルに対して学習可能な制御トークンを導入し、文脈に応じて安全制約を柔軟に組み合わせることで、過剰な拒否を抑制しつつモデルの有用性を維持する「MOSAIC」と呼ばれるモジュール型安全アライメントフレームワークを提案しています。

要約

🏠 今までの問題点：「硬くて重い鍵」か「頼りないメモ」

AI に安全な回答をさせるための方法は、これまで主に 2 つありました。しかし、どちらも完璧ではありませんでした。

1. モデル自体を修正する方法（パラメータレベルのアライメント）

   • 例え： AI の頭脳そのものを「安全な性格」に書き換えてしまうこと。
   • 問題点： 一度書き換えると、その性格は固定されてしまいます。例えば、「子供向けには危険なことを教えるな」と設定した AI を、後で「大人向けには教えていい」と変えようとしても、頭脳を全部書き直す必要があり、とても大変で、他の能力（例えば料理のレシピを教える力）まで壊してしまうリスクがあります。

2. 指示文でお願いする方法（プロンプトベース）

   • 例え： AI に「ねえ、賭け事の話はしないでね」という長いメモを貼り付けておくこと。
   • 問題点： AI はこのメモを「ただのお願い」と捉えてしまい、時々聞き入れ忘れたり、メモが長すぎて本来の質問に集中できなくなったりします。また、「お酒のレシピ」と「ギャンブル」を同時に禁止したい場合、メモがごちゃごちゃになって混乱します。

---

✨ 新しい解決策：MOSAIC（モザイク）

この論文が提案するMOSAICは、**「AI の頭脳は変えずに、小さな『魔法のタグ』を貼り付ける」**というアイデアです。

🧩 1. 小さな「制御トークン」＝ 魔法のタグ

MOSAIC は、AI の頭脳（ベースモデル）を一切変えません。代わりに、**「ギャンブル禁止タグ」や「お酒禁止タグ」**のような、とても小さなデジタルのタグ（制御トークン）を学習させます。

• どう使う？
  • ユーザーが「カジノの勝ち方」を聞いたら、**「ギャンブル禁止タグ」**を AI の質問の前に貼り付けます。AI はそのタグを見て「あ、これは禁止事項だ」と判断し、断ります。
  • ユーザーが「モヒトの作り方」を聞いたら、**「お酒禁止タグ」**を貼り付けます。
  • もし「ギャンブルと酒」の両方が問題になる質問なら、2 つのタグを同時に貼り付けます。

このタグは、AI の頭脳そのものではなく、**「着脱可能なアクセサリー」**のようなものです。必要な時だけつけ、不要な時は外せます。

🎨 2. 「モザイク」のように組み合わせる

これが「MOSAIC（モザイク）」と呼ばれる理由です。
それぞれのタグ（安全ルール）は独立して作られていますが、組み合わせて使えます。

• 「子供向けモード」＝ 賭け事タグ ＋ 暴力タグ
• 「大人向けモード」＝ 賭け事タグのみ
• 「特定の国向けモード」＝ 法律タグ ＋ 文化タグ

このように、必要なルールだけをパズルのように組み合わせて、その場その場に合わせた安全対策を瞬時に作れるのが最大の特徴です。

🛡️ 3. 「やりすぎ防止」の仕組み

これまでの方法だと、「安全対策」を強くしすぎると、 harmless（無害）な質問まで「断る」ようになってしまいます（例：「お酒のレシピ」を聞いたら、単に「お酒」が含まれているだけで断ってしまう）。

MOSAIC は、これを防ぐために**「もしタグがなかったらどう答えるか？」**を AI に学習させます。

• 学習方法： 「タグあり」の時の回答と、「タグなし」の時の回答を比べます。
• 効果： 「この質問はタグがなくても大丈夫な内容だから、タグがあっても断らないでね」という**「過剰な拒否」を防ぐ学習**が行われます。これにより、必要な時はしっかり断りつつ、普段の会話には邪魔にならないようになります。

---

📊 結果：どう変わった？

実験では、新しいテストデータを使ってこの方法が試されました。

• 高い安全性： 危険な質問には、ほぼ 100% の確率で「断る」ことができました。
• 低い過剰拒否： 安全な質問を間違って断ってしまうことが、従来の方法に比べて大幅に減りました。
• 柔軟性： 新しい安全ルール（例えば「新しい法律」）ができたとき、AI の頭脳を全部作り直す必要なく、新しいタグを少し追加するだけで対応できました。

🌟 まとめ

MOSAICは、AI の安全対策を「一度きりの硬い設定」から、**「状況に合わせて組み換えられる、柔軟なブロック」**に変える画期的な方法です。

• AI の能力は守る（頭脳を変えない）。
• 必要な時だけ安全にする（タグを貼る）。
• 邪魔にならない（過剰に断らない）。

これにより、子供にも大人にも、国や状況に合わせて、ちょうどいい安全レベルの AI を提供できるようになる未来が描かれています。

技術概要

MOSAIC: 構成可能な安全アライメントのためのモジュラー制御トークン

技術的サマリー（日本語）

本論文は、大規模言語モデル（LLM）の安全アライメントにおける既存手法の限界を克服し、ユーザー、地域、アプリケーションに応じた条件付きかつ構成可能な安全制御を実現する新しいフレームワーク「MOSAIC (Composable Safety Alignment with Modular Control Tokens)」を提案するものです。

1. 背景と課題

LLM の実用化において、安全ポリシーは年齢、法域、専門職、アプリケーションドメインによって異なります（例：成人向けコンテンツは未成年には制限される、ある国で合法なものが他国では禁止されるなど）。しかし、既存のアライメント手法には以下のような根本的な限界がありました。

• パラメータレベルのアライメント（SFT, RLHF など）: 安全行動がモデルの重みと密接に絡みついており、特定の安全ルールを動的に無効化したり、新しいルールを追加したりするためにモデルの再学習（リトレーニング）が必要となります。これはコストが高く、既存の能力との干渉（Catastrophic Interference）を引き起こすリスクがあります。
• プロンプトベースのアライメント: 自然言語の指示で制約を表現しますが、確率的な解釈に依存するため一貫性が低く、複数の制約を組み合わせるとプロンプトが長くなり、コンテキスト容量を圧迫します。

これらの手法は、安全制約を「明示的・再利用可能・構成可能」な表現として扱えておらず、微細な条件制御を効率的に行うことが困難でした。

2. 提案手法：MOSAIC

MOSAIC は、安全アライメントを「表現学習問題」として再定義し、**学習可能な制御トークン（Control Tokens）**を用いて安全制約をモジュール化します。

核心的な仕組み

• 凍結されたバックボーンモデル: 基盤となる LLM のパラメータは固定（Frozen）したまま、入力に追加される小さな制御トークンのみ学習します。
• モジュラーな制御トークン: 各安全カテゴリ（例：ギャンブル、アルコール、暴力など）は、埋め込み空間内の少数の学習可能なトークン（例：$z_c = \{z_{c,1}, ..., z_{c,m}\}$）で符号化されます。
• 構成可能性（Compositionality）: 推論時に、必要な安全カテゴリに対応するトークンを入力先頭に付加（Prepend）することで、そのカテゴリの拒否行動を活性化します。複数のトークンを組み合わせることで、複数の安全ポリシーを同時に適用できます。

学習戦略と最適化

制御トークンの効率的な学習と「過剰拒否（Over-refusal：安全でない質問ではないのに拒否してしまう）」の抑制のために、以下の 2 つの主要な技術を採用しています。

1. 順序ベースのタスクサンプリング（Order-based Task Sampling）:
   • 安全カテゴリの組み合わせ数は指数的に増加するため、すべての組み合わせを網羅的に学習するのは非現実的です。
   • 代わりに、アクティブなカテゴリの数（Order: 1 次、2 次、3 次...）ごとにトレーニング予算を割り当て、組み合わせを効率的にサンプリングします。これにより、高次な組み合わせに対する制御能力を維持しつつ、計算コストを抑制します。
2. 反事実的知識蒸留（Counterfactual Knowledge Distillation）:
   • 安全な質問（Benign queries）に対して制御トークンが誤って拒否を引き起こすのを防ぐため、反事実的 KDを導入します。
   • 制御トークンありの分布と、同じモデルでトークンなし（ベースライン）の分布を比較し、KL 発散を最小化します。これにより、制御トークンは「必要な場合のみ」介入し、安全な質問に対するモデルの本来の振る舞いを保持するように学習されます。

3. 評価と結果

既存のベンチマークは既存の LLM の学習データと重複しており、真の条件付き制御能力を評価できないため、著者らは新しい現実的な評価ベンチマークを構築しました（1,500 のユーザー要求、5 つの安全カテゴリ、既存の安全アライメント済みモデルでは拒否されないが、特定の条件下では拒否すべき内容）。

主要な結果（Llama-3.1-8B などでの実験）

• 防御成功率（DSR）: 提案手法は、すべてのタスク順序（1 次〜4 次カテゴリの組み合わせ）において、ほぼ完璧な防御成功率（99% 以上）を達成しました。
• 過剰拒否率（OR）の大幅な削減:
  • 従来の SFT（教師あり微調整）は DSR は高いものの、過剰拒否が約 6% 発生しました。
  • MOSAIC は、カテゴリあたり 5 つのトークンを使用する場合、Llama-3.1-8B で過剰拒否率を**1.8%**まで低下させました。
  • 高次のタスク（複数のカテゴリを同時に適用）になるほど、過剰拒否がさらに減少する傾向が確認されました。
• 汎用性の維持: MMLU などの汎用タスクにおける性能低下はほとんど見られず、モデルの有用性が保たれていることが確認されました。
• 拡張性: 新しい安全カテゴリを追加する際、既存のトークンを再学習することなく追加のみで対応でき、性能の劣化は最小限に抑えられました。

4. 主な貢献

1. 安全アライメントの概念の再定義: モノリシックなパラメータ変更ではなく、モジュラーな制約活性化としての「構成可能な表現学習問題」として安全制御を再定義しました。
2. MOSAIC フレームワークの提案: 凍結されたモデル上で学習可能な制御トークンを用いることで、条件付き安全制御とカテゴリの漸増的拡張を実現し、構造的な組み合わせ学習と反事実的 KD によって過剰拒否を抑制しました。
3. 現実的な評価ベンチマークの構築: 既にアライメント済みのモデルに対する条件付き安全活性化を評価するための新しいデータセットを提供し、選択的な制約執行と汎用性の両立を精密に評価可能にしました。

5. 意義と結論

MOSAIC は、LLM の安全制御を「静的なモデル特性」から「動的なコンテキスト依存制御」へと転換させる重要なステップです。

• 実用性: 異なるユーザー層や法域に合わせて、モデルの再学習なしで安全ポリシーを柔軟に切り替え・組み合わせることができます。
• 効率性: パラメータの再学習が不要であり、制御トークンのみで軽量な制御が可能です。
• 精度: 過剰拒否を大幅に削減しつつ、必要な場面での堅牢な拒否を実現し、モデルの有用性を損なわないバランスの取れた安全制御を提供します。

本研究は、多様な価値観や文脈に対応する「Pluralistic Alignment（多元的アライメント）」の実現に向けた、技術的かつ実用的な基盤を築くものです。