기술 요약: MambaNetBurst

문제 제기

네트워크 트래픽 분류 및 침입 탐지는 사이버 보안에 필수적이지만, 최첨단 접근 방식들은 종종 복잡하고 자원을 많이 소모하는 파이프라인에 의존합니다. 현재 주류 방법들은 일반적으로 다음을 포함합니다:

1. 무거운 사전 학습 (Heavy Pre-training): 대규모 데이터셋에서 미세 조정 (fine-tuning) 전에 자기 지도 학습 목표 (예: 마스크된 언어 모델링, 다음 토큰 예측) 를 활용하는 것으로, 비용이 많이 들고 시간이 오래 걸립니다.
2. 조기 집계 (Early Aggregation): 긴 입력 시퀀스를 효율적으로 처리하기 위해 기존 모델들은 종종 토큰화, 패치링 (patching), 또는 스트라이딩 (striding) 을 적용합니다. 이 다운샘플링은 정확한 위협 탐지에 필수적인 세밀한 바이트 수준의 구조적 정보 (예: 프로토콜 헤더, 페이로드 서명) 를 희석시킬 수 있습니다.
3. 아키텍처 비효율성: Transformer 기반 아키텍처는 강력하지만, 이차적인 어텐션 복잡성을 겪거나 선형 어텐션 변형에서도 무거운 피드포워드 레이어를 필요로 하여, 원시 바이트 시퀀스에 대해 계산 비용이 많이 듭니다.

본 논문은 이러한 복잡성들이 불필요할 수 있다고 주장합니다. 특히 효율적인 상태 공간 모델 (SSM) 을 활용할 때, 원시 패킷 바이트를 직접 모델링하는 것이 사전 학습, 엔지니어링된 표현, 또는 조기 다운샘플링의 필요성을 제거할 수 있는지 질문합니다.

방법론: MambaNetBurst

저자들은 MambaNetBurst를 제안합니다. 이는 Mamba-2 백본을 기반으로 구축된 컴팩트한, 토크나이저가 없는 바이트 수준 시퀀스 분류기입니다. 아키텍처는 다음과 같이 작동합니다:

• 입력 구성: 모델은 첫 $n$개의 패킷 (기본값 $n=5$) 을 선택하여 네트워크 플로우를 처리합니다. 각 패킷에서 첫 $m$개의 바이트 (기본값 $m=320$) 를 유지하여 1600 바이트의 고정 길이 버스트를 형성합니다. 이는 패킷 경계를 보존하고 정보량이 풍부한 헤더와 초기 페이로드에 초점을 맞춥니다.
• 토크나이저 없는 임베딩: 바이트를 서브워드 토큰으로 매핑하는 방법과 달리, MambaNetBurst 는 원시 바이트 값 (0–255) 을 직접 학습 가능한 임베딩 벡터 ($d_{model}=256$) 로 매핑합니다. 경량 2 층 MLP 가 시퀀스 모델링 전에 이러한 임베딩을 투영하여 로컬 특징을 풍부하게 합니다.
• 시퀀스 모델링: 학습 가능한 [CLS] 토큰과 위치 임베딩이 추가된 바이트 시퀀스는 잔차 사전 정규화된 Mamba-2 블록의 스택에 의해 처리됩니다.
  • Mamba-2 세부 사항: 모델은 구조화된 상태 공간 이중성 (SSD) 을 활용하여 핵심 계산을 GEMM 친화적 연산으로 재구성합니다. 중요한 점은 Mamba-2 가 Mamba-1 에서 사용된 유연한 대각 행렬 대신 제약된 전이 행렬 ($A = \alpha I$, 스칼라 곱하기 단위 행렬) 을 사용한다는 것입니다.
• 분류: [CLS] 토큰의 최종 은닉 상태는 선형 레이어와 소프트맥스를 거쳐 엔드 - 투 - 엔드 지도 학습 분류에 전달됩니다.
• 사전 학습 부재: 모델은 자기 지도 사전 학습 단계를 완전히 우회하여 순수하게 지도 교차 엔트로피 손실로 훈련됩니다.

주요 기여

