Keras で LIT を使用して Gemma モデルを分析する

1. はじめに

生成 AI プロダクトは比較的新しいため、アプリケーションの動作は以前のソフトウェアよりも大きく異なる場合があります。そのため、使用している ML モデルを調査し、モデルの動作の例を調べて、想定外の動作を調査することが重要です。

Learning Interpretability Tool(LIT)(ウェブサイトGitHub)は、ML モデルのデバッグと分析を行うためのプラットフォームであり、モデルの動作の理由と方法を理解できます。

この Codelab では、LIT を使用して Google の Gemma モデルを最大限に活用する方法を学びます。この Codelab では、解釈可能性の手法であるシーケンスの重要性を使用して、さまざまなプロンプト エンジニアリング アプローチを分析する方法について説明します。

学習目標:

  1. シーケンスの重要性と、モデル分析での使用方法について説明します。
  2. Gemma の LIT を設定して、プロンプト出力とシーケンスの salience を計算する。
  3. LM Salience モジュールでシーケンスの重要性を使用して、プロンプト設計がモデル出力に与える影響を把握する。
  4. LIT でプロンプトの改善案をテストし、その効果を確認する。

注: この Codelab では、Gemma の KerasNLP 実装と、バックエンドに TensorFlow v2 を使用します。GPU カーネルを使用してこのチュートリアルを行うことを強くおすすめします。

LIT UI デモ

2. シーケンスの重要性とモデル分析での使用

Gemma などの Text-to-Text 生成モデルは、トークン化されたテキスト形式の入力シーケンスを受け取り、その入力の一般的なフォローアップまたは完了である新しいトークンを生成します。この生成は 1 つのトークンずつ行われ、モデルが停止条件に達するまで、新しく生成された各トークンを入力と以前の生成結果に(ループで)追加します。たとえば、モデルがシーケンスの終了(EOS)トークンを生成するときや、事前定義された最大長に達したときなどです。

セレンシティ メソッドは、出力のさまざまな部分について、モデルにとってどの入力部分が重要であるかを示す Explainable AI(XAI)手法の一種です。LIT は、さまざまな分類タスクの salience メソッドをサポートしています。これは、一連の入力トークンが予測ラベルに与える影響を説明します。シーケンスの重要度は、これらの方法を Text-to-Text 生成モデルに一般化し、生成されたトークンに先行するトークンが及ぼす影響を説明します。

ここでは、シーケンスの顕著性に Grad L2 Norm メソッドを使用します。このメソッドは、モデルの勾配を分析し、先行する各トークンが出力に与える影響の大きさを提供します。この方法はシンプルで効率的であり、分類などの設定で優れたパフォーマンスが示されています。セリーランス スコアが高いほど、影響力は大きくなります。この方法は、解釈可能性の研究コミュニティでよく理解され、広く使用されているため、LIT で使用されています。

より高度な勾配ベースの salience メソッドには、Grad ⋅ Input統合勾配があります。LIMESHAP などのアブレーションベースの方法もあります。これらの方法はより堅牢ですが、計算コストが大幅に高くなります。さまざまな視覚的特徴量の詳細な比較については、こちらの記事をご覧ください。

注目度の科学的な方法について詳しくは、注目度に関するインタラクティブな入門コンテンツをご覧ください。

3. インポート、環境、その他の設定コード

この Codelab は、新しい Colab で実施することをおすすめします。モデルをメモリに読み込むため、アクセラレータ ランタイムを使用することをおすすめします。ただし、アクセラレータ オプションは時間の経過とともに変化し、制限が適用される点に注意してください。より強力なアクセラレータをご利用になりたい場合は、Colab の有料サブスクリプションをご利用いただけます。また、マシンに適切な GPU がある場合は、ローカル ランタイムを使用することもできます。

注: 次のような警告が表示される場合があります

ERROR: pip's dependency resolver does not currently take into account all the
packages that are installed. This behaviour is the source of the following
dependency conflicts. bigframes 0.21.0 requires scikit-learn>=1.2.2, but you
have scikit-learn 1.0.2 which is incompatible. google-colab 1.0.0 requires
ipython==7.34.0, but you have ipython 8.14.0 which is incompatible.

無視してかまいません。

