NASimJax: GPU-Accelerated Policy Learning Framework for Penetration Testing

1. 문제 정의 (Problem)

사이버 보안 분야에서 침투 테스트 (Penetration Testing) 는 시스템의 취약점을 식별하기 위해 사이버 공격을 시뮬레이션하는 복잡한 순차적 의사결정 작업입니다. 강화학습 (RL) 을 통해 이를 자동화하려는 시도가 증가하고 있지만, 다음과 같은 근본적인 병목 현상에 직면해 있습니다.

• 시뮬레이션 속도의 한계: 기존 침투 테스트 시뮬레이터 (예: NASim) 는 CPU 기반의 Python 구현체로, 현실적인 대규모 네트워크 시나리오를 학습할 만큼의 경험을 생성하는 속도가 매우 느립니다. RL 알고리즘은 수백만 번의 환경 상호작용이 필요하므로, 현대의 GPU 가속기를 효율적으로 활용하지 못합니다.
• 일반화 (Generalization) 부족: 기존 시뮬레이터는 고정된 시나리오에 치중하거나 다양성이 부족하여, 학습된 정책이 훈련되지 않은 새로운 네트워크 환경 (Zero-shot) 에 적용될 때 성능이 급격히 저하됩니다.
• 행동 공간 (Action Space) 의 확장성 문제: 네트워크 호스트 수가 증가함에 따라 행동 공간이 선형적으로 급증하여 탐색 (Exploration) 이 어려워지고 학습이 느려집니다.

2. 방법론 (Methodology)

저자들은 위 문제를 해결하기 위해 NASimJax를 제안했습니다. 이는 기존 Python 기반의 Network Attack Simulator (NASim) 를 JAX 라이브러리를 사용하여 완전히 재구현하고 확장한 프레임워크입니다.

2.1 NASimJax 아키텍처 및 핵심 기술

• JAX 기반 가속화: 전체 학습 파이프라인 (환경 스텝핑 및 정책 업데이트) 을 하드웨어 가속기 (GPU/TPU) 에서 실행되도록 설계했습니다.
  • JIT 컴파일 및 vmap: 수천 개의 병렬 환경 인스턴스를 벡터화하여 처리함으로써 CPU-GPU 통신 오버헤드를 제거하고, 기존 시뮬레이터 대비 최대 100 배의 처리량 (Throughput) 향상을 달성했습니다.
• 맥락적 POMDP (Contextual POMDP) 공식화:
  • 침투 테스트를 부분 관측 마코프 의사결정 과정 (POMDP) 으로 모델링하되, 각 에피소드가 특정 네트워크 인스턴스 (맥락, Context) 에 조건부로 주어지는 **맥락적 MDP (CMDP)**로 정의했습니다.
  • 이는 에이전트가 특정 네트워크를 암기하는 것이 아니라, 다양한 네트워크 구조 분포에서 작동하는 정책을 학습하여 Zero-shot 일반화를 가능하게 합니다.
• 네트워크 생성 파이프라인:
  • 구조적으로 다양하고, 현실적이며, 해결이 보장된 (Guaranteed-solvable) 시나리오를 생성하는 새로운 생성기를 도입했습니다.
  • 토폴로지 밀도, 취약점 밀도, 민감한 호스트 비율 등을 제어하여 학습 난이도를 조절할 수 있습니다.

2.2 대규모 행동 공간 처리: 2SAS (Two-Stage Action Selection)

네트워크 크기가 커질수록 행동 공간이 폭발적으로 증가하는 문제를 해결하기 위해 2 단계 행동 선택 (2SAS) 메커니즘을 제안했습니다.

• 기존 방식 (Flat Masking): 모든 가능한 (호스트, 행동) 조합을 하나의 거대한 행동 공간으로 취급합니다.
• 2SAS 방식: 의사결정을 두 단계로 분해합니다.
  1. 호스트 선택: 접근 가능한 미발견 호스트 중 하나를 선택.
  2. 행동 선택: 선택된 호스트에 대해 실행 가능한 행동 (스캔, 익스플로잇 등) 을 선택.
• 이 방식은 각 단계의 행동 공간 크기를 줄여 탐색 효율성을 극대화합니다.

2.3 학습 전략 및 보상 스케일링

