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

1. はじめに

生成 AI プロダクトは比較的新しく、アプリケーションの動作は以前の形式のソフトウェアよりもばらつきがあります。そのため、使用されている ML モデルの精査、モデルの動作例の検証、想定外の挙動を調査することが重要です。

Learning Interpretability Tool(LIT、ウェブサイトGitHub)は、ML モデルのデバッグと分析を行い、モデルが期待どおりに動作する理由と方法を理解するためのプラットフォームです。

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

学習目標:

  1. シーケンスの顕著性と、モデル分析におけるその用途を理解する。
  2. プロンプト出力とシーケンスの顕著性を計算するために、Gemma の LIT を設定します。
  3. LM Salience モジュールでシーケンスの顕著性を使用して、プロンプト設計がモデルの出力に与える影響を理解する。
  4. LIT で仮定のプロンプト改善をテストし、その影響を確認する。

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

LIT UI のデモ

2. シーケンスの顕著性とモデル分析におけるその用途

Gemma などのテキストからテキストへの生成モデルは、トークン化されたテキストの形式で入力シーケンスを受け取り、その入力に続く典型的な新しいトークンを生成します。この生成は一度に 1 つのトークンを行い、モデルが停止条件に達するまで、新たに生成された各トークンを入力と前の世代に追加します(ループで)。たとえば、モデルがシーケンス終了(EOS)トークンを生成する場合や、事前定義された最大長に達した場合などです。

顕著性メソッドは Explainable AI(XAI)手法の一種で、出力のさまざまな部分について、入力のどの部分がモデルにとって重要であるかを示すことができます。LIT は、入力トークンのシーケンスが予測されたラベルに与える影響を説明する、さまざまな分類タスクの顕著性メソッドをサポートしています。シーケンスの顕著性は、これらの手法をテキストからテキストの生成モデルに一般化し、生成されたトークンに対する前のトークンの影響を説明します。

ここでは、シーケンスの顕著性に Grad L2 Norm メソッドを使用します。これにより、モデルの勾配を分析して、先行する各トークンが出力に与える影響の大きさがわかります。この方法はシンプルで効率的であり、分類やその他の設定でうまく機能することが実証されています。顕著性スコアが大きいほど、影響も大きくなります。この方法は、解釈可能性研究コミュニティ全体で広く利用されており、広く利用されているため、LIT 内で使用されています。

より高度な勾配ベースの顕著性手法には、Grad ⋅ 入力統合勾配などがあります。また、LIMESHAP など、アブレーション ベースの手法も使用できます。これらはより堅牢ですが、計算コストが大幅に高くなります。さまざまなサリエンス メソッドの詳細な比較については、こちらの記事をご覧ください。

salience に関するインタラクティブな入門ガイドで、salience メソッドの科学についてより詳しく学ぶことができます。

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

新しい Colab でこの Codelab に従うことをおすすめします。モデルをメモリに読み込むため、アクセラレータ ランタイムの使用をおすすめします。ただし、アクセラレータ オプションは時間とともに変化し、制限事項が適用されることに注意してください。より強力なアクセラレータを利用したい場合は、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.8.0)および lit-nlp(1.1)と Kaggle アカウントが必要です。

!pip install -q -U lit-nlp
!pip uninstall -y umap-learn
!pip install -q -U keras-nlp
!pip install -q -U keras

Kaggle アクセス

Kaggle にログインするには、kaggle.json 認証情報ファイルを ~/.kaggle/kaggle.json に保存するか、Colab 環境で次のコマンドを実行します。詳細については、kagglehub パッケージのドキュメントをご覧ください。

import kagglehub

kagglehub.login()

また、Gemma の使用許諾契約にも必ず同意してください。

4. LIT モデルの設定

import os

os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import keras_nlp

# Run at half precision.
keras.config.set_floatx("bfloat16")
model_name = 'gemma_instruct_2b_en'
gemma_model = keras_nlp.models.GemmaCausalLM.from_preset(model_name)