LIT と Keras NLP をインストールする

この Codelab では、最新バージョンの keras(3)、keras-nlp(0.14)、lit-nlp(1.2)と、ベースモデルをダウンロードするための Kaggle アカウントが必要です。

pip install -q -U 'keras >= 3.0' 'keras-nlp >= 0.14' 'lit-nlp >= 1.2'

Kaggle へのアクセス

Kaggle で認証するには、次のいずれかの方法があります。

  • 認証情報をファイル(~/.kaggle/kaggle.json など)に保存します。
  • KAGGLE_USERNAME 環境変数と KAGGLE_KEY 環境変数を使用する。
  • Google Colab などのインタラクティブな Python 環境で次のコードを実行します。
import kagglehub

kagglehub.login()

詳細については、kagglehub のドキュメントをご覧ください。また、Gemma ライセンス契約に同意してください。

Keras の構成

Keras 3 は、Tensorflow(デフォルト)、PyTorch、JAX などの複数のディープ ラーニング バックエンドをサポートしています。バックエンドは KERAS_BACKEND 環境変数を使用して構成されます。この変数は、Keras ライブラリをインポートする前に設定する必要があります。次のコード スニペットは、インタラクティブな Python 環境でこの変数を設定する方法を示しています。

import os

os.environ["KERAS_BACKEND"] = "tensorflow"  # or "jax" or "torch"

4. LIT の設定

LIT は、Python Notebook またはウェブサーバー経由で使用できます。この Codelab ではノートブックのユースケースに焦点を当てています。Google Colab でこの Codelab に沿って作業することをおすすめします。

この Codelab では、KerasNLP プリセットを使用して Gemma v2 2B IT を読み込みます。次のスニペットは、Gemma を初期化し、LIT Notebook ウィジェットにサンプル データセットを読み込みます。

from lit_nlp.examples.prompt_debugging import notebook as lit_pdbnb

lit_widget = lit_pdbnb.make_notebook_widget(
    ['sample_prompts'],
    ["gemma2_2b_it:gemma2_instruct_2b_en"],
)

ウィジェットを構成するには、2 つの必須位置引数に渡される値を変更します。

  • datasets_config: データセット名と読み込み元のパスを含む文字列のリスト(「dataset:path」形式)。パスには URL またはローカル ファイルパスを使用できます。次の例では、特別な値 sample_prompts を使用して、LIT ディストリビューションに提供されているサンプル プロンプトを読み込みます。
  • models_config: モデル名と読み込み元のパスを含む文字列のリスト(「model:path」形式)。パスには、URL、ローカル ファイルパス、構成済みのディープラーニング フレームワークのプリセット名を指定できます。

目的のモデルを使用するように LIT を構成したら、次のコード スニペットを実行してノートブックでウィジェットをレンダリングします。

lit_widget.render(open_in_new_tab=True)

独自のデータを使用する

Gemma はテキストからテキストへの生成モデルとして、テキスト入力を受け取ってテキスト出力を生成します。LIT は、独自の API を使用して、読み込まれたデータセットの構造をモデルに伝えます。LIT の LLM は、次の 2 つのフィールドを提供するデータセットを処理するように設計されています。

  • prompt: テキストを生成するモデルへの入力。
  • target: 人間のレーティング担当者からの「グラウンド トゥルース」レスポンスや、別のモデルから事前に生成されたレスポンスなど、オプションのターゲット シーケンス。

LIT には、この Codelab と LIT の拡張 プロンプト デバッグ チュートリアルをサポートする次のソースのサンプルを含む sample_prompts の少量のセットが含まれています。

独自のデータを簡単に読み込むこともできます。フィールド prompt とオプションで target を含むレコードを含む .jsonl ファイルとして読み込むか()、LIT の Dataset API を使用して任意の形式から読み込むことができます。

以下のセルを実行して、サンプル プロンプトを読み込みます。

5. LIT で Gemma の少数ショット プロンプトを分析する

現在、プロンプトは科学であると同時に芸術でもあります。LIT は、Gemma などの大規模言語モデルのプロンプトを経験的に改善するのに役立ちます。以降では、LIT を使用して Gemma の動作を調査し、潜在的な問題を予測して安全性を向上させる方法の例を紹介します。

