기술 요약: MARGIN – 불균형 취약점 탐지를 위한 마진 인식 정규화 기하학

1. 문제 제기

본 논문은 실제 시나리오에서 소프트웨어 취약점 탐지의 중요한 과제를 다루며, 여기서 딥러닝 모델은 두 가지 서로 다르지만 누적되는 형태의 데이터 불균형에 직면합니다:

• 빈도 불균형: 취약점 데이터셋은 심각한 장꼬리 분포를 보입니다. 소수의 공통 취약점 분류 (CWE) 카테고리 (예: CWE-79, CWE-89) 가 데이터를 지배하는 반면, 많은 다른 카테고리 (예: CWE-400, CWE-98) 는 희귀합니다. 표준 손실 함수는 모든 샘플을 동등하게 취급하여 모델이 다수 클래스로 편향되고 소수 취약점을 간과하게 만듭니다.
• 난이도 불균형: 빈도 외에도 서로 다른 CWE 카테고리는 고유한 학습 난이도 수준을 갖습니다. 높은 구조적 다양성을 가진 취약점 (예: 스택/힙 오버플로우 및 오프-바이-원 오류를 포함하는 CWE-119) 은 특징 표현에서 높은 클래스 내 분산을 나타내어, 반복적인 패턴을 가진 취약점 (예: CWE-415) 보다 학습하기 어렵습니다.

저자들은 기존 방법들이 주로 (재샘플링 또는 재가중치를 통해) 빈도 불균형만 다루지만 난이도 불균형은 고려하지 못한다고 주장합니다. 그들은 이러한 문제들을 임베딩 기하학적 관점에서 재해석하여, 두 가지 불균형 모두 초구면 표현 공간에서 기하학적 왜곡을 유발한다고 관찰합니다. 구체적으로, 소수 또는 난이도가 높은 클래스는 그 프로토타입 주변으로 빽빽하게 수렴하지 못하여, 확률 질량 영역이 다른 클래스의 보로노이 셀 (Voronoi cells) 로 침입하게 됩니다. 이 "왜곡 영역"은 오분류와 불안정한 결정 경계를 초래합니다.

2. 방법론: MARGIN

제안된 프레임워크인 MARGIN(불균형 취약점 탐지를 위한 마진 인식 정규화 기하학, Margin-Aware Regularized Geometry for Imbalanced Vulnerability DetectioN) 은 판별력 있는 취약점 표현을 학습하기 위해 단위 초구면에서 작동하는 메트릭 기반 접근법입니다. 이는 적응형 마진 학습과 초구면 프로토타입 모델링을 통합합니다.

핵심 구성 요소:

1. 초구면 임베딩 공간:
   이 방법은 특징과 분류기 가중치가 $\ell_2$-정규화되어 단위 초구면 $S^{d-1}$ 위에 놓이도록 코사인 소프트맥스 손실 (Cosine Softmax Loss) 을 활용합니다. 클래스는 프로토타입으로 표현되며, 목표는 해당 프로토타입과 샘플 임베딩을 정렬하면서 각도적 분리를 유지하는 것입니다.

2. vMF 분포를 통한 기하학적 모델링:
   클래스 분포의 집중도를 정량화하기 위해 저자들은 폰 미세스 - 피셔 (von Mises–Fisher, vMF) 분포를 사용하여 임베딩을 모델링합니다. 집중도 파라미터 $\kappa_i$는 샘플이 클래스 평균 주변에 얼마나 빽빽하게 군집하는지 추정합니다.

   • 보로노이 셀: 이상적인 균형 조건 (Neural Collapse) 하에서 클래스들은 등각 조밀 프레임 (Equiangular Tight Frame, ETF) 구조를 형성하여 구를 동등한 보로노이 셀로 분할합니다.
   • 왜곡 영역: 불균형 설정에서 클래스의 확률 질량 (vMF 꼭짓점 각도 $\theta^{vMF}_i$로 정의됨) 은 종종 해당 보로노이 셀 ($\theta^{cell}$) 을 초과합니다. 이 질량이 다른 클래스의 결정 영역으로 침입하는 영역이 "왜곡 영역"을 구성하며, 이는 직접적으로 거짓 양성 (False Positives) 과 거짓 음성 (False Negatives) 을 초래합니다.

3. 적응형 마진 메커니즘:
   왜곡 영역을 줄이기 위해 MARGIN 은 각 클래스에 대해 동적 각도 마진 $m_i$를 도입합니다.

   • 마진은 클래스의 vMF 꼭짓점 각도와 이상적인 보로노이 꼭짓점 각도 간의 차이를 기반으로 계산됩니다: $m_i = \max(\frac{\theta^{vMF}_i - \theta^{cell}}{2}, \dots)$.
   • 이 메커니즘은 제약을 동적으로 조정합니다: 높은 분산 (높은 난이도 또는 낮은 빈도) 을 가진 클래스는 확률 질량을 보로노이 셀 내부로 되돌리기 위해 더 큰 마진을 받으며, 컴팩트한 클래스는 더 작은 마진을 받습니다.

