본문으로 건너뛰기

Griffin의 절충안: Gated Linear Recurrence와 Local Attention을 함께 쓰는 법

· 약 6분
p4r4d0xb0x
p4r4d0xb0x
Rustacean, AI, OSS Enthusiast

선형 recurrence와 지역 어텐션이 교차하는 시퀀스 모델 구조

Transformer는 모든 위치를 직접 비교하는 능력으로 언어 모델링의 표준이 됐지만, 생성 시 과거의 키와 값을 계속 보관해야 한다. 반대로 전통적인 순환신경망(RNN)은 토큰마다 상태 하나를 갱신하므로 추론 비용이 작지만, 시간축 병렬화와 대규모 학습이 어렵다. Griffin 논문은 이 대립을 “둘 중 하나를 고르는 문제”로 다루지 않는다. Hawk라는 gated linear recurrence와, Hawk에 local attention을 섞은 Griffin을 함께 제안해 서로 다른 시간 범위의 기억을 분담시킨다.

논문의 초록은 Hawk가 보고된 Mamba 성능을 넘고, Griffin이 Llama-2와 경쟁적인 성능을 6배 이상 적은 학습 토큰으로 달성했다고 보고한다. 또한 학습 중 Transformer와 비슷한 하드웨어 효율, 추론 중 낮은 지연과 높은 처리량, 학습 길이보다 긴 시퀀스에 대한 extrapolation을 주장한다. 다만 이 글은 초록과 원문 실험 조건을 근거로 쓰며, 외부 재현 결과나 새로운 비교를 추가하지 않는다.

Hawk: RNN의 상태를 게이트로 다시 설계하기​

선형 recurrence는 현재 입력과 이전 상태의 선형 결합으로 이해할 수 있다. 간단히 쓰면 h_t = a_t ⊙ h_{t-1} + b_t ⊙ x_t와 같은 형태다. 여기서 게이트는 과거를 얼마나 지우고 새 입력을 얼마나 쓸지 조절한다. Griffin 논문의 Hawk는 이 아이디어를 언어 모델 블록 안에 넣되, 입력 의존적 게이트와 안정적인 정규화를 사용해 학습 가능성을 높인다.

이 설계의 양면성은 분명하다. 생성 시에는 현재 토큰과 고정 크기 상태만 읽고 쓰므로 시퀀스 전체를 다시 순회할 필요가 없다. 반면 시간축 recurrence는 본질적으로 순서 의존적이어서 학습 병렬화가 어려워질 수 있다. 논문은 gated linear recurrence를 parallel scan으로 계산할 수 있는 형태로 정리하고, 실제 모델이 Transformer와 비슷한 학습 하드웨어 효율을 보이도록 구현·구조를 함께 설계한다.

용어 해설: 게이티드 선형 recurrence

쉬운 정의: 과거 상태와 새 입력을 더할 때, 학습된 게이트가 각각의 비중을 조절하는 순환 계산이다.

예시: 냉장고에 남은 음식을 기록할 때 오래된 메모를 전부 버리거나 전부 보존하지 않고, 유통기한이 지난 항목은 지우고 새로 산 항목은 크게 반영하는 방식이다.

용어 해설: parallel scan

쉬운 정의: 순서가 있는 누적 계산을 여러 묶음으로 나누어 병렬 처리한 뒤, 묶음 사이의 요약 상태를 연결하는 알고리즘이다.

예시: 긴 줄의 합계를 한 사람이 처음부터 더하는 대신 구간별 합을 여러 사람이 계산하고 마지막에 합치는 것과 같다. 단, 각 구간의 결과를 다음 구간 계산에 전달해야 한다.

Griffin: 최근 기억은 attention, 긴 흐름은 recurrence​

Griffin은 Hawk만 쌓은 모델이 아니다. 원문은 recurrent block과 local multi-query attention(MQA) block을 교대하는 구조를 설명한다. 구체적으로 recurrent block 두 개 뒤에 local attention을 사용하는 residual block을 배치하며, 기본 local attention window는 1024 tokens로 제시된다. 최근 1024토큰 안의 세밀한 관계는 어텐션이 직접 보고, 그보다 긴 흐름은 recurrence의 상태가 압축해 전달하는 그림이다.

이 조합은 단순한 평균이 아니다. local attention은 전역 KV cache보다 저장 범위가 작고, recurrence는 창 밖의 정보를 고정 상태로 넘긴다. 따라서 “긴 문맥을 전부 attention으로 본다”와 “모든 정보를 하나의 벡터로 압축한다” 사이의 설계 공간이 생긴다. 하지만 어떤 정보가 창 안에 남고 어떤 정보가 상태로 요약되는지는 학습 결과에 달려 있다. 긴 문서의 임의 위치를 정확히 인용해야 하는 시스템이라면 별도 평가가 필요하다.

