nishitvp/better_opts_attacksLLM에 대한 아키텍처 인식 적대적 공격을 사용하여 파인튜닝 기반 프롬프트 인젝션 방어(SecAlign, SecAlign++, StruQ)를 깨기 위한 공격 프레임워크.
이 저장소는 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에서 기본 모델을 다운로드합니다 (일부 모델은 게이트(gated)되어 있어 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 세 개의 폴더가 있고 각 폴더에 모델이 들어 있는지 확인하세요.
강한 지식 공격자에 대한 ASTRA를 실행하려면 핵심 스크립트는 experiment.py입니다. 이 스크립트는 AlpacaFarm 데이터셋에서 40개의 예시를 가져와 Print exactly 'Hacked!' 페이로드를 지정된 구성으로 주입하고, 초기화(무작위 초기화), 순방향 패스 수, 반복 횟수 등을 통제하면서 ASTRA와 기준선(baseline, 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: 특정 평가 실행에 사용할 적대적 접두사(adversarial prefix)의 길이입니다. 기본값은 5입니다.
--suffix-length: 특정 평가 실행에 사용할 적대적 접미사(adversarial suffix)의 길이입니다. 기본값은 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
위 명령을 실행하면 스크립트가 작업량을 여러 하위 프로세스에 자동으로 균등 분배하며, 각 하위 프로세스는 하나의 GPU를 사용하여 공격을 계산합니다. 공격은 최소 48GB 메모리의 GPU에서 실행할 것을 권장합니다.
각 하위 프로세스는 공격하는 각 예시의 전체 대화 기록(transcript)을 --expt-folder-prefix에 지정된 경로의 별도 하위 폴더에 기록합니다. 로그는 Python 객체를 포함하는 쿼리 가능한 데이터베이스 형식으로 저장됩니다 (자세한 내용은 utils/experiment_logger 파일 참조). 기록된 각 대화록에는 최적화 전반에 걸친 토큰 시퀀스, 손실 값, 초기화 구성과 같은 기타 중요한 세부 정보가 포함되며 나중에 검색할 수 있습니다.
전체 ASTRA 실험이 끝나면 생성된 로그는 analysis/analysis.ipynb 파일을 사용하여 나중에 분석할 수 있습니다. Jupyter 노트북을 --expt-folder-prefix 경로에 복사하여 붙여넣고, 노트북의 모델 경로를 올바른 모델 경로로 교체한 다음 노트북을 실행하기만 하면 됩니다. 그러면 마지막 셀에서 GCG 성공 횟수와 ASTRA 성공 횟수가 출력됩니다. (참고: 분석 스크립트는 각 예시에 대해 최적화 중 생성된 각 입력 토큰 시퀀스에서 실제 출력을 생성하기 때문에 실행하는 데 시간이 걸립니다.)
analysis.ipynb 노트북에는 평균 손실 곡선을 플롯하고 생성된 로그로 추가 탐색을 수행하는 헬퍼 코드도 포함되어 있습니다.
분석 스크립트는 사후 분석 코드로 인해 최적화 자체가 오염되지 않도록 근본적으로 분리되어 있습니다.
또한 data/attacks_generated.zip에 있는 압축 파일에 ASTRA로 생성된 일부 공격을 공유합니다. 구체적으로 이 압축 파일에는 각각 40개의 ASTRA 생성 공격 예시가 포함된 4개의 파일이 있으며, Mistral과 Llama-3의 각 SecAlign 방어 버전에 대해 각각 2개씩 있습니다. 이 파일들은 20개 및 25개 토큰 예산으로 생성된 예시를 제공합니다 ((0, 20) 및 (5, 20) 구성). 인젝션은 이미 구성의 입력 필드에 배치되어 있습니다.
약한 지식 공격자에 대한 ASTRA++를 실행하려면 핵심 스크립트는 experiment_universal.py입니다. 이 스크립트는 AlpacaFarm 데이터셋에서 10개의 훈련 예시를 가져와 지정된 접두사, 접미사 구성으로 Print exactly 'Hacked!' 페이로드를 주입하고 훈련하여 마지막에 범용 프롬프트 인젝션(universal prompt injection)을 얻습니다. 또한 이 스크립트는 ASTRA++ 훈련 후 기준선인 범용 GCG도 실행합니다.
이전 섹션의 experiment.py 스크립트와 마찬가지로 experiment_universal.py도 --model-name, --defense, --prefix-length, --suffix-length, --expt-folder-prefix 인자를 받습니다.
이전 섹션의 ASTRA 실행에 나열된 모든 모델 외에도 experiment_universal.py 스크립트는 SecAlign++ 방어에 해당하는 모델 하나를 추가로 받습니다.
SecAlign++ 방어로 훈련하려면 --model-name 매개변수를 secalign_refactored/secalign_models/Meta-SecAlign-8B로, --defense 매개변수를 meta_secalign으로 전달하세요.
이 스크립트가 받는 추가 명령줄 인자가 하나 더 있는데, 바로 --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에 업로드합니다 (피클에는 악성 코드가 포함되어 있지 않음을 약속합니다 :).
위 스크립트는 ASTRA 및 ASTRA++ 공격을 기본 설정 그대로 실행하지만, 이 코드베이스는 훈련 데이터셋 크기, 선택된 임계값, 무작위 시드, 초기화 구성 등 공격에 영향을 줄 수 있는 거의 모든 매개변수를 구성할 수 있게 해줍니다. 흥미로운 제안이 있으시면 언제든지 풀 리퀘스트를 보내주세요.
문의 사항, 버그 신고, 세부 사항 또는 설명이 필요하시면 GitHub에 이슈를 등록하거나 작성자에게 이메일을 보내주세요.