4. 집중도 인식 로짓 스케일링:
   프레임워크는 클래스 집중도 $\kappa_i$에 적응된 로짓을 위한 스케일링 인자 $s_i$를 포함합니다.

   • 낮은 집중도 (높은 분산, 학습이 어려움) 를 가진 클래스는 결정 경계를 날카롭게 하고 기울기 신호를 증폭시키기 위해 더 큰 스케일링 인자를 할당받습니다.
   • 높은 집중도를 가진 클래스는 훈련 과정을 지배하는 것을 방지하기 위해 더 작은 인자를 할당받습니다.
   • 이는 로그 - 집중도 공간에서의 소프트 순위 전략을 통해 달성됩니다.

5. 기하학적 중앙값을 통한 추론:
   추론 단계에서 MARGIN 은 학습된 가중치 벡터에만 의존하는 대신, 각 클래스에 대한 정규화된 훈련 임베딩의 기하학적 중앙값을 계산하여 데이터 기반 프로토타입을 구성합니다. 이는 경험적 특징 분포를 더 잘 반영하는 더 견고한 프로토타입을 제공합니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 제시합니다:

1. 불균형의 기하학적 재해석: 저자들은 빈도 및 난이도 불균형을 재검토하여 임베딩 기하학의 렌즈를 통해 그 효과를 분석합니다. 그들은 이러한 불균형이 클래스 조건부 확률 질량 영역과 이상적인 보로노이 결정 영역 사이의 불일치로 나타난다는 것을 증명합니다.
2. MARGIN 프레임워크: 그들은 두 가지 불균형을 동시에 완화하는 통합 메트릭 학습 프레임워크를 제안합니다. 적응형 마진 (확률 질량을 보로노이 셀과 정렬하기 위해) 과 집중도 인식 스케일링 (기울기 기여도를 균형시키기 위해) 을 통합함으로써, MARGIN 은 복잡한 재샘플링 없이도 결정 경계를 안정화합니다.
3. 실증적 검증: 공개 데이터셋 (BigVul, MegaVul, ReposVul) 에 대한 광범위한 실험은 MARGIN 이 그래프 기반 방법 (ReVeal, LIVABLE) 과 최신 표현 학습 모델 (HCL-VC, CLeVeR) 을 포함한 강력한 기준선보다 일관되게 우월함을 보여줍니다. 특히 세분화된 CWE 분류와 희귀 클래스에서 두드러집니다.

4. 실험 결과

평가 규모와 장꼬리 분포가 다양한 세 가지 데이터셋에서 수행되었습니다. 주요 발견 사항은 다음과 같습니다:

• 성능 향상: MARGIN 은 모든 데이터셋에서 최첨단 (state-of-the-art) 결과를 달성했습니다. BigVul에서 이진 F1 점수는 86.57% 였으며, CWE 분류에 대한 매크로 F1 은 70.72% 였습니다. 이는 다음으로 가장 좋은 기준선 (CLeVeR) 보다 매크로 F1 에서 7% 포인트 이상 우월한 성과입니다.
• 희귀 클래스에 대한 강건성: 심각한 불균형을 보이는 MegaVul 데이터셋에서 MARGIN 은 이진 F1 51.06% 와 매크로 F1 34.65% 를 달성하여 희귀 CWE 유형에 어려움을 겪은 기준선들을 크게 능가했습니다.
• 절대 실험 (Ablation Studies): 적응형 마진 또는 적응형 로짓 스케일링을 제거하면 성능이 저하되어 두 메커니즘 모두 필수적임을 확인했습니다. 마진은 특히 정밀도와 재현율 사이의 균형을 개선하는 반면, 스케일링은 소수 클래스에 대한 학습 역학을 안정화했습니다.
• 기하학적 분석: 사례 연구는 MARGIN 이 "왜곡 영역" (FPR/FNR 동역학을 통해 측정됨) 을 성공적으로 줄이고, 개선된 클러스터링 지표 (NMI, ARI) 와 감소된 그람 행렬 조건수 (Gram matrix condition number) 로 입증된 바와 같이 임베딩 공간을 안정적이고 잘 분리된 ETF 구조로 유도함을 보여주었습니다.

5. 중요성과 주장

본 논문은 MARGIN 이 취약점 탐지에서 불균형에 대한 원칙적인 기하학적 해석을 제공한다고 주장합니다. 임베딩 공간을 명시적으로 모델링하고 확률 질량을 결정 영역과 정렬하도록 기하학을 정규화함으로써, 이 방법은 분류 정확도뿐만 아니라 모델의 강건성과 해석 가능성도 향상시킵니다.

저자들은 그들의 접근법이 데이터 재가중치를 통한 증상 치료보다는 불균형 설정에서 편향의 근본 원인인 임베딩 공간의 기하학적 왜곡을 해결한다고 강조합니다. 그들은 현재 평가가 데이터 가용성으로 인해 C/C++ 데이터셋으로 제한되어 있지만, 기하학적 프레임워크는 백본에 구애받지 않으며 이론적으로 다른 언어에도 적용 가능하다고 지적합니다. 이 작업은 실용적인 불균형 처리와 이론적 매니폴드 학습 사이의 간극을 메워 보안이 중요한 취약점 탐지 시스템에 더 안정적인 기반을 제공합니다.