Flux Attention: 효율적인 LLM 추론을 위한 문맥 인식형 하이브리드 어텐션

긴 문맥을 처리할 때 발생하는 연산량 폭증 문제를 해결하기 위해, 입력된 질문의 성격에 따라 각 레이어의 연산 방식을 실시간으로 결정하는 동적 시스템을 제안한다. 기존의 고정된 방식보다 효율적이면서도 정보 검색 능력을 유지하여 실제 하드웨어에서 체감할 수 있는 속도 향상을 이끌어냈다.

Transformer의 핵심인 Attention은 문장의 모든 단어를 서로 비교하므로 문장이 길어질수록 계산량이 제곱으로 늘어난다. 이를 해결하기 위해 일부만 계산하는 Sparse Attention이 등장했지만, 모든 상황에 일률적으로 적용하면 중요한 정보를 놓치는 '성능 붕괴'가 발생한다. 특히 특정 정보를 찾아내야 하는 검색(Retrieval) 작업은 모든 단어를 꼼꼼히 봐야 하는 반면, 전체적인 흐름을 파악하는 요약 작업은 듬성듬성 봐도 충분하다.

Flux Attention은 이 문제를 해결하기 위해 '레이어 라우터'라는 지능형 스위치를 도입한다. 이 스위치는 입력된 질문이 꼼꼼한 검색이 필요한지, 아니면 대략적인 파악으로 충분한지를 먼저 판단한다. 그 후, 모델의 각 층(Layer)에 '이 질문은 꼼꼼하게 계산해라' 혹은 '이 질문은 건너뛰며 계산해라'라는 명령을 실시간으로 내린다.

기존 연구들이 개별 연산 헤드 단위로 복잡하게 계산량을 조절하려다 오히려 하드웨어 정렬이 깨져 속도가 느려졌던 것과 달리, Flux Attention은 층 전체의 모드를 통째로 바꾼다. 이는 GPU가 메모리에서 데이터를 읽어올 때 끊김 없이 연속적으로 처리할 수 있게 하여, 이론적인 연산량 감소를 실제 추론 속도 향상으로 직결시킨다.

Flux Attention은 사전 학습된 LLM을 고정한 상태에서 가벼운 Layer Router를 추가하는 구조이다. 라우터는 입력 쿼리 텐서 xQ [s × h × d']를 입력받아 Prefill-Suffix Pooling을 수행한다. 이 과정에서 프롬프트의 시작과 끝부분 토큰 100개씩을 추출하여 [200 × h × d' → 1 × h × d']로 압축해 전체 문맥의 특징을 요약한다.

추출된 특징은 MLP 기반의 Context Encoder와 Router Head를 거쳐 각 레이어의 라우팅 로짓 πFA와 πSA를 생성한다. 학습 시에는 Gumbel-Softmax relaxation을 적용하여 r_soft = exp((πFA + gFA)/τ) / Σ exp((π_i + g_i)/τ)와 같이 미분 가능한 확률값을 계산하고, 이를 통해 Full Attention과 Sparse Attention의 출력을 가중합하여 학습한다. 추론 시에는 argmax를 통해 하나의 모드를 확정하는 Hard Routing을 수행한다.

학습 목적 함수는 언어 모델링 손실(L_language)과 희소성 규제 손실(L_diff)의 조합으로 구성된다. Lagrangian relaxation 기법을 사용하여 max_λ min_θ L_lang + λ₁L_diff + λ₂L_diff² 식을 최적화하며, 이를 통해 각 작업의 난이도에 맞는 최적의 희소성 수준을 모델이 스스로 찾아내도록 유도한다.

Qwen3-4B/8B 및 Llama-3.1-8B 모델을 대상으로 실험한 결과, LongBench-E 벤치마크에서 Full Attention 모델과 대등하거나 이를 상회하는 성능을 유지하면서도 상당한 효율성을 확보했다. 특히 검색 집약적인 작업에서 기존의 정적 희소성 모델들이 성능 붕괴를 겪는 것과 달리, Flux Attention은 높은 정확도를 유지했다.

추론 속도 측면에서는 256K 컨텍스트 길이 기준으로 Prefill 단계에서 최대 2.8배, Decode 단계에서 최대 2.0배의 속도 향상을 기록했다. 이는 레이어 단위 라우팅이 GPU의 실행 패턴과 잘 일치하여 동기화 지연(Synchronization long-tail)을 제거했기 때문이다.

RULER 벤치마크를 통한 문맥 확장 능력 테스트에서도 256K 범위까지 정보 검색 및 추론 능력을 안정적으로 유지함을 확인했다. 또한 수학적 추론(GSM8K, AIME24) 작업에서도 성능 저하 없이 효율적인 연산이 가능함을 입증했다.

Flux Attention의 아키텍처는 Frozen Backbone과 Trainable Layer Router의 분리 구조가 핵심이다. 라우터는 각 레이어의 어텐션 블록 직전에 위치하여 입력 쿼리를 가로채 의사결정을 내린다. 하드웨어 효율성을 극대화하기 위해 Sparse 모드 선택 시 해당 레이어의 KV 캐시 로딩을 완전히 우회(Bypass)하도록 설계되었다.

수학적으로는 Gumbel-Softmax의 온도 파라미터 τ를 학습 초기에는 높게 설정하여 탐색을 유도하고, 점진적으로 낮추어 추론 시의 Hard Routing과의 간극을 줄이는 Annealing 전략을 사용한다. 또한 Lagrangian multiplier λ를 작업별로 동적으로 조정하여, 검색 성능이 중요한 작업에서는 λ를 낮춰 Full Attention 사용을 장려하고 요약 등에서는 λ를 높여 희소성을 강제한다.

구현 측면에서는 Block-Sparse-Attention 커널을 활용하여 64 단위의 블록 희소성을 제어하며, 128 토큰 규모의 Sink Token을 유지하여 생성의 안정성을 보장한다. 이러한 설계는 이론적인 FLOPs 감소를 실제 Wall-clock time 단축으로 전환하는 데 결정적인 역할을 한다.

라우터의 성능이 입력 프롬프트의 처음과 마지막 100개 토큰에 크게 의존하므로, 중요한 지시 사항이 문서의 중간에 위치할 경우 라우팅 결정이 최적화되지 않을 가능성이 있다. 또한 가벼운 MLP 구조의 라우터는 복잡한 문맥 노이즈를 필터링하는 데 한계가 있어 풀링 윈도우를 과도하게 넓힐 경우 오히려 성능이 저하되는 현상이 관찰되었다.