ThreatVisionAI: A Hybrid CNN-ViT Framework for Image-Based Malware Classification

기술 요약: ThreatVisionAI

문제 정의

전통적인 악성코드 탐지 방식, 특히 시그니처 기반 시스템은 난독화되거나 다형성을 가진 제로데이 위협에 대해 일반화하는 데 어려움을 겪습니다. 휴리스틱 접근 방식은 더 넓은 범위를 커버하지만, 종종 높은 오탐률(false-positive rates) 문제를 야기합니다. 이미지 기반 악성코드 분류는 바이너리 파일을 시각적 표현으로 변환하여 구조적 관계를 포착함으로써 유망한 대안으로 부상했습니다. 그러나 기존의 딥러닝 프레임워크는 다음과 같은 중대한 한계에 직면해 있습니다:

1. 주파수 정보의 손실: 표준 합성곱 신경망(CNN) 및 비전 트랜스포머(ViT)는 일반적으로 픽셀 공간에서만 작동합니다. 풀링(pooling)과 스트라이드 컨볼루션(strided convolutions)을 통해, 이들은 서로 밀접하게 연관된 악성코드 패밀리(예: 시각적으로 유사한 변종인 Swizzor.gen!E와 Swizzor.gen!I)를 구별하는 데 필수적인 고주파 질감(high-frequency textural) 세부 정보를 버립니다.
2. 구조적 격차: 하이브리드 모델(예: LeViT-MC, ConvNeXt–Swin)은 CNN과 트랜스포머 브랜치를 결합하지만, 전용 주파수 영역 브랜치가 결여되어 있습니다. 결과적으로 방향성 질감 차이를 활용하지 못합니다.
3. 해석 가능성: 많은 AI 기반 프레임워크는 모델 결정에 대한 명확한 설명을 제공하지 못하며, 분석가들에게 왜 특정 오분류가 발생하는지에 대한 통찰력을 제공하지 못합니다.

방법론: ThreatVisionAI 프레임워크

저자들은 공간적, 주파수 영역 및 전역적 관계 특징을 보완적으로 포착하도록 설계된 하이브리드 분류 프레임워크인 ThreatVisionAI를 제안합니다. 이 아키텍처는 세 개의 병렬 브랜치로 구성되며, 각 브랜치는 독립적으로 학습된 후 가중 소프트 보팅(weighted soft voting)을 통해 융합됩니다.

1. 원본 이미지 CNN 브랜치 (Raw-Image CNN Branch)

• 입력: 바이트 수준의 시각화를 나타내는 그레이스케일 악성코드 이미지(224 × 224)를 직접 처리합니다.
• 아키텍처: 저수준 에지(edge)부터 패밀리 특유의 구조적 시그니처까지 공간적 특징을 학습하기 위해, ImageNet 사전 학습 없이 처음부터 학습된 ResNet-18 백본을 사용합니다.
• 기능: 바이너리의 시각적 구조에 내재된 국소적 공간 패턴을 포착합니다.

2. 웨이브릿 기반 CNN 브랜치 (Wavelet-Based CNN Branch)

• 입력: 단일 레벨의 2D Haar 웨이브릿 분해를 거친 동일 이미지를 처리합니다. 이는 224 × 224 이미지를 네 개의 112 × 112 계수 서브밴드로 변환합니다:
  • 근사(cA): 저주파 콘텐츠.
  • 세부(cH, cV, cD): 수평, 수직, 대각선 방향의 고주파 변화.
• 아키텍처: 4채널 입력을 받도록 수정된 ResNet-18 백본을 사용합니다.
• 기능: 픽셀 공간에서는 보이지 않는 다중 스케일 주파수 구조와 방향성 질감 단서를 추출하여, 전역적 외형이 유사한 패밀리 간의 구분을 돕습니다.

3. 비전 트랜스포머 (ViT) 브랜치

• 입력: 웨이브릿 CNN과 동일한 4채널 웨이브릿 변환 입력을 받습니다.
• 아키텍처: timm을 통해 ImageNet-21k로 사전 학습된 ViT-Tiny 모델을 사용합니다. 패치 투영 레이어는 4채널을 수용하도록 적응되었습니다.
• 기능: 셀프 어텐션 메커니즘을 사용하여 이미지 전체의 장거리 의존성을 모델링하며, 국소적 합성곱 필터가 놓치는 전역적 문맥을 포착합니다.

융합 전략 (Fusion Strategy)

• 가중 소프트 보팅 (Weighted Soft Voting): 학습된 특징 결합 대신, 각 브랜치의 25개 클래스 확률 벡터를 가중 합산하여 결합합니다.
• 가중치: 검증 세트에 대한 철저한 그리드 탐색을 통해 결정되었습니다: 0.50 (Raw CNN), 0.40 (Wavelet CNN), 0.10 (ViT).
• 근거: 이 방식은 추가적인 학습 파라미터를 요구하지 않고, 브랜치 기여도의 투명성을 유지하며, 작은 데이터셋에 대한 과적합을 완화합니다.

