これは、私たちの論文「Your Agentic LLMs Secretly Encode Latent Signals of Indirect Prompt-Injection Exposure」のコードリポジトリです。
ArXiv版と論文リンク: https://arxiv.org/abs/2608.02657
このリポジトリは、潜在的なIPI露出シグナルのプロービングパイプライン(AgentDojoトレース収集、IPIリスクラベリング、隠れ状態のフィーチャー化、IPI露出プローブのトレーニング/評価)を実装しています。
@misc{dong2026agenticllmssecretlyencode,
title={Your Agentic LLMs Secretly Encode Latent Signals of Indirect Prompt-Injection Exposure},
author={Jianshuo Dong and Yiming Liu and Maosen Zhang and Nan Deng and Xu Peng and Xiaoping Zhang and Tianwei Zhang and Jie Zhang and Han Qiu},
year={2026},
eprint={2608.02657},
archivePrefix={arXiv},
primaryClass={cs.CR},
url={https://arxiv.org/abs/2608.02657},
}
まずCPUフレンドリーなチェックから始めましょう:
git clone https://github.com/jianshuod/IPI-exposure-signal.git
cd IPI-exposure-signal
# Install uv if needed.
curl -LsSf https://astral.sh/uv/install.sh | sh
# This repository targets Python 3.12.
uv python install 3.12
uv sync --extra dev
uv run pytest
テストスイートはGPUリソースなしで実行できるように設計されています。コアメッセージのシリアライゼーション、AgentDojoのロングホライズン攻撃登録、ラベリングロジック、プローブデータセットの構築、トレーニングメトリクス、トークン集計、vLLMパッチのパッケージングをチェックします。
収集では、コレクターをOpenAI互換のモデルサーバーに向けます:
export IPI_AWARE_UPSTREAM_BASE_URL=http://127.0.0.1:8000/v1
export IPI_AWARE_UPSTREAM_API_KEY=EMPTY
隠れ状態の抽出とプローブトレーニングでは、モデル依存関係とCUDA環境に一致するPyTorchビルドをインストールします:
uv sync --extra dev --extra models
主なワークフローは ipi-signal-probe で公開されています:
例:
uv run ipi-signal-probe collect \
--suite workspace \
--attack direct \
--injected \
--max-cases 4 \
--results-dir results/probe_traces/v2
uv run ipi-signal-probe label \
--root results/probe_traces/v2/<run-name> \
--labeling-protocol risk_faced
uv run ipi-signal-probe featurize \
--root results/probe_traces/v2/<run-name> \
--model Qwen/Qwen3-8B \
--backend transformers_hook \
--model-family qwen3 \
--selected-position -1
uv run ipi-signal-probe partition \
--root results/probe_traces/v2/<run-name> \
--dataset broad
# Training consumes a JSON config. The Qwen example below generates the
# train/eval configs used for its full run.
サポートされているラベリングプロトコルは、risk_faced、risk_visible、risk_actual です。
examples/qwen3_5_0_8b/ ディレクトリには、Qwen/Qwen3.5-0.8B 用の実行可能な小規模モデルの再現スクリプトとオープンソースのユーザビリティテストが含まれています。
8つの1GPU vLLMサーバーを起動します:
examples/qwen3_5_0_8b/serve_qwen3_5_0_8b_8x.sh start
スモークテストを実行します:
QWEN3_5_0_8B_MODE=smoke \
examples/qwen3_5_0_8b/run_probe_pipeline.sh start
完全な例を実行します:
QWEN3_5_0_8B_MODE=full \
examples/qwen3_5_0_8b/run_probe_pipeline.sh start
完全な例では、8つの独立したTP=1 vLLMサーバー、8つのコレクタープロセス、プロセス/サーバーあたり256ワーカー、8つのフィーチャー化シャード、1GPUのプローブトレーニングシャードを使用します。すべての環境変数、レジュームモード、キャッシュ設定、ステータス/停止コマンドについては、examples/qwen3_5_0_8b/README.md を参照してください。
デフォルトのフィーチャー化バックエンドは transformers_hook です。vLLM直接隠れ状態キャプチャバックエンドには、チェックインされたvLLM 0.19.0オーバーレイが必要です:
uv sync --frozen --extra dev --extra vllm
uv run python scripts/apply_vllm_ipi_aware_patch.py --check
使用前に docs/vllm_patches_v0.19.0.md に従ってください:
uv run ipi-signal-probe featurize \
--root results/probe_traces/v2/<run-name> \
--model Qwen/Qwen3-8B \
--backend vllm_direct \
--model-family qwen3 \
--selected-position -1
このリポジトリは、vendor/agentdojo の下にパッチ適用済みのAgentDojoチェックアウトを同梱しています。アップストリームのAgentDojoは、この論文で必要とされるロングホライズン攻撃をネイティブに公開していないため、収集コマンドはこのソースチェックアウトから実行するか、vendor/agentdojo/src を PYTHONPATH 上のインストール済みアップストリームAgentDojoより前に配置して実行する必要があります。
IPI-exposure-signal/
├── ipi_aware/
│ ├── data_collection/ # AgentDojo traces and decision points
│ ├── message_content.py # Chat-message normalization used by collection/features
│ └── probes/ # labels, features, partitions, training, eval
├── examples/
│ └── qwen3_5_0_8b/ # 8-GPU Qwen/Qwen3.5-0.8B reproduction scripts
├── patches/ # vLLM 0.19.0 hidden-state capture overlay
├── vendor/agentdojo/ # patched AgentDojo source for long-horizon attacks
├── docs/ # vLLM patch notes
├── scripts/ # utility scripts
└── tests/ # CPU-friendly regression tests
このプロジェクトはApache-2.0でライセンスされています。LICENSE を参照してください。
| Command | Purpose |
|---|
collect | AgentDojoタスクを実行し、トレースと決定ポイントを保存する |
label | 収集された決定ポイントにIPIリスクラベルを割り当てる |
featurize | プローブ例のモデル隠れ状態を抽出する |
partition | トレイン/検証/評価の分割を構築する |
train | 隠れ状態の特徴量で軽量プローブをトレーニングする |
eval | 設定された分割でトレーニング済みプローブを評価する |
eval-groups | ラベリングプロトコルとホールドアウト設定のグループ評価を実行する |
| 変数 | 目的 |
|---|
IPI_AWARE_UPSTREAM_BASE_URL | 収集で使用されるOpenAI互換モデルエンドポイント |
IPI_AWARE_UPSTREAM_API_KEY | アップストリームエンドポイントのAPIキー。ローカルvLLMの場合は EMPTY |
PYTHONPATH | このチェックアウト外でパッチ適用済みAgentDojoソースを使用する場合は、vendor/agentdojo/src を先頭に配置する |
QWEN3_5_0_8B_MODEL | 例スクリプト内のデフォルトのQwen/Qwen3.5-0.8Bチェックポイントパスを上書きする |
QWEN3_5_0_8B_CACHE_DIR | モデルとローカルキャッシュをデータファイルシステムにリダイレクトする |
QWEN3_5_0_8B_FEAT_BACKEND | vLLMオーバーレイを適用した後、vllm_direct に設定する |