KAIST 전기및전자공학부 김준모 교수 연구팀, AI 안전성 검증 기술 ‘스테이블-지플로우넷’ 개발
튼튼한 성을 만들려면 먼저 약한 곳을 찾아야 한다. 생성형 인공지능(AI)도 마찬가지다. KAIST 연구진이 AI의 숨은 취약점을 기존보다 약 7배 더 다양하게 찾아내는 안전성 검증 기술을 개발했다. 더욱 안전하고 신뢰할 수 있는 AI 개발의 기반이 될 것으로 기대된다.
KAIST(총장 배충식)는 전기및전자공학부 김준모 교수 연구팀이 대규모 언어모델(Large Language Model, LLM)의 안전성을 검증하는 레드팀(Red Teaming, AI를 의도적으로 공격해 취약점을 찾는 안전성 검증 과정) 기술의 한계를 극복한 새로운 프레임워크 ‘스테이블-지플로우넷(Stable-GFlowNet)'을 개발했다고 밝혔다.
KAIST 김준모 교수, 권민찬 박사과정
연구 배경은
대규모 언어모델(LLM)이 다양한 실제 서비스에 빠르게 도입되면서, 모델이 유해하거나 위험한 답변을 내놓지 않도록 하는 안전성 확보가 핵심 과제로 떠올랐다. 이를 위해 배포 전에 모델의 취약점을 능동적으로 찾아내는 '레드팀(Red-Teaming)'이 필수 과정으로 자리 잡았다.
효과적인 레드팀을 위해서는 AI가 유해한 답변을 유도하는 '공격 프롬프트'를 스스로 생성하도록 학습시켜야 한다. 이때 단순히 성공률이 높은 공격만 찾는 것이 아니라, 서로 다른 유형의 약점을 폭넓게 발굴하는 '다양성'이 중요하다. 다양한 취약점을 미리 찾아낼수록, 이후 안전 학습(safety fine-tuning)을 통해 더 많은 위험을 차단할 수 있기 때문이다.
그러나 기존 방법들은 효과성과 다양성을 동시에 달성하기 어려웠다. 보상을 최대화하는 강화학습(RL) 기반 방법은 가장 효과가 큰 소수의 공격에만 수렴하는 '모드 붕괴(mode collapse)' 현상에 취약했고, 다양성을 강조하는 방법(Quality-Diversity)은 정해진 틀 안에서만 공격을 생성해 높은 성공률을 확보하기 어려웠다.
이러한 한계를 보완할 대안으로 '생성 흐름 네트워크(Generative Flow Networks, GFN)'가 주목받았다. GFN은 보상에 비례해 다양한 결과를 골고루 생성하도록 설계되어 레드팀에 적합하지만, 학습 과정에서 전체 보상의 총합을 추정하는 '분배 함수(Z)'를 별도로 학습해야 하는 구조였다. 이 추정이 부정확하면 학습이 불안정해지고 모드 붕괴가 가속화되는 문제가 있었다. 더욱이 유해성 판별기가 무의미한 문장에도 잘못된 보상을 부여하는 등 보상 신호에 잡음이 많아, 학습이 엉뚱한 방향으로 무너지기 쉬웠다.
S-GFN 메소드에 대한 전반적인 모식도. 공격을 생성하는 attacker llm을 훈련시키며, vicitm llm이 공격을 받고 그 응답을 toxic classifier가 채점한다. 채점된 보상은 기울기 계산에 사용되며, S-GFN의 3가지 컴포넌트 (Min-K Fluency Stabilizer, Noisy Gradient Pruning, Constrastive Trajectory Balance)를 모두 보여준다.
연구 내용은
연구는 이러한 GFN의 구조적 불안정성을 해결한 새로운 프레임워크 '스테이블-지에프엔(Stable-GFN, S-GFN)'을 제안한다. 핵심은 세 가지 기법으로 요약된다.
첫째, '대조 궤적 균형(Contrastive Trajectory Balance, CTB)'이다. 추정이 까다로운 분배 함수 Z를 직접 학습하는 대신, 같은 정책에서 뽑은 두 개의 결과물을 서로 짝지어 상대적으로 비교하는 방식을 도입했다. 이 과정에서 Z가 자연스럽게 상쇄되어 학습 불안정성을 근본적으로 줄이면서도, 기존 방식(Trajectory Balance)과 수학적으로 동일한 최적 정책에 도달함을 이론적으로 증명했다.
둘째, '노이즈 경사 가지치기(Noisy Gradient Pruning, NGP)'다. 두 결과물의 보상 차이가 의미 없이 작은 경우는 학습에서 제외하고, 뚜렷한 차이를 보이는 중요한 신호에만 집중하도록 했다. 이를 통해 유해성 판별기나 대상 모델의 무작위적 변동에 과도하게 휘둘리지 않고, 잡음이 많은 환경에서도 안정적인 학습이 가능해졌다.
셋째, '최소-K 유창성 안정화 장치(Min-K Fluency Stabilizer, MKS)'다. AI가 뜻 없이 깨진 문장(횡설수설)으로 보상을 가로채는 '보상 해킹'을 막기 위해, 문장에서 가장 부자연스러운 일부 단어들의 생성 확률을 기준으로 비정상적인 문장을 걸러냈다. 이를 통해 자연스러운 문장 형태를 유지하면서도 훨씬 넓은 공격 공간을 탐색할 수 있게 했다.
실험 결과, 스테이블-지에프엔은 기존 GFN 방식이 찾은 고유 공격 유형(17개)의 약 7배인 134개를 발굴하면서도 92%의 높은 공격 성공률을 유지했다. 또한 한 방법으로 찾은 공격을 막도록 모델을 학습시킨 뒤 서로 공격해 보는 '교차 공격' 평가에서, 본 기술로 학습한 방어 모델은 다른 공격들을 폭넓게 막아낸 반면, 본 기술의 공격은 다른 방법으로 방어 학습한 모델에도 여전히 유효했다. 이는 본 기술이 더 포괄적인 취약점을 찾아내 방어 성능의 일반화로 이어짐을 보여준다.
나아가 연구팀은 본 기술이 레드팀에 국한되지 않음을 확인했다. 신약 후보 물질을 찾는 '분자 생성'이나 잡음이 포함된 분포 매칭 문제 등 다른 영역에서도 기존 GFN 계열 방법보다 빠르고 안정적으로 수렴해, 폭넓게 확장 가능한 범용 기술임을 입증했다.
앞으로는
연구는 GFN의 고질적 한계였던 학습 불안정성과 모드 붕괴를 해결함으로써, 적은 데이터와 잡음이 많은 현실적 환경에서도 AI의 다양한 취약점을 안정적으로 발굴할 수 있는 길을 열었다. 이는 LLM 배포 전 더 폭넓은 위험을 사전에 점검하고 방어할 수 있게 해, 안전하고 신뢰할 수 있는 AI 개발에 핵심적으로 기여할 것으로 기대된다.
또한 제안된 방법론은 레드팀뿐 아니라 신약 개발, 소재 탐색 등 '다양하고 우수한 후보를 찾아야 하는' 다른 분야로도 확장될 수 있어, 향후 폭넓은 응용이 가능할 것으로 전망된다.
<저작권자(c)스마트앤컴퍼니. 무단전재-재배포금지>
