ReWY가 델타룰 어텐션 역전파의 행렬 곱을 7개에서 3개로 줄인다
ReWY: Faster Gradient Computation for Delta-Rule Attention Family
무엇인가
Kimi Linear는 전체의 약 3분의 1을 Kimi Delta Attention(KDA)로 채우고, Qwen3-Next는 48개 레이어 중 36개에 Gated DeltaNet(GDN)을 넣으며 Qwen3.5-0.8B 텍스트 백본은 24개 중 18개에 GDN을 쓴다. 이 모델들은 같은 델타룰 연산자로 순환 행렬 상태를 갱신하는데, 청크 단위 구현은 행렬 병렬성을 노출하는 대신 역전파에서 최종 입력 기울기에 필요 없는 보조 adjunct를 계산한다. 기존 WY 분해는 W=AE, U=AZ를 만들어 상태 스캔에 쓰고, 역전파에서 W̄=-R̄S^T, Ā=R̄Z^T+W̄E^T, 그리고 역행렬 풀백 -A^TĀA^T를 차례로 만든다. 청크 64, 키/값 폭 128에서 이 국소 WY 그래프는 행렬 곱 7개와 청크·헤드당 734만 dense FLOPs를 쓴다.
어떻게 동작하나
ReWY의 출발점은 청크 보정이 M R = Z - E S 하나의 삼각 잔차 시스템을 만족한다는 관찰이다. 여기서 M=I+L, L은 E Y^T의 순엄격 하삼각이고 A는 M의 역행렬이다. 저자들은 선형 solve의 고전적 역미분을 이 합성 그래프에 적용해 공유 민감도 P=A^T R̄=M^{-T}R̄를 계산하면 역행렬 adjunct 연쇄가 통째로 불필요해진다는 것을 보인다. 이때 Z̄=P, Ē_dir=-P S^T, S̄_res=-E^T P, D=L̄=-L_<(P R^T), Ē_intra=D Y, Ȳ_res=D^T E가 나오고, 새로 필요한 곱은 A^T R̄, P R^T, P S^T 세 개뿐이다. 국소 WY 곱이 7개에서 3개로 줄고, 제거되는 dense 연산량은 ΔF=4n²d_k+4n³, n=64·d_k=d_v=128에서 734만에서 419만 FLOPs(42.9%)다.
무엇과 다른가
남은 상호작용에는 두 가지 재사용 변환이 붙는다. 하나는 벡터 게이트용 행렬 intra로, 토큰 쌍마다 exp(G_ik-G_jk)를 곱하던 pairwise 처리를 exp(G_ik)·exp(-G_jk)로 분리해 U_q=H_q Y, U_e=D Y, U_y=H_q^T Q_g + D^T E 같은 행렬 수축으로 바꾼다. 수치 안정성을 위해 특징별 중심 c_k*=(M_k+m_k)/2를 잡아 최대 로그 크기를 Δ_k/2로 최소화하고, γ와 γ^{-1} 대신 α=exp(G-c), ζ=exp(c-G)를 쓴다. 다른 하나는 스칼라 게이트 GDN용 Gram 제거로, 게이트 기울기가 K K^T의 가중 행·열 합만 필요하다는 점을 이용한다. A_K=ΓK, B_K=Γ^T diag(β)K로부터 Ḡ_i^(L)=β_i⟨K_i,(A_K)_i⟩-⟨K_i,(B_K)_i⟩, β̄_i^(L)=⟨K_i,(A_K)_i⟩를 얻어 약 2n²d_k FLOPs의 dense K K^T와 n×n 중간체를 없앤다. 두 변환 모두 ReWY가 만든 계수 adjunct D를 소비자 쪽에서 재사용한다.
어떻게 쓰나
실험은 RTX 5070Ti, A800-SXM4-80GB, Ascend 910C, H100 80GB HBM3, H200, Ascend 950DT에서 진행됐고, 베이스라인은 동결된 FLA 구현이다. (B,T,H) 6개 형상 (1,512,4), (1,2048,4), (1,8192,4), (1,2048,16), (2,2048,16), (1,8192,16)에 C=64, d_k=d_v=128을 쓰고, 재계산(R)과 전체 저장(A) 정책을 양쪽에 맞춰 비교했다. NVIDIA GPU에서 기하평균 완전 저수준 역전파 1.15~1.81배, H100/H200에서는 네 연산자 모두 1.18~1.70배(원래 FP32 작업 행렬 유지), Ascend 950DT 공개 autograd에서는 1.09~2.89배가 나왔다. DeltaNet은 1.1배 미만으로 남고 KDA의 잔차 단독 이득은 완만하며, GDN2가 가장 큰 이득을 보였다.
전제와 한계
컴포넌트 실험은 ReWY와 두 변환의 기여를 분리한다. GDN에서 Gram 제거 단독은 RTX 5070Ti R/A 1.02/1.03배, Ascend 910C 1.05/1.07배였고, Gram을 고정한 채 ReWY를 더하면 각각 1.05/1.07배, 1.07/1.09배가 됐다. 할당 실험에서는 GDN/NPU의 잔차 단독 구현이 64MiB, Gram 제거를 더하면 29MiB를 추가로 아꼈다. FLA 전체 모델 역전파(A800)에서는 KDA 10.04%, GDN 0.28%, DeltaNet 약 0.95%, GDN2 레이어 포트 7.92~8.55% 개선이 나왔고, 코어 역전파는 수정 전 FLA KDA 모델의 측정된 역전파 구간 중 25.7~31.3%를 차지한다. 수치 검증은 CPU FP64 순환 참조와 비교해 일반 출력·기울기 상대 L2 0.8%, 게이트·베타·erase/write 기울기 2%, 근접 영역 절대 10^-6 임계를 적용했다.
이 논문이 개발자에게 주는 메시지는 커널을 손보기 전에 미분 그래프부터 줄이라는 것이다. ReWY는 아키텍처나 목적함수를 바꾸지 않고 DeltaNet, GDN, KDA, GDN2 네 연산자에 그대로 적용되는 연산자 알고리즘 변경이며, FLA처럼 이미 커뮤니티가 쓰는 구현에 얹을 수 있다. 학습 파이프라인에서 델타룰 어텐션의 역전파가 병목인지, 재계산/저장 정책이 속도와 메모리에 어떤 상호작용을 만드는지 점검할 때 참고할 만하다.
저자들이 명시한 한계도 분명하다. exact arithmetic에서의 등가성이 동일한 부동소수점 궤적을 보장하지 않고, KDA 폴백에는 bitwise-determinism 제약이 있으며, Ascend 950DT 스위트는 FP16, 가변 길이, graph-replay, 학습 수렴을 검증하지 않는다. 950DT에서는 전체 모델 역전파 측정이 없어 연산자 비율로 모델 전체 속도를 추론하지 않는다. GDN2는 FLA의 공식 전체 모델이 아니라 레이어 포트이고, Kimi Linear나 Qwen3-Next 전체 배포를 측정한 것도 아니다. A800의 일부 DeltaNet 궤적 진단은 허용 임계를 넘었고, 연산자 검증 통과가 전체 모델 수치 수용이나 학습 수렴을 보장하지 않는다. 저자들은 또한 FlashQLA 개정판이 이미 잔차 우변을 합치지만 양측 역행렬 풀백은 유지하며, 이 논문이 그 구현과 성능을 비교하지 않는다고 밝힌다.