WiMamba: 선형 복잡도로 무선 채널의 미래를 열다

이 논문은 무선 통신 시스템, 특히 채널 상태 정보(CSI) 표현 학습을 위한 효율적인 기반 모델(FM)인 'WiMamba'를 소개한다. 기반 모델은 대규모 데이터로 사전 학습된 후 소량의 데이터만으로 다양한 다운스트림 작업(예: 빔형성 예측, 위치 추정, 채널 분류)에 적응할 수 있는 범용 표현을 학습하는 모델이다. 현재 무선 FM 연구는 두 갈래로 진행되어 왔다: 첫째, GPT나 비전 모델 등 범용 FM을 차용하는 방법(높은 복잡도와 부적합한 모달리티 문제 있음), 둘째, 무선 채널에 특화된 트랜스포머 기반 FM을 구축하는 방법(2차 복잡도로 인한 확장성 한계 있음). WiMamba는 후자의 접근법을 취하되, 복잡도 문제를 근본적으로 해결하기 위해 최근 등장한 선형 복잡도 시퀀스 모델 'Mamba'를 채택했다. 논문의 시스템 모델은 N개의 안테나와 M개의 서브캐리어를 가진 MIMO 채널 행렬 H를 입력으로, 이를 D차원의 잠재 표현 x로 매핑하는 함수 f_θ를 기반 모델로 정의한다. 이 모델은 대규모 무라벨 채널 데이터셋 D_pre로 사전 학습되며, 이후 제한된 라벨 데이터 D_p_task로 각 작업 p에 맞는 작은 헤드 네트워크만을 추가 학습(fine-tuning)하여 사용된다. 좋은 무선 FM의 요구사항으로 작업 불변성(R1), 다양한 채널 차원 처리 능력(R2), 낮은 지연/메모리 사용(R3)을 제시한다. WiMamba의 아키텍처는 다음과 같다. 먼저, 입력 채널 H의 실수/허수 성분을 벡터화하고, 미리 정의된 토큰 차원 L_e에 따라 패치로 나눠 토큰 시퀀스를 생성한다. 이때 '적응적 세분화'를 통해 L_e를 조절하면 토큰 시퀀스 길이 T_e가 변하여 복잡도와 표현 세부성을 유연하게 제어할 수 있다. 토큰은 공유 임베딩 레이어를 거쳐 D차원으로 투영된 후, 핵심인 '양방향 Mamba 백본'에 입력된다. 채널 데이터는 시간적 순서가 없으므로, 하나의 임베딩을 반으로 나누어 각각 순방향과 역방향 Mamba 레이어 스택(Q개)으로 처리한 후 결과를 결합한다. Mamba 레이어는 선형 복잡도의 선택적 SSM을 사용하여 트랜스포머의 어텐션을 대체한다. 사전 학습은 마스킹 언어 모델링(MLM) 방식의 자기 지도 학습으로 진행된다. 무라벨 채널 데이터를 토큰화한 후, 일부 토큰을 마스킹하고(실수/허수 쌍을 함께 마스킹), WiMamba 백본이 생성한 임베딩으로 마스킹된 토큰을 재구성하는 선형 디코더를 함께 학습한다. 이 손실 함수를 최소화함으로써 모델은 채널의 내재적 구조를 포착하는 풍부한 표현을 학습하게 된다. 추론 단계에서는 사전 학습된 WiMamba 백본의 가중치를 고정한 채, 다운스트림 작업을 수행한다. 백본이 생성하는 임베딩(클래스 토큰 x_cls, 평균 풀링 x_mean, 평탄화 x_flatten 중 선택)을 입력으로, 해당 작업에 대한 소규모 라벨 데이터로만 작은 헤드 네트워크(예: MLP)를 학습시킨다. 복잡도 분석에 따르면, WiMamba의 전체 추론 복잡도는 O(QD * ⌈2NM/L_e⌉)로 채널 크기 N, M에 대해 선형적으로 증가하며, 이는 트랜스포머의 O(T_e^2) 대비 엄청난 이점이다. 실험은 DeepMIMO 데이터셋을 사용하여 네 가지 다운스트림 작업(분류 및 회귀)에 대해 수행되었다. 결과는 WiMamba가 트랜스포머 기반 FM(예: Channel Transformer)과 비교하여 동등하거나 더 나은 정확도를 유지하면서, 추론 지연 시간과 GPU 메모리 사용량을 크게 절감했음을 보여준다. 이는 제안된 아키텍처가 요구사항 R1-R3을 모두 충족시키며, 대규모 실시간 무선 시스템에 FM을 적용하는 데 있어 실질적인 장벽을 해소했음을 의미한다.