실험 설정 및 결과

프레임워크는 심각한 클래스 불균형이 존재하는 25개 악성코드 패밀리와 9,465개의 그레이스케일 이미지를 포함하는 Malimg 데이터셋을 통해 평가되었습니다. 데이터는 70/15/15(학습/검증/테스트) 비율로 분할되었습니다.

성능 지표

• 전체 하이브리드 정확도: 98.01%
• 가중 F1 점수 (Weighted F1 Score): 0.9742
• 절제 연구 통찰 (Ablation Insights):
  • 웨이브릿 CNN 단독 성능(0.9733 F1)이 Raw CNN(0.9662 F1)보다 우수하여, 주파수 영역 특징의 가치를 입증했습니다.
  • ViT 브랜치 단독 성능(0.9492 F1)은 낮았으나, CNN 앙상블에 추가되었을 때 일관되고 반복 가능한 이득을 제공하며 보완적인 전역 문맥을 기여했습니다.
  • 하이브리드 모델은 2개 브랜치 구성인 CNN+ViT 앙상블(0.9692 F1)보다 뛰어난 성능을 보였으며, 이는 주파수 영역 특징이 이 데이터셋에서 매우 중요함을 확인시켜 줍니다.

구체적 개선 사항

• Swizzor 변종: 웨이브릿 브랜치는 시각적으로 유사한 패밀리 간의 구별 능력을 크게 향 향상시켰습니다. Swizzor.gen!E의 F1은 Raw CNN 대비 0.7368에서 0.8571로, Swizzor.gen!I는 0.5263에서 0.8276으로 증가했습니다.
• 실패 분석 (Grad-CAM): 모델은 Autorun.K를 Yuner.A로 지속적으로 오분류했습니다 (F1 = 0.0000). Grad-CAM 히트맵은 두 패밀리의 시각적 구조와 활성화 패턴이 인간 전문가에게도 거의 동일하다는 것을 보여주었습니다. 저자들은 이것이 모델의 결함이 아닌 근본적인 데이터셋의 한계라고 결론지었습니다.

강건성 및 효율성

• 적대적 강건성 (Adversarial Robustness): FGSM 화이트박스 공격 하에서 정확도가 급격히 하락하여 (98.01%에서 $\epsilon=0.05$일 때 7.84%로), 적대적 섭동에 취약함을 나타냈습니다.
• 효율성: 전체 하이브리드 모델은 27.91M 파라미터를 포함하며, GPU에서 이미지당 5.4 ms의 추론 시간을 달성하여 비실시간 스캐닝 워크플로우에 운영상 수용 가능한 수준임을 보여주었습니다.

주요 기여 및 의의

본 논문은 다음과 같은 기여를 주장합니다:

1. 새로운 하이브리드 아키텍처: 공간적 CNN 및 전역 문맥 ViT와 함께 전용 주파수 영역(웨이브릿) 브랜치를 통합한 최초의 악성코드 분류 프레임워크인 ThreatVisionAI를 도입했습니다.
2. 주파수 특징의 가치 입증: 절제 연구를 통해 웨이브릿 기반 특징이 공간 필터만으로는 분리할 수 없는 밀접하게 연관된 악성코드 패밀리를 구별하는 데 측정 가능한 이득을 제공함을 확인했습니다.
3. 해석 가능성: Grad-CAM 적용을 통해 체계적인 오분류(특히 Autorun.K vs. Yuner.A)가 모델의 결함이 아닌 내재적인 클래스 간 시각적 유사성에 의해 발생한다는 시각적 증거를 제공하여, 분석가들에게 모델의 한계를 명확히 이해할 수 있게 했습니다.

한계 및 향후 과제

저자들은 다음과 같은 한계를 명시적으로 인정합니다:

• 데이터셋 범위: 평가는 2011년에 수집된 Malimg 데이터셋에 국한되어 있어, 현대의 악성코드 생태계의 규모와 다양성을 반영하지 못합니다. MaleVis나 BIG2015와 같은 벤치마크에 대한 일반화 여부는 검증되지 않았습니다.
• 적대적 취약성: 모델은 화이트박스 공격에 취격하며, PGD와 같은 더 강력한 공격이나 데이터셋 변화(dataset shifts)에 대한 테스트가 수행되지 않았습니다.
• 향후 방향: 저자들은 다중 데이터셋 평가, 적대적 학습, 대안적인 웨이브릿 기저(wavelet bases) 탐색, 그리고 멀티모달 프레임워크를 위한 동적 행동 신호의 통합을 향후 과제로 제안합니다.