FlashRT: 프롬프트 주입 및 지식 오염 레드팀 테스트를 위한 계산 및 메모리 효율적 프레임워크

긴 문맥을 지원하는 LLM이 확산됨에 따라 프롬프트 주입과 같은 보안 위협을 평가하는 레드팀 테스트의 중요성이 커지고 있지만, 기존 최적화 기반 공격 방식은 막대한 GPU 메모리와 계산 시간을 요구한다. FlashRT는 선택적 재계산과 그래디언트 근사 기법을 통해 자원이 제한된 환경에서도 대규모 보안 취약점 진단을 가능하게 하여 LLM 안전성 연구의 진입 장벽을 낮춘다.

Transformer 아키텍처의 Self-Attention 메커니즘은 시퀀스 길이가 길어질수록 연산량과 메모리 요구량이 제곱 단위로 증가하며, 특히 최적화 기반 공격에서는 수천 번의 반복 계산이 필요해 비용 문제가 심각하다. 기존의 KV-Caching은 이전 토큰의 계산 결과를 재사용하지만, 문맥 중간에 공격용 텍스트가 삽입되면 그 이후의 모든 토큰에 대한 Attention 연산을 다시 수행해야 하는 한계가 있다.

FlashRT는 공격 텍스트가 전체 긴 문맥에서 차지하는 비중이 매우 작다는 점에 착안하여, 모든 후속 토큰을 재계산하는 대신 최종 출력 확률에 가장 큰 영향을 미치는 '중요 토큰'들만 골라 Attention을 업데이트한다. 이는 Attention Weights를 Influence Score로 활용하여 모델의 판단에 기여도가 낮은 배경 지식 토큰들의 재계산을 생략하는 원리이다.

또한, 가중치 업데이트를 위한 그래디언트 계산 시에도 전체 문맥을 다 참조하지 않고 무작위로 샘플링된 부분 문맥만 활용하여 역전파를 수행한다. 이러한 근사 방식은 정확한 그래디언트 값과는 차이가 있을 수 있으나, 이산적인 토큰 최적화 과정에서 유효한 탐색 방향을 제시하기에는 충분하다는 점을 활용해 메모리 병목 현상을 해결했다.

FlashRT의 핵심은 순전파 효율화를 위한 Selective Recomputing과 역전파 효율화를 위한 Gradient Approximation으로 구성된다. 순전파 단계에서는 먼저 현재 최적의 공격 텍스트(T_best)에 대한 KV-Cache를 저장하고, 중간 레이어의 Attention Weights를 평균 내어 각 문맥 토큰의 Influence Score를 산출한다. [문맥 토큰과 타겟 출력 토큰 간의 Attention 값 입력 → 레이어별 평균 연산 → Influence Score 출력 → 점수가 높은 상위 β 비율의 토큰 식별]

식별된 중요 토큰들에 대해서만 Hidden State를 업데이트하고 나머지는 기존 캐시를 유지함으로써 연산량을 줄인다. 이때 Attention 연산은 [선택된 쿼리 행렬 Q'와 업데이트된 K', V' 행렬 입력 → Softmax(Q'K'T/√dk)V' 연산 → 업데이트된 Attention Output 출력] 과정을 거치며, 이는 기존의 최적화된 커널과 호환되어 시스템 수준의 가속을 지원한다.

역전파 단계에서는 메모리 점유를 줄이기 위해 문맥을 여러 세그먼트로 나눈 뒤, 하이퍼파라미터 γ 비율만큼 무작위로 샘플링하여 부분적인 입력값(X_tilde)을 구성한다. [샘플링된 토큰 시퀀스 입력 → 축소된 그래프 기반 역전파 수행 → 근사 그래디언트 산출 → 공격 후보 토큰 생성] 과정을 통해 메모리 사용량을 억제하며, 최적화가 정체될 경우 다시 샘플링을 수행하는 Gradient Resampling 전략을 사용하여 지역 최적점(Local Minima) 탈출을 돕는다.

Llama-3.1-8B 모델을 사용한 NarrativeQA 데이터셋 실험에서 FlashRT는 기존 SOTA 방식인 nanoGCG 대비 실행 시간을 2736.9초에서 1039.5초로 약 2.6배 단축했으며, GPU 메모리 사용량은 164.8GB에서 53.7GB로 약 3배 절감했다. 특히 공격 성공률(ASR)은 오히려 10% 상승하는 결과를 보였다.

지식 오염(Knowledge Corruption) 공격에서도 NQ 데이터셋 기준 nanoGCG가 84.0GB의 메모리를 사용할 때 FlashRT는 29.8GB만을 사용하여 동일한 1.0의 ASR을 달성했다. 또한 4개의 H100 GPU로도 실행이 불가능했던 Llama-3.1-70B 모델에 대한 레드팀 테스트를 FlashRT는 성공적으로 수행하여 대규모 모델에 대한 확장성을 입증했다.

FlashRT는 Transformer의 Attention 메커니즘이 가진 희소성(Sparsity)을 활용한다. 연구진은 실험을 통해 긴 문맥 중 극히 일부의 토큰만이 최종 출력 생성에 유의미한 Attention Weights를 가진다는 점을 발견했으며, 이를 Influence Score로 정량화했다. 특히 모델의 중간 레이어(32레이어 모델 기준 15-19번 레이어)가 토큰 간의 의존 관계를 가장 잘 반영한다는 선행 연구에 기반하여 해당 레이어의 Attention 정보를 활용한다.

메모리 관리 측면에서 vLLM의 PagedAttention과 같은 시스템 최적화 도구들이 단일 단계 로그 확률 계산보다는 다단계 디코딩에 최적화되어 있음을 지적하며, FlashRT는 PyTorch의 SDPA(Scaled Dot Product Attention)를 기반으로 구현되어 단일 단계 계산 효율을 극대화했다. 또한 역전파 시의 메모리 절감을 위해 입력 토큰을 샘플링하는 방식은 활성화 값(Activations) 저장 공간을 직접적으로 줄여 긴 시퀀스에서도 선형적인 메모리 증가를 억제한다.

공격 텍스트가 문맥의 앞부분이나 중간에 위치할 때 가장 큰 효율 개선을 보이며, 문맥의 맨 끝에 위치할 경우 KV-Caching의 이점이 줄어들어 성능 향상 폭이 상대적으로 작아진다. 또한 근사 그래디언트를 사용하므로 정확한 그래디언트가 필수적인 특정 정밀 최적화 시나리오에서는 수렴 속도가 느려질 가능성이 있다.