次のコードは、Gemma モデルで顕著性をサポートするために LIT ラッパーを初期化します。LIT フレームワークではこれらをモデルと呼びますが、このケースでは、上記で読み込んだものと同じ基盤 gemma_model に対する異なるエンドポイントです。これにより、LIT はオンデマンドで生成、トークン化、サリエンスを計算できます。

from lit_nlp.examples.models import instrumented_keras_lms

batch_size = 1
max_sequence_length = 512
init_models = instrumented_keras_lms.initialize_model_group_for_salience
models = init_models(model_name, gemma_model,
                     batch_size=batch_size,
                     max_length=max_sequence_length)

5. LIT データセットの設定

Gemma は、テキスト入力を受け取ってテキスト出力を生成するテキストからテキストの生成モデルです。LIT のモデルでは、データセットが生成をサポートするために次のフィールドを提供することを前提としています。

  • prompt: KerasGenerationModel への入力。
  • target: オプションのターゲット シーケンス(「グラウンド トゥルース」など)(ゴールド)回答、またはモデルから事前に生成された回答。

LIT には、次のようないくつかの異なるソースからの例を含む少数の sample_prompts のセットが含まれています。

  • GSM8K: 少数ショットの例で小学校の数学の問題を解決します。
  • Gigaword Benchmark: 一連の短い記事の見出しを生成します。
  • 憲法プロンプト: ガイドライン/境界のあるオブジェクトの使用方法に関する新しいアイデアを生み出します。

また、独自のデータを簡単に読み込むこともできます。.jsonl ファイルにはフィールド prompt と、必要に応じて target を含むレコード()が含まれるレコード、または LIT の Dataset API により任意の形式から読み込むことができます。

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

from lit_nlp.examples.datasets import lm as lm_data

datasets = {
  'sample_prompts': lm_data.PromptExamples(
      lm_data.PromptExamples.SAMPLE_DATA_PATH
  ),
}

6. LIT UI の設定

LIT はインタラクティブなモデル理解ツールであり、人間参加型の評価とモデルの動作のプロービングを可能にします。LIT UI では、次の操作が可能であり、このやり取りを容易にします。

  • データセットとモデルの出力をライブで可視化できます。
  • 顕著性メソッドを実行して、モデルの動作を制御する入力トークンを把握する
  • 仮説を検証するための反事実的条件を 立てることです

LIT では、これらすべてを同じインターフェース内で実現できるため、異なるツールを切り替える際の煩わしさが軽減されます。これは、この Codelab の後半で注目するプロンプト エンジニアリングのようなタスクで特に有用です。

この UI レイアウトは、他の任意の生成言語モデルに使用できます。ここに記載されていない機能に関心をお持ちの場合は、こちらでその一覧をご確認ください。

from lit_nlp.api import layout
modules = layout.LitModuleName

LM_SALIENCE_LAYOUT = layout.LitCanonicalLayout(
    left={
        'Data Table': [modules.DataTableModule],
        'Datapoint Editor': [modules.DatapointEditorModule],
    },
    upper={  # if 'lower' not specified, this fills the right side
        'Salience': [modules.LMSalienceModule],
    },
    layoutSettings=layout.LayoutSettings(leftWidth=40),
    description='Custom layout for language model salience.',
)

このコードは LIT サーバーを初期化します。また、サンプル プロンプトでモデルを実行して結果をキャッシュに保存するため、数秒かかることがあります。

from lit_nlp import notebook as lit_notebook

lit_widget = lit_notebook.LitWidget(
    models=models,
    datasets=datasets,
    layouts={'default': LM_SALIENCE_LAYOUT},
    default_layout='default',
)

これで、UI を表示できるようになりました。

lit_widget.render(height=800)

LIT UI のデモ

LIT は新しいタブでフルページとして開くこともできます。次のコードを実行して、表示されたリンクをクリックします。