• 보상 스케일링: 생성된 네트워크의 크기에 따라 보상의 절대값이 달라지는 문제를 해결하기 위해, 이론적 최대 보상으로 보상을 정규화하여 학습 신호의 안정성을 확보했습니다.
• UED (Unsupervised Environment Design) 비교:
  • DR (Domain Randomization): 매 에피소드마다 무작위 네트워크를 생성.
  • PLR (Prioritized Level Replay): 에이전트의 후회 (Regret) 가 큰 레벨을 우선적으로 재학습하여 자연스러운 커리큘럼을 형성.

3. 주요 기여 (Key Contributions)

1. NASimJax 프레임워크: 침투 테스트 RL 학습을 위한 첫 번째 순수 JAX 구현체로, 100 배의 속도 향상을 통해 대규모 실험을 가능하게 함.
2. 2SAS 메커니즘: 대규모 네트워크에서의 행동 공간 확장 문제를 해결하고, 평면적 행동 마스킹 (Flat Masking) 보다 월등히 우수한 성능을 입증.
3. 일반화 및 커리큘럼 학습 통찰:
   • 희소한 (Sparse) 토폴로지에서 학습하는 것이 밀집된 (Dense) 토폴로지보다 Zero-shot 일반화에 더 효과적임을 발견 (암묵적 커리큘럼 효과).
   • PLR 이 DR 보다 대규모 네트워크에서 더 강력한 일반화 성능을 보임.
4. 새로운 실패 모드 발견: PLR 의 에피소드 리셋 동작과 2SAS 의 신용 부여 (Credit Assignment) 구조 간의 상호작용으로 인해 발생하는 특정 조건에서의 학습 붕괴 현상을 규명하고 분석함.

4. 실험 결과 (Results)

• 성능 비교: 단일 엔트리 레벨 GPU 에서 **초당 160 만 스텝 (1.6M steps/sec)**을 처리하여 기존 NASim 대비 100 배 빠른 속도를 기록했습니다.
• 행동 공간 확장성 (40 호스트 네트워크):
  • 2SAS: 42% 해결률 (Success Rate) 달성.
  • Flat Masking: 14% 해결률에 그침.
  • 네트워크가 커질수록 2SAS 의 우위가 극명하게 나타났습니다.
• Zero-shot 일반화 (ZSPT):
  • 학습 분포의 영향: 훈련 시 낮은 밀도 (Sparse) 네트워크에서 학습한 정책이, 훈련 분포보다 더 밀집된 테스트 환경에서도 더 잘 일반화되었습니다.
  • PLR vs DR: PLR 기반 방법이 DR 보다 일관된 성능을 보였으며, 특히 PLR-2SAS 조합이 높은 밀도 환경에서 우수한 성능을 발휘했습니다.
• 실패 모드 분석: 40 호스트 네트워크의 특정 밀도 (td=0.15) 에서 PLR-2SAS 가 완전히 붕괴되는 현상이 관찰되었습니다. 이는 PLR 의 에피소드 초기화가 2SAS 의 단계별 학습을 방해하고, 성공적인 에피소드가 없어 리플레이 버퍼가 채워지지 않기 때문임이 규명되었습니다. (PLR⊥-2SAS 는 이를 우회하여 성공함).

5. 의의 및 결론 (Significance)

NASimJax 는 RL 기반 침투 테스트 연구의 새로운 기준을 제시합니다.

• 실용성: 기존에는 계산 비용으로 인해 불가능했던 대규모 네트워크 (40 개 이상 호스트) 에 대한 RL 실험을 가능하게 하여, 실제 사이버 보안 환경에 더 가까운 규모의 정책을 연구할 수 있는 토대를 마련했습니다.
• 연구 플랫폼: 고정된 벤치마크가 아닌, 연구자가 네트워크 복잡도, 행동 공간 크기, 토폴로지 다양성 등을 자유롭게 조절할 수 있는 구성 가능한 플랫폼을 제공합니다.
• 이론적 통찰: 대규모 행동 공간 처리를 위한 2SAS 와 같은 새로운 아키텍처의 필요성과, 훈련 분포 설계 (희소한 네트워크 학습) 가 일반화에 미치는 중요한 영향을 명확히 보여주었습니다.

결론적으로, 이 논문은 RL 기반 침투 테스트의 확장성과 일반화 문제를 해결하기 위한 컴퓨팅 인프라와 알고리즘적 혁신을 동시에 제시하며, 향후 더 복잡하고 현실적인 사이버 보안 자동화 연구의 기반이 될 것으로 기대됩니다.