동적 순위 강화학습을 이용한 적응형 저차원 다중 헤드 셀프 어텐션

초록

DR‑RL은 강화학습 에이전트를 활용해 입력 시퀀스와 레이어 특성에 따라 실시간으로 어텐션 행렬의 저차원(rank)을 조절한다. 온라인 행렬 섭동 이론을 기반으로 안전한 rank 변화를 보장하고, 배치 SVD와 경량 Transformer 정책망을 결합해 FLOPs를 40 % 이상 절감하면서도 SST‑2 등 GLUE 벤치마크에서 풀‑랭크와 동등한 정확도를 유지한다.

상세 분석

본 논문은 대규모 언어 모델(Large Language Model, LLM)의 핵심 연산인 Multi‑Head Self‑Attention(MHSA)의 계산 복잡도를 저차원 근사화로 경감시키면서, 정적인 rank 설정이 초래하는 표현력 손실을 동적으로 보완하는 새로운 프레임워크인 Dynamic Rank Reinforcement Learning(DR‑RL)을 제안한다. 핵심 아이디어는 rank 선택을 “시퀀스‑레벨 정책”으로 모델링하여, 강화학습 에이전트가 각 레이어·헤드·토큰 구간마다 최적의 rank r을 실시간으로 결정하도록 하는 것이다. 이를 위해 다음과 같은 기술적 토대를 마련한다.

1. MDP 정의와 보상 설계: 상태 s는 현재 입력 토큰의 통계(예: 토큰 다양성, 평균 attention entropy)와 레이어·헤드 인덱스를 포함한다. 행동 a는 선택 가능한 rank 집합(예: {16,32,64,128}) 중 하나이며, 보상 R는 두 항목의 가중합으로 구성된다. 첫 번째 항은 섭동 이론에 기반한 근사 오차 감소량(σ_{r+1}·‖V‖)이며, 두 번째 항은 FLOPs 절감 비율에 대한 페널티이다. 이렇게 하면 에이전트는 “정확도‑효율성” 트레이드‑오프를 직접 학습한다.

2. 온라인 행렬 섭동 이론: 저차원 근사 A_r ≈ U_r V_rᵀ에서 rank를 r→r′(r′>r) 로 증가시키는 경우, 섭동 Δ = A_{r′}−A_r 의 Frobenius norm은 Σ_{k=r+1}^{r′}σ_k² 로 정확히 계산된다. 논문은 이 식을 이용해 rank 변동이 attention 출력 Y에 미치는 상한 ‖Y_{r′}−Y_r‖ ≤ σ_{r+1}‖V‖ 를 추정하고, 이를 보상 함수에 직접 삽입한다. 따라서 에이전트는 “안전 구역(safe region)”을 벗어나지 않는 rank 선택만을 수행한다는 이론적 보장을 얻는다.

3. 경량 정책 네트워크: 정책망은 2‑layer Transformer 인코더(각 레이어 64‑dim hidden)와 작은 MLP 헤드로 구성돼, 입력 상태를 빠르게 인코딩한다. 파라미터 수가 수천 개 수준에 불과해, 실제 inference 단계에서 추가 연산 비용이 무시할 정도이다.

4. 배치 부분 SVD: 전체 attention 행렬에 대해 매번 완전 SVD를 수행하면 O(n³) 비용이 발생한다. 저자는 배치 Partial SVD(Top‑k Lanczos) 알고리즘을 활용해 각 헤드당 O(n²·r) 비용으로 상위 r개의 특이값·벡터만 추출한다. 이 과정은 GPU‑friendly하게 구현돼, rank 업데이트가 실시간으로 이루어질 수 있다.

5. 실험 결과: L > 4096인 긴 시퀀스 환경에서 DR‑RL은 평균 FLOPs를 42 % 절감하면서, perplexity와 GLUE 전반(특히 SST‑2 92.78 % 정확도)에서 풀‑랭크 Transformer와 통계적으로 유의미한 차이가 없음을 보였다. 또한 기존 정적 저차원 방법(Performer, Nyströmformer) 대비 8‑12 % 정도의 정확도 향상을 기록했다.

6. 한계와 향후 과제: 현재 정책망은 각 토큰 단위가 아니라 배치 단위로 rank을 결정한다는 점에서 미세한 토큰‑레벨 최적화가 제한된다. 또한, 하드웨어‑특화 최적화(예: FPGA, ASIC)와의 연동을 위해 rank 선택을 더 세분화하거나, 연산량을 직접 측정하는 메타‑러닝 기법을 도입할 여지가 있다.

전반적으로 DR‑RL은 “동적 차원 축소”라는 새로운 패러다임을 제시하며, 강화학습과 선형대수 이론을 결합한 설계가 LLM 추론 효율성을 크게 향상시킬 수 있음을 실증한다.