lit_widget.render(open_in_new_tab=True)

注: 通常の .py スクリプトで LIT を使用している場合は、LitWidget ではなく lit_nlp.dev_server.Server() を使用してスタンドアロン サーバーを起動します。詳しくは、LIT のドキュメントをご覧ください。

7. LIT での Gemma の少数ショット プロンプトの分析

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

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

高品質の LLM ベースのプロトタイプとアプリケーションのための最も重要なプロンプト手法の 2 つは、少数ショット プロンプト(プロンプト内の望ましい行動の例を含む)と思考の連鎖(LLM の最終出力前の説明または推論の形態を含む)です。しかし、効果的なプロンプトの作成は依然として難しいものです。

自分の嗜好に基づいて食べ物が好きかどうかを評価する場合を手助けする例を考えてみましょう。最初のプロトタイプの 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 モジュールを使用してプロンプトを調べることができます。

8. デバッグにシーケンス サリエンスを使用する

顕著性は可能な限り小さいレベル(入力トークンごと)で計算されますが、LIT ではトークンの顕著性をより解釈しやすい大規模なスパン(行、文、単語など)に集約できます。顕著性の詳細と、それを使用して意図しないバイアスを特定する方法については、Saliency Explorable をご覧ください。

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

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 データポイント エディタ

もう一つの方法は、目的のプロンプトでウィジェットを直接再レンダリングすることです。

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 の顕著性

提案部分を誤ってコピーして貼り付け、更新に失敗するという人的ミスが明らかになります。

次に「最適化案」の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 を使用した設計に関する People + AI Guidebooks のガイドブックをご覧ください。

修正された少数ショット プロンプトは良いものの、まだ正しくありません。ユーザーに卵を避けるように正しく伝えていますが、推論は正しくありません。実際にはユーザーが卵を食べられないと述べているのに、卵が好きではないと言っています。次のセクションでは、どうすればよいか見ていきます。

9. 仮説を検証してモデルの動作を改善する

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 の顕著性

おすすめの方がはるかに安全であることがわかります。さらに [次のサービスに合わない] オプションには食事制限に従って適合性を明確に述べるという原則と、分析(いわゆる「思考の連鎖」)の影響を受けます。これにより、正しい理由で出力が行われているという確信が高まります。

10. モデルのプローブと探索に非技術系チームを含める

解釈可能性は、XAI、ポリシー、法務などの専門知識を幅広く活用し、チームで取り組むことを意図しています。

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

LIT を通じて、このパラダイムを変えることができるという希望があります。この Codelab で見てきたように、LIT の視覚的な媒体とインタラクティブな機能によって顕著性を検証し、例を探索することで、さまざまな関係者が調査結果を共有し、伝えることができます。これにより、モデルの探索、プローブ、デバッグに幅広いチームメイトを参加させることができます。こうした技術的な手法に触れることで、モデルの仕組みに対する理解を深められます。さらに、初期のモデルテストでより多様な専門知識のセットによって、改善可能な望ましくない結果を発見することもできます。

11. まとめ

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

  • LIT UI にはインタラクティブなモデル実行用のインターフェースが用意されており、ユーザーは出力を直接生成して「What if」をテストできる説明します。これは、さまざまなプロンプト バリエーションをテストする際に特に便利です。
  • LM Salience モジュールは、顕著性を視覚的に表現し、制御可能なデータ粒度を提供するため、モデル中心の構造(トークンなど)ではなく、人間中心の構造(文や単語など)について伝えることができます。

モデル評価で問題のある例が見つかった場合は、デバッグのために LIT に取り込みます。まず、モデリング タスクに論理的に関連すると考えられる最大のコンテンツ単位を分析し、可視化を使用して、モデルがプロンプト コンテンツに正しくまたは誤って対応している箇所を確認します。次に、コンテンツを小さな単位にドリルダウンして、修正候補を特定します。

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