하이브리드 어텐션 모델로의 변형

Transformer의 full (softmax) attention은 시퀀스 길이에 따라 연산이 𝒪(T²)로 증가하고 autoregressive 추론에서 KV cache가 길이에 비례해 커져 긴 문맥 확장이 비효율적이다. 순수 선형 어텐션은 계산·메모리 면에서 유리하지만 long-range recall 성능에서 full attention을 대체하기에 부족한 측면이 존재한다. FlashMorph는 사전학습된 Transformer를 효율적 하이브리드 구조로 변환하면서 레이어 선택 비용을 크게 줄이고 긴 문맥 검색 능력을 보존하는 실용적 통로를 제공한다.

Transformer의 self-attention은 입력 시퀀스 내 모든 토큰 쌍을 비교해 가중합을 계산하므로 시퀀스 길이 T에 대해 연산 복잡도가 𝒪(T²)이며 autoregressive 환경에서는 KV cache가 길이에 비례해 증가한다. 이로 인해 매우 긴 문맥을 다루는 추론에서 메모리와 시간 비용이 급증해 실무적 한계가 발생한다. 반대로 linear attention은 쿼리·키에 feature map을 적용하고 누적 키-값 통계를 recurrent state로 유지해 연산을 𝒪(T)로 줄이고 고정 크기 상태 캐시를 가능하게 만든다. 그러나 linear attention은 전역적 상호작용을 softmax로 직접 모델링하지 않으므로 검색·리콜 성능에서 full attention 대비 열세가 생기기 쉽다. FlashMorph는 이 두 접근의 장점을 취합하기 위해 사전학습된 Transformer의 각 레이어에 학습된 linear 대체 브랜치를 붙이고 레이어별 스칼라 게이트를 도입해 두 브랜치 출력을 보간한다. 구체적으로 게이트 α(l)이 1이면 해당 레이어는 full-attention 출력을, 0이면 linear-attention 출력을 따르며 중간값은 혼합 표현을 생성한다. 이때 모델 본체와 linear 브랜치 파라미터는 고정하고 오직 게이트만 합성된 장기 검색 신호로 학습해 어떤 레이어가 전역 구성 하에서 full attention을 유지해야 하는지를 공동으로 평가한다. 공동 최적화는 개별 레이어의 중요도를 독립적으로 평가하는 기존 방식과 달리 보존 레이어들 사이의 보완성·중복을 포착해 동일 예산 K 내에서 더 효율적이고 효과적인 하이브리드 구성을 찾아낸다.

전체 파이프라인은 세 단계로 나뉜다. 첫 번째 단계에서는 pretrained full-attention 모델을 교사로 삼아 각 레이어별 linear-attention 브랜치를 hidden-state L2 정렬 손실로 학습해 all-linear 학생을 얻고, 이렇게 학습된 브랜치를 원본 full-attention 레이어와 쌍으로 묶어 morphable model을 구성했다. 두 번째 단계에서는 morphable model의 모든 파라미터를 고정하고 레이어별 게이트 α(l)만 최적화해 합성 검색 데이터에서 정답 토큰 위치의 hidden-state 정렬 손실(L_align)과 선형화 규제(L_reg=Σα)를 동시에 최소화했다. 이 목적은 L_total = L_align + λ L_reg으로 표현되며 논문 기본값으로 λ=0.1을 사용했다. 게이트는 초기값 α(l)=1으로 시작해 최적화 후 값이 큰 상위 K개 레이어를 선택해 이산적인 하이브리드 구조를 만든다. 마지막 단계에서는 선택된 하이브리드 모델을 교사의 로그릿스 분포를 KL로 맞추는 logits distillation과 장문 맥락에서의 language-model finetuning으로 미세조정해 성능을 회복했다. 구현상 실험은 Qwen3-0.6B와 Qwen3-1.7B 백본을 사용했고 Lightning Attention, GLA, GDN 같은 선형 어텐션 변형을 평가했으며 모든 선택 단계는 모델 파라미터를 고정한 상태에서 게이트만 업데이트해 선택 비용을 최소화했다.

Needle-in-a-Haystack 실험에서 FlashMorph는 20M 토큰의 레이어 선택 예산으로 0.6B와 1.7B 백본에서 전반적으로 강한 복원 성능을 보였다. 특히 1.7B 백본에서는 NIAH-Single-1에 대해 모든 문맥 길이에서 완전한 정확도를 달성했고 NIAH-Single-2·3에서도 기존 방법 대비 유의미한 개선이 관찰되었다. 이 결과는 매우 적은 선택 예산으로도 joint gate 최적화가 장거리 검색에 핵심적인 레이어들을 잘 식별함을 의미한다. Commonsense 및 리콜 집약적 태스크 평가에서 FlashMorph는 Lightning/GLA/GDN 백본 전반에 걸쳐 원본 Transformer의 추론 능력을 대체할 수 있는 수준을 유지했고, 특히 리콜 관련 평균 점수에서 우위를 보였다. 효율성 측면에서 FlashMorph 하이브리드는 prefilling에서 128K 길이에서 2.24배, 256K에서 2.81배의 지연 시간 개선을 달성했으며 decoding에서는 256K에서 1.56배, 512K에서 2.07배의 개선을 보였다. 또한 FlashMorph는 동일 환경에서 더 적은 GPU 메모리를 사용해 512K·1M 토큰 수준의 실행을 가능하게 했고, 레이어 선택 비용은 Qwen3-1.7B 기준 FLOPs 2.5×10^17 및 2.1 GPU시간으로 보고되어 기존 방법들보다 수십에서 수천 배 효율적이었다.

전체 아키텍처는 각 full-attention 레이어에 학습된 linear-attention 브랜치를 결합해 morphable layer를 구성하고 게이트 α(l)을 통해 두 브랜치 출력을 보간하는 구조로 설계되었다. 혼합 출력은 H_mix(l)=α(l)·H_full(l)+(1−α(l))·H_lin(l)로 계산되며 최종 이산화는 TopK(α)로 이루어져 선택된 K 레이어만 full attention을 유지한다. 레이어 선택 단계에서는 원본 full-attention 모델과 linear 대체 브랜치의 파라미터를 고정하고 오직 게이트만 최적화해 selection overhead를 극소화했다. 수식적으로 hidden-state alignment 손실 L_hidden = (1/L)Σ_l ||H_lin(l)−H_full(l)||_2^2로 linear 브랜치를 초기 학습하고, 선택 단계에서는 정답 토큰 위치 기반 정렬 손실 L_align과 선형화 규제 L_reg=Σ_l α(l)을 합쳐 L_total=L_align+λL_reg을 최소화해 게이트를 학습했다. 이 구성은 개별 레이어의 중요도를 독립적으로 측정하는 기존 layerwise 복원·교체 방법과 달리 레이어들 간의 상호작용과 중복을 반영해 전역 최적화를 수행한다. 구현상 실험은 8×8 GPUs, BFloat16 정밀도로 수행되었고 하이브리드 기본 비율은 3:1로 설정했으며 λ의 기본값은 0.1로 보고되었다.