A Unified Adaptive Feature Composition Framework for Multi-Task Generalization in Wireless Foundation Models

문제 정의

무선 파운데이션 모델(Wireless Foundation Models, WFMs)은 보편적인 채널 표현을 학습하는 데 잠재력을 보여주었으나, 현재의 패러다임 하에서는 다양한 다운스트림 태스크로의 적응 과정에서 상당한 병목 현상에 직면해 있습니다. 기존의 접근 방식은 크게 두 가지 범주로 나뉘며, 두 방식 모두 한계를 가지고 있습니다:

1. 전체 미세 조정(Full Fine-tuning): 유연성은 높지만, 상당한 계산 및 저장 오버헤드를 초라하여 이질적인 태스크들을 위한 배포와 유지보수를 복잡하게 만듭니다.
2. 고정된 특징 추출(Frozen Feature Extraction): 이 패러다임은 일반적으로 단일 고정 표현(주로 마지막 레이어의 출력)에 의존합니다. 그러나 이는 성능 병목을 초래하는데, 서로 다른 무선 태스크들은 서로 다른 수준의 추상화를 요구하기 때문입니다. 예를 들어, 채널 추정(Channel Estimation)은 미세한 시공간-주파수 및 다중 경로 구조(저수준 특징)를 필요로 하는 반면, 빔 관리(Beam Management) 및 위치 추정(Localization)은 전역적인 공간 및 기하학적 패턴(고수준 특징)을 필요로 합니다.

본 논문은 마지막 레이어에만 의존하는 것이 특징 과도 평활화(Feature Over-smoothing) 문제를 일으킨다고 주장합니다. 즉, 깊은 층의 표현이 특정 태스크에 필요한 미세한 세부 사항을 잃게 된다는 것입니다. 또한, 자연어 처리(NLP) 및 컴퓨터 비전에서 가져온 기존의 다층 융합 기술들은 너무 무겁거나(과도한 파라미터 도입) 단일 태스크에 밀접하게 결합되어 있어, WFM을 위한 통합적이고 경량화된 인터페이스를 제공하지 못합니다.

방법론

이러한 문제를 해결하기 위해, 저자들은 **라우팅 어댑터 기반 특징 구성(Routing Adapter for Feature Composition, RAFC)**이라는 경량 모듈을 중심으로 한 **통합 적응형 특징 구성 프레임워크(Unified Adaptive Feature Composition Framework)**를 제안합니다.

1. 핵심 개념: 적응형 다층 집계(Adaptive Multi-Level Aggregation)
마지막 레이어의 출력만을 추출하는 대신, 이 프레임워크는 서로 다른 트랜스포머 깊이(depth)의 은닉 상태(hidden states)를 재사용 가능한 다수준 특징 풀(pool)로 취급합니다. RAFC 모듈은 다운스트림 태스크의 구체적인 요구 사항에 따라 이러한 계층적 표현을 동적으로 집계합니다.

2. RAFC 아키텍처
RAFC 모듈은 세 가지 순차적인 단계를 거치는 모델 불가지론적(model-agnostic) 플러그인으로 작동합니다:

• 레이어 압축 기술(Layer Compact Description): 전체 토큰 시퀀스를 처리할 때 발생하는 이차적인 계산 비용을 피하기 위해, 선택된 레이어의 토큰 시퀀스는 먼저 토큰별 평균 풀링(token-wise mean pooling)을 통해 압축됩니다. 이를 통해 압축된 전역 레이어 기술자(global layer descriptors)가 생성됩니다.
• 학습 가능한 가중치 할당(Learnable Weights Allocation): 이 기술자들은 서로 연결되어 경량 병목 멀티 레이어 퍼셉트론(MLP)에 입력됩니다. MLP는 특징을 투영하고, ReLU 활성화 함수를 적용한 뒤, 다시 확장하여 라우팅 가중치($\alpha$) 세트를 생성합니다. 소프트맥스 정규화를 통해 이 가중치들의 합은 1이 됩니다.
• 가중 합산(Weighted Summation): 원래의 계층적 특징 텐서들은 생성된 가중치를 통해 샘플 적응형 가중 합산 방식으로 융합됩니다. 이를 통해 국부적 구조와 전역적 의미론의 균형을 맞춘 합성 표현 $z^{(q)}$가 생성됩니다.

3. 레이어 선택 전략
프레임워크는 중복성과 오버헤드를 피하기 위해 모든 레이어를 사용하지 않습니다. 대신, 저-, 중-, 고수준 특징을 포착하기 위해 대표적인 레이어 집합 $S = \{2, \lfloor N/2 \rfloor + 1, N\}$을 선택합니다.