논문은 residual pattern과 MLP block을 Transformer baseline과 공유한다고 설명한다. 이 점은 비교를 해석할 때 중요하다. 성능 차이가 recurrence와 local attention에서 왔더라도, 전체 블록·학습 레시피·데이터 구성이 동일한지는 표와 실험 절을 함께 확인해야 한다.

논문이 보고한 실험을 읽는 법​

원문은 100M부터 7B까지 모델 크기를 바꾸고, 14B Griffin도 추가해 scaling curve를 제시한다. MassiveText 데이터셋, 2048 토큰 시퀀스, AdamW가 사용됐다고 기술한다. validation loss와 학습 FLOPs 관계에서 Griffin이 여러 예산 구간에서 자체 Transformer baseline보다 낮은 loss를 보였다고 보고하며, Hawk는 상대적으로 격차가 있으나 큰 예산에서 좁혀진다고 설명한다.

downstream 비교에서는 자체 MQA Transformer baseline과 Mamba-3B, Llama-2를 사용한다. Hawk-3B가 보고된 Mamba-3B보다 강한 결과를 보였고, Griffin-7B와 14B가 Llama-2와 경쟁적인 평균을 보였다는 것이 논문의 요지다. 하지만 외부 baseline은 서로 다른 데이터와 토큰 수로 학습됐다고 원문이 직접 주의를 준다. 특히 Llama-2와 Mamba가 Griffin보다 훨씬 많은 토큰을 보았다는 점은 결과를 읽을 때 반드시 함께 적어야 한다.

시스템 관점: 빠른 한 토큰이 전체 서비스 비용을 결정한다​

추론에서 Transformer의 비용은 context와 KV cache에 따라 커진다. Griffin의 recurrent block은 고정 크기 상태를 사용하고 local attention은 제한된 창을 사용하므로, 긴 시퀀스에서 메모리와 지연을 줄일 가능성이 있다. 논문은 낮은 latency와 높은 throughput을 보고하지만, 실제 값과 유리한 조건은 하드웨어·배치·시퀀스 길이·커널에 의존한다. 따라서 운영 환경에서는 평균 지연뿐 아니라 첫 토큰 지연, 토큰 간 지연, 동시성별 처리량, 상태 메모리까지 측정해야 한다.

학습 측면에서는 parallel scan과 local attention 모두 GPU 병렬성을 활용할 수 있다. 분산 학습에서는 14B 모델을 확장하고 shard 방법을 설명했다는 점이 실무적인 관심사다. 그렇다고 기존 Transformer 학습 스택을 그대로 재사용할 수 있다는 뜻은 아니다. recurrence state의 체크포인트, sequence parallel 통신, variable-length 배치 처리, fused kernel 지원을 구현 단계에서 검증해야 한다.

논문 구조 분석​

논문은 Introduction에서 RNN의 추론 장점과 학습 난점, Transformer의 표현력과 긴 시퀀스 비용을 대비한다. Background에서 선형 recurrence와 local attention의 구성 요소를 설명한 뒤 Hawk를 만들고, 두 블록을 섞어 Griffin을 정의한다. 이후 scaling, downstream, 긴 시퀀스 extrapolation, 속도·메모리 평가로 가는 문제→설계→검증 순서다. 결론은 모델이 경쟁적임을 강조하지만, 구조적 선택이 어떤 작업에 유리한지에 대한 해석은 실험 조건과 함께 읽어야 한다.

한계와 도입 전 체크리스트​

첫째, local window가 최근 정보를 우선한다는 편향은 장거리의 정확한 회수 능력과 별개다. 둘째, recurrence 상태는 작지만 정보 압축의 손실 경로를 갖는다. 셋째, 논문의 강한 baseline 비교는 학습 토큰 수와 데이터셋이 같지 않은 경우가 있으므로 절대적인 우열로 번역하면 안 된다. 넷째, 논문이 보고한 속도 이점이 모든 추론 엔진에서 자동으로 나타나는 것은 아니다.

실험한다면 다음을 같은 조건으로 비교하는 것이 좋다.

  • 문맥 길이와 batch size별 prefill·decode 지연
  • 사용자 수 증가에 따른 상태·KV 메모리 피크
  • 최근 구간 질문과 장거리 associative recall을 분리한 정확도
  • 동일 토큰 예산·데이터·튜닝 예산에서의 perplexity와 downstream 점수
  • 체크포인트 복구와 스트리밍 입력에서 recurrent state의 일관성

Griffin의 메시지는 RNN으로 돌아가자는 것이 아니다. 긴 범위를 요약하는 recurrence와 가까운 범위의 정확한 비교를 담당하는 attention을 한 모델 안에 배치해, 알고리즘적 효율과 표현력 사이의 경계를 조정하자는 제안이다. 제공된 자료에서 확인되는 성능은 유망하지만, 실제 도입 전에는 원문·코드·벤치마크를 확인해야 한다. 특히 서비스의 문맥 접근 패턴이 1024토큰 창과 상태 압축 가정에 맞는지가 먼저다.

Sources​

// COMMENTS

댓글