複雑なプロンプトのエラーを特定する

高品質の LLM ベースのプロトタイプとアプリケーションに最も重要なプロンプト手法は、少数ショット プロンプト(プロンプトで望ましい動作の例を含む)と思考の連鎖(LLM の最終出力に先立って説明や推論の形式を含む)の 2 つです。しかし、効果的なプロンプトを作成することは、多くの場合、依然として困難です。

たとえば、好みに基づいて食べ物が好みに合うかどうかを判断する手助けをする例を考えてみましょう。初期プロトタイプの Chain-of-Thought プロンプト テンプレートは次のようになります。

def analyze_menu_item_template(food_likes, food_dislikes, menu_item):
  return f"""Analyze a menu item in a restaurant.

## For example:

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Onion soup
Analysis: it has cooked onions in it, which you don't like.
Recommendation: You have to try it.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Baguette maison au levain
Analysis: Home-made leaven bread in france is usually great
Recommendation: Likely good.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Macaron in france
Analysis: Sweet with many kinds of flavours
Recommendation: You have to try it.

## Now analyze one more example:

Taste-likes: {food_likes}
Taste-dislikes: {food_dislikes}
Suggestion: {menu_item}
Analysis:"""

このメッセージに問題があることに気付きましたか?LIT は、LM Salience モジュールを使用してプロンプトを調べるのに役立ちます。

6. デバッグにシーケンスの重要性を使用する

セレンスは可能な限り小さいレベル(入力トークンごとに)で計算されますが、LIT はトークンのセレンシーを、行、文、単語など、より解釈しやすい大きなスパンに集約できます。セリグナティと、それを使用して意図しないバイアスを特定する方法について詳しくは、セリグナティの探索ツールをご覧ください。

まず、プロンプト テンプレート変数に新しいサンプル入力を指定します。

food_likes = """Cheese"""
food_dislikes = """Can't eat eggs"""
menu_item = """Quiche Lorraine"""

prompt = analyze_menu_item_template(food_likes, food_dislikes, menu_item)
print(prompt)

fewshot_mistake_example = {'prompt': prompt}  # you'll use this below
Analyze a menu item in a restaurant.

## For example:

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Onion soup
Analysis: it has cooked onions in it, which you don't like.
Recommendation: You have to try it.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Baguette maison au levain
Analysis: Home-made leaven bread in france is usually great
Recommendation: Likely good.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Macaron in france
Analysis: Sweet with many kinds of flavours
Recommendation: You have to try it.

## Now analyze one more example:

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis:

上記のセルまたは別のタブで LIT UI を開いている場合は、LIT のデータポイント エディタを使用して、次のメッセージを追加できます。

LIT Datapoint Editor

別の方法として、目的のプロンプトを使用してウィジェットを直接再レンダリングすることもできます。

lit_widget.render(data=[fewshot_mistake_example])

モデルの予期しない完了に注目してください。

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis: A savoury tart with cheese and eggs
Recommendation: You might not like it, but it's worth trying.

食べられないと明確に言ったものを、モデルが食べるように提案するのはなぜですか?

シーケンスの重要度は、少数ショットの例にある根本的な問題を特定するのに役立ちます。最初の例では、分析セクション it has cooked onions in it, which you don't like の思考の流れの推論が、最終的な推奨事項 You have to try it と一致していません。

LM Salience モジュールで [Sentences] を選択し、推奨事項の行を選択します。UI は次のようになります。

LIT セレンシティ

これは、推奨事項の部分を誤ってコピーして貼り付け、更新しなかったというヒューマン エラーを示しています。

最初の例の「Recommendation」を Avoid に修正して、もう一度試してみましょう。LIT には、このサンプルがサンプル プロンプトにプリロードされているため、このユーティリティ関数を使用して取得できます。

def get_fewshot_example(source: str) -> str:
  for example in datasets['sample_prompts'].examples:
    if example['source'] == source:
      return example['prompt']
  raise ValueError(f'Source "{source}" not found in the dataset.')

lit_widget.render(data=[{'prompt': get_fewshot_example('fewshot-fixed')}])

モデルの完了は次のようになります。

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis: This dish contains eggs and cheese, which you don't like.
Recommendation: Avoid.