• 왜 레이어 2인가? 저자들은 레이어 1이 관측 노이즈(높은 섭동 비율)에 매우 민감하다는 것을 경험적으로 입증했습니다. 반면 레이어 2는 저수준 노이즈를 효과적으로 억제하면서도 미세한 물리적 변동을 유지하므로, 최적의 저수준 인터페이스 역할을 합니다.
• 왜 모든 레이어를 사용하지 않는가? 균등 샘플링은 표현의 다양성과 배포 효율성 사이의 실용적인 절충안을 제공합니다.

4. 학습 패러다임
사전 학습된 WFM 백본은 엄격하게 고정된 상태를 유지합니다. 오직 태스크 특화형 RAFC 모듈과 다운스트림 태스크 헤드의 파라미터만이 최적화됩니다. 이는 프레임워크가 무시할 만한 수준의 추가 파라미터(약 0.045M)만을 도입하면서도, 파운데이션 모델의 저지연 이점을 보존하도록 보장합니다.

주요 기여

• 통합 프레임워크: 단일 레이어 추출에서 동적 다층 집계로 적응 패러다임을 전환하여, 깊은 WFM에서의 특징 과도 평활화 문제를 완화하는 모델 불가지론적 특징 구성 프레임워크를 제시합니다.
• 경량 설계: RAFC 모듈은 토큰별 압축 레이어와 작은 MLP로만 구성되어, 무시할 만한 수준의 계산 및 저장 오버헤드(<50K 추가 파라미터)를 도입합니다.
• 해석 가능성: 학습된 라우팅 가중치는 물리적 해석 가능성을 제공하며, 서로 다른 태스크가 본질적으로 서로 다른 물리적 표현 수준(예: 재구성을 위한 미세한 세부 사항 vs 공간 인지를 위한 전역적 추상화)을 우선시한다는 점을 밝혀냅니다.

실험 결과

프레임워크는 사전 학습된 트랜스포머 기반 WFM(8개 레이어, 6.62M 파라미터)을 사용하여 네 가지 대표적인 다운스트림 태스크(채널 추정, 채널 예측, 빔 예측, 도시 지역 위치 추정)에 대해 평가되었습니다.

• 성능 향상:
  • 채널 추정: 마지막 레이어 전용 베이스라인 대비 NMSE를 38% 감소시켰습니다.
  • 채널 예측: NMSE를 80% 감소시켰습니다.
  • 빔 예측: Top-1 정확도를 3% 향상시켰습니다.
  • 도시 지역 위치 추정: 절대 오차를 34% 줄였습니다.
• 효율성: 이 방법은 전체 파라미터 크기에 약 0.7%만을 추가하며, 지연 시간 증가도 매우 미미합니다 (0.0002 ms ~ 0.0011 ms).
• 전이 가능성: 프레임워크는 다른 백본인 LWMv1.1에도 성공적으로 적용되어 일관된 성능 향상을 보였으며, 이는 특정 파운데이션 아키텍처로부터의 결합 해제(decoupling)를 증명합니다.
• 절제 연구(Ablation Studies):
  • 특정 레이어(2, 5, 8)를 선택하는 것이 성능-오버헤드 측면에서 저수준만 사용하거나, 고수준만 사용하거나, 혹은 모든 레이어를 사용하는 것보다 우수했습니다.
  • 평균 풀링(Mean pooling)이 어텐션 기반 또는 전체 토큰 라우팅에 비해 최소한의 FLOPs로 최대치에 근접한 성능을 제공하는 최적의 집계 전략임을 확인했습니다.

의의 및 주장

본 논문은 RAFC가 WFM을 다양한 다운스트림 시나리오에 적응시키기 위한 저복잡도, 확장 가능성 및 설명 가능한 인터페이스를 제공한다고 주장합니다.

• 태스크별 선호도: 학습된 가중치 분석에 따르면, 신호 재구성 태스크(추정/예측)는 "거친 단계에서 미세한 단계로(coarse-to-fine)" 학습 궤적을 보이며, 결국 미세한 물리적 변동을 복구하기 위해 저수준 특징에 크게 의존합니다. 반면, 공간 인지 태스크(빔 예측/위치 추정)는 전역적인 기하학적 구조를 포착하기 위해 중·고수준의 추상화에 더 많은 비중을 둡니다.
• 동적 적응성: 라우팅 메커니즘은 정적이지 않고 입력 조건에 따라 적응합니다. 예를 들어, 비가시선(NLoS) 환경에서는 라우터가 거시적 구조를 활용하기 위해 더 깊은 레이어로 가중치를 이동시키며, 단말 속도가 높을 경우 순간적인 신호 변동을 포착하기 위해 얕은 레이어로 가중치를 이동시킵니다.
• 실제 배포: 전체 미세 조정이나 무거운 교차 어텐션(cross-attention) 메커니즘을 피함으로써, 이 프레임워크는 무선 시스템의 엄격한 경량화 및 저지연 요구 사항을 해결하여, 성능 저하 없이 이질적인 태스크 전반에 걸쳐 사전 학습된 WFM을 효율적으로 재사용할 수 있게 합니다.