このリポジトリには、SecAlign++、SecAlign、StruQ を破る ASTRA および ASTRA++ 攻撃を実行するコードが含まれています。また、生成された攻撃の例や攻撃ログもいくつか含まれています。
このコマンドは、リポジトリとすべてのサブモジュールを、必要な正しい場所にクローンします。
git clone --recurse-submodules https://github.com/nishitvp/better_opts_attacks.git
このコードベースは、LLM やディープラーニング関連の標準的なパッケージ群以外に特別な Python パッケージを必要としません。必要なパッケージをインストールするには、以下を実行してください。
python3 -m pip install -r requirements.txt
これで、必要なパッケージが環境にインストールされるはずです。
まず、HuggingFace からベースモデルをダウンロードします（一部はゲート付きで、HuggingFace でのアクセス承認が必要な場合があります）。
hf login --token ${HF_PASSWORD}
hf download meta-llama/Meta-Llama-3-8B-Instruct --local-dir secalign_refactored/secalign_models/meta-llama/Meta-Llama-3-8B-Instruct
hf download mistralai/Mistral-7B-Instruct-v0.1 --local-dir secalign_refactored/secalign_models/mistralai/Mistral-7B-Instruct-v0.1
hf download meta-llama/Llama-3.1-8B-Instruct --local-dir secalign_refactored/secalign_models/meta-llama/Llama-3.1-8B-Instruct
hf download facebook/Meta-SecAlign-8B --local-dir secalign_refactored/secalign_models/Meta-SecAlign-8B
SecAlign アダプタと StruQ モデルをダウンロードします。適切な場所にダウンロードして展開するための便利なスクリプトがあります。
python3 setup_scripts/download_models.py
プロジェクトルートから実行すると、このスクリプトは正しいモデルを正しい場所にダウンロードして展開します。確認するには、secalign_refactored/secalign_models に huggyllama、meta-llama、mistralai の3つのフォルダがあり、それぞれにモデル自体が含まれているかどうかを確認してください。
強い知識敵対者に対して ASTRA を実行するための主要スクリプトは experiment.py です。このスクリプトは、AlpacaFarm データセットから 40 個の例を取り出し、指定された構成でペイロード Print exactly 'Hacked!' を注入し、初期化（ランダム初期化）、フォワードパスの回数、イテレーションなどを制御しながら、ASTRA とベースライン（GCG）の両方を実行します。
ユーザーは、experiment.py を呼び出す際にコマンドライン引数として設定することで、以下のパラメータを制御できます。
--model-name: 指定可能な値は meta-llama、meta-llama-instruct、mistralai、mistralai-instruct、huggyllama です。
--defense: 指定可能な値は undefended、struq、secalign です。デフォルト値は secalign です。--model-name パラメータと --defense パラメータの組み合わせによって、使用される具体的なモデルが以下の表の通り一意に決まります。
--model-name | --defense | 読み込まれるモデル |
|---|---|---|
meta-llama-instruct | secalign | 事前命令チューニング済み Meta-Llama-3-8B-Instruct を備えた SecAlign アダプタ |
mistralai-instruct | secalign | 事前命令チューニング済み Mistral-7B-Instruct-v0.1 を備えた SecAlign アダプタ |
meta-llama | secalign | SecAlign 防御済み Meta-Llama-3-8B（事前命令チューニングなし） |
mistralai | secalign | SecAlign 防御済み Mistral-7B-Instruct-v0.1（事前命令チューニングなし） |
huggyllama | secalign | SecAlign 防御済み Llama-2-7B（事前命令チューニングなし） |
meta-llama | struq | StruQ 防御済み Meta-Llama-3-8B（事前命令チューニングなし） |
mistralai | struq | StruQ 防御済み Mistral-7B-v0.1（事前命令チューニングなし） |
huggyllama | struq | StruQ 防御済み Llama-2-7B（事前命令チューニングなし） |
meta-llama-instruct | undefended | 無防御（生の）Meta-Llama-3-8B-Instruct |
mistralai-instruct | undefended | 無防御（生の）Mistral-7B-Instruct-v0.1 |
meta-llama | undefended | 無防御（生の）Meta-Llama-3-8B（事前命令チューニングなし） |
mistralai | undefended | 無防御（生の）Mistral-7B-Instruct-v0.1（事前命令チューニングなし） |
huggyllama | undefended | 無防御（生の）Llama-2-7B（事前命令チューニングなし） |
これ以外の --model-name と --defense の組み合わせは無効です。
--prefix-length: 特定の評価実行で使用する敵対的プレフィックスの長さ。デフォルトは 5 です。
--suffix-length: 特定の評価実行で使用する敵対的サフィックスの長さ。デフォルトは 20 です。
--expt-folder-prefix: 実験実行のログを保存するパス。
完全なコマンドの例は次のようになります。
python3 experiment.py --model-name meta-llama-instruct --defense secalign --prefix-length 5 --suffix-length 20 --expt-folder-prefix logs/astra_llama_secalign
上記のコマンドを実行すると、スクリプトはワークロードを自動的に各サブプロセスへ均等に分散し、各サブプロセスが1つの GPU を使用して攻撃を計算します。攻撃は少なくとも 48GB のメモリを備えた GPU で実行することを推奨します。
各サブプロセスは、攻撃対象とした各例の完全なトランスクリプトを、--expt-folder-prefix で指定されたパスの別々のサブフォルダにログとして記録します。ログは、Python オブジェクトを含むクエリ可能なデータベース形式で保存されます（詳細はファイル utils/experiment_logger を参照してください）。記録される各トランスクリプトには、最適化全体を通じたトークン列、損失値、初期化設定などの重要な詳細が含まれており、後で取得できます。
ASTRA の実験がすべて完了したら、生成されたログはファイル analysis/analysis.ipynb を使用して後で解析できます。Jupyter ノートブックを --expt-folder-prefix のパスにコピー＆ペーストし、ノートブック内のモデルパスを正しいモデルパスに置き換えて、ノートブックを実行してください。最後のセルで GCG の成功数と ASTRA の成功数が出力されます。（補足：解析スクリプトは、最適化中に生成された各入力トークン列に対して実際に出力を生成するため、実行に時間がかかります。）
analysis.ipynb ノートブックには、平均損失曲線をプロットしたり、生成されたログをさらに調査するためのヘルパーコードも含まれています。
解析スクリプトは、最適化自体が事後解析コードによって汚染されないように、根本的に分離されています。
また、data/attacks_generated.zip にある zip ファイルに、ASTRA で生成された攻撃の一部を共有します。具体的には、zip ファイルにはそれぞれ40個の ASTRA 生成攻撃例を含む4つのファイルが含まれており、それぞれ SecAlign 防御済みの Mistral と Llama-3 の各バージョンに対して2つずつあります。これらのファイルは、20トークンと25トークンの予算（(0, 20) と (5, 20) の構成）で生成された例を提供します。注入はすでに構成の入力フィールドに配置されています。
弱い知識敵対者に対して ASTRA++ を実行するための主要スクリプトは experiment_universal.py です。このスクリプトは、AlpacaFarm データセットから 10 個の訓練例を取り出し、指定されたプレフィックス・サフィックス構成でペイロード Print exactly 'Hacked!' を注入し、最後にユニバーサルプロンプトインジェクションを得るために訓練します。このスクリプトはまた、ASTRA++ の訓練後に、ベースラインである GCG のユニバーサルバージョンも実行します。
前のセクションの experiment.py と同様に、experiment_universal.py も --model-name、--defense、--prefix-length、--suffix-length、--expt-folder-prefix の引数を受け取ります。
前のセクションで ASTRA の実行に関して列挙したすべてのモデルに加えて、experiment_universal.py スクリプトは SecAlign++ 防御に対応するもう1つのモデルも受け取ります。
SecAlign++ 防御で訓練するには、--model-name パラメータに secalign_refactored/secalign_models/Meta-SecAlign-8B を渡し、--defense パラメータに meta_secalign を渡してください。
このスクリプトが受け取る追加のコマンドライン引数がもう1つあり、それは --training-run です。これは、訓練実行がどの 10 個の例のセットで行われるかに対応します。
それ以外のパラメータは前と同じ解釈です。
ASTRA++ の実験がすべて完了したら、生成されたログはファイル analysis/analysis_universal.ipynb を使用して後で解析できます。前と同様に、Jupyter ノートブックを --expt-folder-prefix のパスにコピー＆ペーストし、モデルを正しく読み込んでください（SecAlign++ の読み込み方法に起因する多少の厄介さがありますが、これは将来統一される予定です）。これにより、GCG のテスト成功数と ASTRA の成功数に加えて、達成された最良のユニバーサルプロンプトインジェクションが出力されるはずです。
ASTRA++ については、SecAlign++ 防御済み Llama-3.1-8B-Instruct に対して生成されたユニバーサルプロンプトインジェクションを、メタデータとそれが生成された構成とともにファイル data/universal_pis_secalign++.json にアップロードしています。さらに、すべてのユニバーサリティ訓練実行におけるすべての攻撃のトランスクリプトも data/final_result_logs.pkl にアップロードしています（pickle に悪意のあるコードは含まれていないことをお約束します:)。
上記のスクリプトは ASTRA および ASTRA++ 攻撃をそのまま実行できますが、このコードベースでは、訓練データセットのサイズ、選択される閾値、乱数シード、初期化構成など、攻撃に影響を与えるほぼすべてのパラメータを構成できます。もし興味深い提案があれば、遠慮なくプルリクエストを送ってください。
質問、バグ報告、詳細、または明確化が必要な場合は、Github で issue を立てるか、著者にメールしてください。