この事例から得られる重要な教訓は、早い段階でプロトタイプを作成すると、事前に想定していないリスクを明らかにできますということです。また、言語モデルはエラーが発生しやすい性質があるため、エラーに備えて事前に設計する必要があります。詳しくは、AI を使った設計に関する 人間と AI に関するガイドブックをご覧ください。

修正された少数ショット プロンプトは改善されていますが、まだ正しくありません。卵を食べないようにとユーザーに正しく伝えていますが、その理由が正しくありません。ユーザーは卵を食べられないと述べていますが、プロンプトでは卵が嫌いだと述べています。次のセクションでは、改善方法について説明します。

7. 仮説をテストしてモデルの動作を改善する

LIT を使用すると、同じインターフェース内でプロンプトの変更をテストできます。この例では、コンスティチューションを追加してモデルの動作を改善することをテストします。コンスティチューションは、モデルの生成をガイドする原則を含む設計プロンプトを指します。最近のメソッドでは、憲法原則のインタラクティブな導出も可能になっています。

このアイデアを使って、プロンプトをさらに改善してみましょう。プロンプトの上部に、生成の原則を記載したセクションを追加します。プロンプトは次のように始まります。

Analyze a menu item in a restaurant.

* The analysis should be brief and to the point.
* The analysis and recommendation should both be clear about the suitability for someone with a specified dietary restriction.

## For example:

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Onion soup
Analysis: it has cooked onions in it, which you don't like.
Recommendation: Avoid.

...

lit_widget.render(data=[{'prompt': get_fewshot_example('fewshot-constitution')}])

この更新により、例を再実行すると、まったく異なる出力を確認できます。

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis: This dish contains eggs, which you can't eat.
Recommendation: Not suitable for you.

プロンプトの視認性を再検討して、この変化が発生している理由を把握できます。

LIT セレンシティ

推奨事項ははるかに安全です。また、「不適切」は、分析(いわゆる思考の連鎖)とともに、食事制限に応じて適切性を明確に示すという原則の影響を受けています。これにより、出力が正しい理由で発生していることをさらに確信できます。

8. モデルのプローブと探索に技術チーム以外のチームを含める

解釈可能性は、XAI、ポリシー、法務などの専門知識を横断するチームの取り組みです。

従来、開発の初期段階でモデルを操作するには、高度な技術的な専門知識が必要でした。そのため、一部の共同作業者はモデルにアクセスしてプローブすることが困難でした。これまで、これらのチームが初期のプロトタイピング フェーズに参加できるツールはありませんでした。

LIT を通じて、このパラダイムを変革したいと考えています。この Codelab で学んだように、LIT の視覚的なメディアと、視認性を調べて例を探索するインタラクティブな機能は、さまざまな関係者が調査結果を共有して伝えるのに役立ちます。これにより、モデルの探索、プローブ、デバッグに多様なチームメイトを参加させることができます。これらの技術的な方法を学ぶことで、モデルの仕組みをより深く理解できます。また、初期モデルのテストでより多様な専門知識を活用することで、改善が必要な望ましくない結果を特定することもできます。

9. まとめ

まとめると次のようになります。

  • LIT UI は、インタラクティブなモデル実行用のインターフェースを提供します。これにより、ユーザーは出力を直接生成し、「what if」シナリオをテストできます。これは、さまざまなプロンプトのバリエーションをテストする場合に特に便利です。
  • LM Salience モジュールは、salience を視覚的に表し、制御可能なデータの粒度を提供します。これにより、モデル中心のコンストラクト(トークンなど)ではなく、人間中心のコンストラクト(文や単語など)について説明できます。

モデル評価で問題のある例を見つけたら、LIT に持ち込んでデバッグします。まず、モデリング タスクに論理的に関連する、考えられる最も大きなコンテンツ単位を分析します。可視化を使用して、モデルがプロンプト コンテンツに正しくまたは正しく対応している場所を確認します。次に、コンテンツのより小さな単位にドリルダウンして、表示されている誤った動作を詳しく記述し、考えられる修正方法を特定します。

最後に、Lit は常に改善されています。機能の詳細やご意見・ご要望については、こちらをご覧ください。