TabNSA: 효율적인 테이블 데이터 학습을 위한 네이티브 희소 어텐션

초록

TabNSA는 테이블 데이터의 고차원·이질적 특성을 고려해, 특성을 토큰으로 취급하고 계층적 희소 어텐션(Native Sparse Attention)을 적용한다. 압축·선택·슬라이딩 윈도우 세 단계로 불필요한 특성을 제거해 O(N²) 복잡도를 크게 낮추고, 독립 파라미터를 가진 MLP‑Mixer인 TabMixer와 결합해 전역·국부 상호작용을 동시에 모델링한다. 또한 사전학습된 LLM(Gemma)을 활용해 Few‑Shot 및 전이 학습 성능을 강화한다. 실험 결과, 기존 딥러닝 기반 테이블 모델과 GBDT를 모두 능가하는 정확도와 효율성을 입증하였다.

상세 분석

TabNSA는 테이블 데이터의 특성을 ‘토큰’으로 재해석함으로써 시퀀스‑기반 어텐션 메커니즘을 그대로 적용할 수 있는 기반을 마련한다. 기존의 전통적인 어텐션은 입력 차원이 N일 때 O(N²)의 연산 비용을 요구하지만, TabNSA는 Native Sparse Attention(NSA)의 세 가지 핵심 단계—Token Compression, Token Selection, Sliding Window—를 계층적으로 배치해 복잡도를 크게 감소시킨다.

1. Token Compression 단계에서는 연속적인 특성 블록을 하나의 압축 토큰으로 집계한다. 이때 블록 길이(l)와 스트라이드(d)를 조절해 지역적 상관관계를 유지하면서 차원을 축소한다. 압축 과정은 위치 인코딩을 포함한 MLP fₖᶜᵐᵖ 로 구현되어, 각 블록 내의 중요 정보를 보존한다.

2. Token Selection 단계는 압축 토큰에 대한 어텐션 스코어를 Softmax로 정규화한 뒤, 상위 k 개의 블록을 선택한다. 선택된 블록은 원본 특성 차원으로 복원(concatenation)되어, 모델이 인스턴스별로 가장 의미 있는 특성 집합에 집중하도록 만든다. 이 과정은 순위 기반(rank) 필터링을 통해 동적 sparsity를 구현하므로, 데이터마다 다른 특성 조합을 학습한다.

3. Sliding Window 단계는 최근 w 개의 토큰을 고정‑크기 윈도우로 유지해 지역적 패턴을 포착한다. 이는 장거리 의존성을 완전히 포기하지 않으면서도 연산량을 제한하는 트레이드오프를 제공한다.

희소 어텐션의 출력은 3가지 브랜치(압축, 선택, 윈도우) 각각에 대해 별도의 어텐션을 수행하고, 게이트 g₍c₎(sigmoid + MLP)로 가중합한다. 이때 게이트는 입력 특성에 기반해 각 브랜치의 기여도를 동적으로 조절한다.

TabNSA는 이러한 희소 어텐션 모듈을 TabMixer 백본과 결합한다. TabMixer는 기존 MLP‑Mixer와 달리 채널‑와 토큰‑별로 독립적인 파라미터를 갖는 두 개의 병렬 MLP를 사용한다. SiLU와 GeLU 활성화 함수를 연속 적용해 비선형성을 강화하고, Residual 연결을 통해 학습 안정성을 확보한다. 결과적으로 TabMixer는 전역적인 특성‑특성 상호작용과 샘플‑샘플 간 관계를 동시에 모델링한다.

또한 TabNSA는 Gemma와 같은 사전학습된 대형 언어 모델(LLM)을 파인‑튜닝해 Few‑Shot 학습 능력을 확장한다. 구조화된 테이블 데이터를 텍스트 프롬프트로 변환한 뒤, Gemma의 디코더‑전용 트랜스포머를 고정하고 마지막 레이어만 미세조정한다. 이렇게 얻은 텍스트 임베딩은 Dual‑Pooling 후 선형 변환을 거쳐 TabNSA의 최종 분류기에 통합된다. LLM의 풍부한 세계 지식과 추론 능력이 테이블 데이터의 제한된 라벨 상황에서도 일반화 성능을 크게 끌어올린다.

실험에서는 12개 이상의 공개 테이블 벤치마크(예: Adult, Higgs, YearPrediction 등)와 다양한 데이터 규모(수천~수십만 샘플)에서 TabNSA가 기존 SOTA 딥러닝 모델(TabNet, SAINT, FT‑Transformer 등) 및 GBDT(XGBoost, LightGBM)를 일관되게 앞선다. 특히 저‑샘플(≤ 1 %) 및 전이 학습 시, 희소 어텐션이 불필요한 특성을 걸러내어 과적합을 방지하고, LLM 보조가 극히 적은 라벨에서도 높은 정확도를 유지한다는 점이 강조된다.

요약하면, TabNSA는 (1) 계층적 희소 어텐션을 통해 연산 효율성을 극대화하고, (2) TabMixer의 강력한 비선형 믹싱으로 복합 상호작용을 포착하며, (3) LLM 기반 프롬프트 학습으로 Few‑Shot 능력을 보강한다는 세 축을 성공적으로 결합한 새로운 테이블 데이터 학습 프레임워크라 할 수 있다.