Heads of Security는 학습된 어텐션 게이트로 프롬프트 인젝션을 막는다
Heads of Security: Learned Attention Gates Improve Prompt Injection Robustness
무엇인가
간접 프롬프트 인젝션은 모델이 외부 콘텐츠에 심어진 지시를 사용자 명령으로 착각할 때 발생한다. LLM이 웹 검색 결과, 툴 출력, 문서를 읽고 행동하는 에이전트로 쓰일수록 이 공격면은 권한 사고로 이어진다. 논문은 지시-데이터 분리(instruction-data separation)를 방어 원리로 잡는다. 신뢰된 지시는 모델 행동을 조종할 수 있지만, 비신뢰 데이터는 읽고 분석할 수는 있어도 지시로 따르지 않게 만드는 경계다. 저자들은 이 경계를 모델 아키텍처 안에서 강제하는 Heads of Security(HoS)를 제안한다.
어떻게 동작하나
핵심은 각 어텐션 헤드에 붙는 작은 채널별 게이트다. 토큰을 trusted, untrusted, template 세 채널로 나누고, 레이어 l과 헤드 k마다 C×C 게이트 로짓 g를 둔다. g_{c_i,c_j}는 쿼리 채널 c_i가 키 채널 c_j를 얼마나 볼지 조절한다. template 채널은 채팅 템플릿 토큰, 흔히 어텐션 싱크로 쓰이는 BOS를 포함한다. 채널 할당은 외부에서 주어진다. 예를 들어 시스템·사용자 프롬프트와 모델 생성은 trusted, 툴 출력은 untrusted다. 마지막 토큰의 어텐션은 softmax(qK^T/√d_h + {log σ(g_{c_n,c_j})})로 계산된다. σ(g)가 1에 가까우면 정상적으로 보고, 0에 가까우면 그 채널을 닫는다. softmax 안에서 log σ(g)를 더하는 것은 해당 키의 비정규화 어텐션 가중치에 σ(g)를 곱하는 것과 같다. 모든 키에 같은 게이트가 걸리면 softmax가 공통 이동을 상쇄하므로, 채널 간 차이만 출력에 영향을 준다. 그래서 모든 토큰이 같은 채널인 입력은 원모델과 동일하게 처리된다. GPT-OSS처럼 명시적 어텐션 싱크가 있는 모델은 싱크를 키 전용 채널로 하나 더 두어 g를 C×(C+1)로 확장한다.
무엇과 다른가
학습하는 파라미터는 게이트 로짓뿐이다. Llama-3.3-70B(L=80, h=64)에서는 46,080개 스칼라로, SecAlign의 349M, 전체 파인튜닝과 ASIDE의 70B와 대비된다. GPT-OSS-20B(L=24, h=64)에서는 싱크 키 채널을 포함해 18,432개로, SecAlign의 38M, ASIDE의 21B보다 훨씬 작다. 훈련은 SecAlign++ 레시피를 따른다. DPO 목적과 변형된 Cleaned-Alpaca 데이터를 쓰고, 각 선호 쌍은 사용자 지시와 비신뢰 구간에 삽입된 인젝션으로 구성된다. 선호 응답은 사용자 지시를 따르고, 거부 응답은 인젝션을 따른다. 즉 모델 전체를 넓게 파인튜닝하는 대신, 어텐션 라우팅만 학습해 지시-데이터 분리를 익히게 한다.
어떻게 쓰나
적응형 공격 평가가 이 논문의 핵심이다. 블랙박스 RL-Hammer 공격에서 GPT-OSS-20B는 무방어 97%, SecAlign 85%, HoS 10%의 공격 성공률(ASR)을 기록했다. Llama-3.3-70B 계열은 무방어 99%, Meta-SecAlign-70B 82%, HoS 1%다. 저자들은 HoS가 GPT-5.6-Luna보다도 낮은 ASR을 보였다고 적는다. 화이트박스에서는 Claudini 공격을 GPT-5.6-Sol로 50회 자동연구 루프를 돌려 HoS에 맞춰 강화한 V38 변형을 만들었다. 100개 AgentDojo 사용자-인젝션 과제와 100개 InjecAgent 목표에서 120토큰 접미사를 10^17 FLOPs 예산으로 최적화했을 때, HoS는 50% ASR에 도달하는 데 기준 모델이나 SecAlign보다 약 3배의 공격 연산이 필요했다. AgentDojo에서는 약 6×10^16 FLOPs 대 기준 2.2–2.3×10^16, InjecAgent에서는 3.0×10^16 대 8.5–9.5×10^15였다. 이는 각각 2.6–2.7배, 3.2–3.6배다.
전제와 한계
유틸리티 손실은 작다. GPT-OSS-20B의 5-shot MMLU-Pro에서 HoS는 67%, 원모델은 68%였고 SecAlign은 59%, ASIDE는 52%로 떨어졌다. Llama-3.3-70B에서는 HoS 55%, 원모델 56%, SecAlign 16%다. Meta-SecAlign-70B는 1-shot GSM8K 2.0%, 5-shot MMLU 28.4%, 5-shot MMLU-Pro 16.1%까지 내려가는 경우가 있었다. 3턴 AlpacaFarm 대화에서는 원모델이 97–100% 유틸리티를 유지했고 HoS가 그에 근접했다. SecAlign은 입력 필드가 2턴에 나오면 43%, 3턴에 나오면 75%로 떨어졌고, 1턴에 입력이 있으면 이후 2·3턴 유틸리티가 71%, 82%였다. 에이전트 벤치마크인 AgentDojo, AgentDyn, LLMail-Inject, WASP에서는 HoS와 SecAlign이 대체로 비슷했고, HoS는 대부분 평가에서 유틸리티 손실 없이 ASR을 크게 낮췄다. GPT-OSS-20B에서 HoS는 AgentDojo와 WASP의 SecAlign 프런티어를 파레토 지배하고 AgentDyn과 LLMail에서는 파레토 동등했다. 반면 ASIDE는 툴 호출 능력을 잃어 세 벤치마크에서 거의 0, AgentDojo에서 10% 유틸리티에 그쳤다. SecAlign은 LLMail에서 사용자가 요약만 요청해도 29.95% 생성에서 툴 호출을 했고, 이 때문에 ASR이 높았다.
게이트가 실제로 무엇을 배웠는지도 분석한다. 신뢰 쿼리에서 신뢰 키와 비신뢰 키로 가는 게이트 차이 Δ를 보면, 대부분 헤드는 비신뢰 키에 대한 상대적 어텐션을 줄인다. 저자들은 이를 red 헤드라 부른다. 반대로 비신뢰 콘텐츠 쪽으로 어텐션을 늘리는 소수 헤드도 있는데, 이를 blue 헤드라 하고 인젝션 탐지에 기여한다고 해석한다. LLMail-Inject 하위 과제에서 네 모델 모두 이메일 전송률은 1% 미만이었지만 탐지율은 달랐다. 원모델 11.03%, blue 게이트만 켠 경우 28.20%로 약 2.5배, red 게이트만 켠 경우 3.60%, 전체 HoS 게이트는 3.11%였다. 인과적 헤드 절제와 역할 혼동(role confusion) 프로브 실험에서는 HoS가 툴 출력의 역할 인식을 개선하고, 공격된 AgentDojo에서 인젝션 토큰의 Userness를 낮추는 것으로 나타났다.
개발자 관점에서 HoS는 에이전트가 툴 출력, 웹 페이지, 문서를 데이터로만 읽어야 하는 상황에 맞는 방어 프리미티브다. 전체 모델을 DPO로 넓게 파인튜닝하지 않고도 헤드별 게이트만 학습해 지시-데이터 경계를 넣을 수 있고, 파라미터 수가 작아 배포·실험 비용이 낮다. 도입 전에 토큰 채널 라벨링이 정확한지, 채팅 템플릿과 BOS 같은 어텐션 싱크가 올바른 채널로 분류되는지 확인해야 한다. 또한 정적 공격뿐 아니라 RL-Hammer나 화이트박스 최적화 같은 적응형 공격으로 자체 평가해야 한다. HoS는 신뢰 경계를 강제할 뿐, 어떤 소스가 신뢰할 수 있는지 판단하지는 않는다.
저자들이 밝힌 한계는 분명하다. 지시-데이터 분리는 종단 간 방어가 아니다. HoS는 신뢰 경계를 강제하지만, 어떤 소스를 신뢰할지나 사용자가 어떤 지시를 승인했는지는 결정하지 않는다. 패키지 설치처럼 외부 콘텐츠의 지시를 따라야 하는 작업에는 엄격한 분리가 맞지 않는다. 훈련 데이터는 단일 턴 Alpaca 기반으로 고정되어 있어, 다중 턴·다중 지시·다양한 템플릿 데이터가 양쪽 방법을 개선하고 SecAlign의 분포 밖 실패를 줄일 수 있다고 본다. 다만 동일 데이터에서 SecAlign은 지름길을 학습했고 HoS는 그렇지 않았다는 점을 근거로, 학습 메커니즘을 제한하면 훈련 분포의 의도치 않은 편향을 습득할 능력이 줄어든다고 주장한다. 공격자는 방어를 본 뒤 적응할 수 있으므로 유한한 평가로 보안을 확정할 수 없다. 마지막으로 HoS는 표준 dense attention용이라, 특정 키 토큰에 대한 명시적 어텐션 계수가 없는 선형 어텐션 아키텍처에는 직접 호환되지 않는다.