1. 최초의 직접 바이트 - 분류 Mamba-2: MambaNetBurst 는 토큰화, 패치링, 또는 사전 학습 없이 원시 네트워크 패킷 바이트에 Mamba-2 를 직접 적용하여 트래픽 분류를 수행한 첫 번째 연구입니다.
2. 제약된 역학의 검증: 논문은 Mamba-2 의 제한된 전이 구조 (스칼라 $\times$ 단위 행렬) 가 패킷 바이트 모델링에 충분할 뿐만 아니라, 더 유연한 Mamba-1 에 비해 데이터셋 간 일관성을 향상시키는 유익한 정규화제로 작용함을 입증합니다.
3. 사전 학습 제거: 이 연구는 구별력 있는 트래픽 표현이 완전한 지도 학습 설정에서 원시 바이트로부터 직접 학습될 수 있음을 증명하여, 문헌에서 주류를 이루는 비용이 많이 드는 2 단계 (사전 학습 + 미세 조정) 파이프라인의 필요성을 제거합니다.
4. 포괄적인 애블레이션: 저자들은 임베딩 전략, 위치 인코딩, 스트라이딩, 모델 깊이, 상태 크기 등에 대한 광범위한 분석을 제공하여, 바이트 수준의 시간 해상도를 보존하는 것이 중요하며 중간 정도의 상태 크기 ($d_{state}=16-64$) 가 충분함을 밝힙니다.

실험 결과

모델은 암호화된 모바일 앱 식별, VPN/Tor 트래픽, 멀웨어 분류, IoT 공격을 다루는 6 개의 공개 벤치마크에서 평가되었습니다.

• 성능: MambaNetBurst 는 일관되게 강력한 결과를 달성하여, 종종 훨씬 더 무겁고 사전 학습된 베이스라인 (예: ET-BERT, YaTC, NetMamba) 을 능가하거나 그에 필적했습니다.
  • ISCXTor2016에서 0.9990 의 매크로 -F1 을 달성하여 NetMamba (0.9986) 와 ET-BERT (0.9967) 를 능가했습니다.
  • USTC-TFC2016(멀웨어) 에서 0.9954 의 매크로 -F1 에 도달하여 NetMamba (0.9957) 와 ET-BERT (0.9910) 를 능가했습니다.
  • CrossPlatform(Android/iOS) 에서 각각 0.9824 와 0.9851 의 매크로 -F1 점수를 달성하여 사전 학습된 Transformer 들과 경쟁하거나 우위였습니다.
• 효율성:
  • 훈련 속도: Mamba-2 는 Mamba-1 에 비해 역전파 (backward passes) 에서 30–60% 속도 향상을 보였으며, FlashAttention-2 를 사용하는 선형 Transformer 보다 훨씬 빠른 추론 속도를 보여주었습니다.
  • 메모리: Transformer 변형보다 2–4 배 적은 GPU 메모리를 사용하여 상용 하드웨어 (예: RTX 3090) 에서 더 큰 배치 크기를 가능하게 했습니다.
• 애블레이션 통찰:
  • 스트라이딩은 해롭다: 조기 다운샘플링 (스트라이딩) 이 가장 심각한 성능 저하를 초래하여, 세밀한 바이트 순서가 중요함을 확인시켰습니다.
  • Mamba-2 대 Mamba-1: Mamba-2 는 데이터셋 전반에서 더 나은 평균 성능과 더 낮은 분산을 보여주어, 그 제약들이 암묵적 정규화제로 작용함을 시사합니다.
  • 깊이와 크기: 중간 정도의 깊이 (4 레이어) 와 상태 크기 (16–64) 가 충분했으며, 상태 크기를 128 로 증가시키면 성능이 저하되어 이 모달리티에는 거대한 잠재 메모리가 불필요함을 나타냈습니다.

중요성과 주장

본 논문은 MambaNetBurst가 네트워크 트래픽 분석을 위한 새로운 실용적 패러다임을 확립한다고 주장합니다:

• 직접 학습의 실현 가능성: 고비용의 사전 학습 파이프라인이 최첨단 트래픽 분류에 필수적이라는 가정에 도전합니다. 컴팩트한 SSM 을 통한 직접적이고 희석되지 않은 바이트 - 분류 학습이 매우 효과적임이 입증되었습니다.
• 효율성과 배포 가능성: 사전 학습을 피하고 Mamba-2 의 하드웨어 효율성을 활용함으로써, 이 접근법은 특히 자원 제약 환경이나 빈번한 재훈련이 필요한 시나리오 (예: NIDS) 에 더 간단하고 빠르며 배포 가능한 솔루션을 제공합니다.
• 바이트 수준 모델링의 견고성: 결과는 잠재 역동 시스템의 유연성을 극대화하거나 데이터를 조기에 집계하는 것보다 패킷 바이트의 네이티브 해상도를 보존하는 것이 더 중요함을 시사합니다.

저자들은 이 작업이 사이버 보안 분야에서 바이트 수준 SSM 을 위한 유망한 방향을 열어주며, 현재 해당 분야를 지배하는 복잡하고 사전 학습에 의존하는 파이프라인에 대한 더 간단한 대안을 제공한다고 결론지었습니다.