Griffin
Griffin : 효율적인 언어 모델을 위해 Gated Linear Recurrence와 Local Attention의 혼합
arXiv
모델 아키텍처
우리의 모든 모델에는 (i) 잔차 블록, (ii) MLP 블록 및 (iii) 시간 혼합 블록과 같은 구성 요소가 포함되어 있습니다. (i)와 (ii)는 모든 모델에서 동일하지만 글로벌 MQA(Multi-Query Attention), 로컬(슬라이딩 윈도우) MQA 및 제안된 반복 블록의 세 가지 시간적 혼합 블록을 고려합니다. 순환 블록의 일부로 우리는 선형 순환 단위 Orvieto et al., 2023b에서 영감을 받은 새로운 순환 계층인 RG-LRU(Real-Gated Linear Recurrent Unit)를 사용합니다.
그림 2(a)에 표시된 잔차 블록은 모델의 전역 구조를 정의하며 사전 표준 Transformers(Xiong et al., 2020)에서 영감을 받았습니다. 입력 시퀀스를 삽입한 후 이를 전달합니다. $N$ 그러한 블록( $N$ 모델 깊이를 나타냄) 그런 다음 RMSNorm Zhang 및 Sennrich, 2019를 적용하여 최종 활성화를 생성합니다. 토큰 확률을 계산하기 위해 최종 선형 레이어와 소프트맥스를 적용합니다. 이 레이어의 가중치는 입력 임베딩 레이어와 공유됩니다.
잔여 블록
그림 2: a) 모드 아키텍처의 주요 백본은 누적된 잔여 블록입니다. $N$ 타임스. b) 우리가 사용하는 게이트 MLP 블록. c) MQA(Multi Query Attention)의 대안으로 제안하는 순환 블록입니다. 이는 섹션 2.4에 정의된 제안된 RG-LRU 계층을 사용합니다.
잔차 블록에는 순서대로 적용되는 두 가지 구성 요소가 포함됩니다. 첫 번째 구성 요소는 숨겨진 상태를 취합니다. $chi$ RMSNorm Zhang 및 Sennrich, 2019를 적용한 다음 시간 혼합 블록을 적용합니다. 그런 다음 출력을 건너뛰기 연결과 병합합니다. $chi$ 추가를 통해. 마찬가지로 두 번째 구성 요소는 RMSNorm을 적용한 다음 MLP 블록을 적용한 다음 해당 출력을 RMSNorm 입력의 건너뛰기 연결과 병합합니다. 이 블록은 그림 2(a)에 나와 있습니다.
MLP 블록
우리는 차원 입력으로부터 두 개의 분기를 생성하는 게이트 MLP 블록 Dauphin et al., 2017(그림 2(b) 참조)을 사용합니다. $D$ . 출력 차원이 있는 선형 레이어를 적용합니다. $MD$ 각 지점에서 $M$ 확장 계수를 나타냅니다. 단순화를 위해 우리는 $M=3$ 이 작업 내내. GeGeLU Shazeer, 2020과 유사하게 요소별 곱셈을 통해 병합하기 전에 분기 중 하나에 GeLU 비선형성 Hendrycks 및 Gimpel(2016)을 적용합니다. 그러나 MLP 블록에서는 출력 차원이 있는 최종 선형 레이어를 적용합니다. $D$ GeGeLU 레이어의 출력에.
시간적 혼합 블록
시간 혼합 블록은 시퀀스의 서로 다른 시간 위치에서 숨겨진 레이어 활성화를 집계하는 모델의 구성 요소입니다. 우리는 글로벌 MQA Shazeer, 2019, 로컬 MQA Beltagy et al., 2020 및 우리가 제안한 Recurrent 블록이라는 세 가지 시간적 혼합 블록을 고려합니다.
글로벌 멀티 쿼리 주목
달리 명시하지 않는 한, 우리는 Transformer 기준선 Shazeer(2019)의 추론 속도를 향상시키기 위해 MHA 대신 MQA를 사용합니다. 우리는 고정 헤드 치수를 사용합니다. $D_{머리}=128$ , 주의 헤드 수를 고정합니다. $H$ 그렇게 $HD_{헤드}=D$ . 이를 위해서는 모델 차원이 필요합니다. $D$ 우리는 절대 위치 임베딩을 사용하지 않지만 상대 위치 임베딩으로 RoPE(Rotary Position Embedding) Su et al., 2021을 사용합니다.
로컬 슬라이딩 윈도우 주의
전역 주의를 사용할 때의 주요 단점 중 하나는 계산 복잡성이 시퀀스 길이에 따라 2차적으로 증가한다는 것입니다. 이 문제를 해결하기 위해 여러 작업에서 슬라이딩 윈도우 주목이라고도 알려진 지역적 관심 Beltagy et al., 2020을 채택하기 시작했습니다. 이를 통해 각 위치는 과거에 고정된 수의 토큰에만 참석할 수 있습니다. 이는 계산 FLOP를 줄일 뿐만 아니라 KV 캐시의 크기를 창 크기로 제한하여 더 이상 시퀀스 길이가 2차가 되지 않도록 합니다. 기타 모든 세부 사항은 글로벌 MQA와 동일합니다.
반복 차단
우리의 순환 블록(그림 2(c))은 GSS 블록 Mehta et al., 2022 및 Mamba Gu and Dao, 2023에서 사용하는 블록과 유사합니다. $D$ 출력 차원을 사용하여 두 개의 선형 레이어를 적용합니다. $D_{RNN}$ 동시에 두 개의 분기를 생성합니다. 첫 번째 분기에서는 H3 Dao et al., 2022b의 Shift-SSM에서 영감을 받아 시간 필터 차원이 4인 작은 분리 가능한 Conv1D 레이어를 적용합니다. 이 Conv1D 레이어는 매우 작습니다. $4D$ 매개변수. 우리는 제안된 RG-LRU 레이어(아래에 정의됨)를 사용하여 Conv1D 레이어를 따릅니다. 두 번째 분기에서는 GeLU 비선형성을 적용한 다음 요소별 곱셈을 통해 분기를 병합합니다. 그런 다음 출력 차원을 사용하여 최종 선형 레이어를 적용합니다. $D$ .
리얼 게이트 선형 순환 장치(RG-LRU)
제안된 RG-LRU 레이어는 LRU(Linear Recurrent Unit) Orvieto et al.(2023b)에서 영감을 받은 단순 반복을 갖지만 비선형 RNN, 특히 LSTM Hochreiter 및 Schmidhuber(1997)에 대한 문헌에서 영감을 받은 게이팅 메커니즘을 통합합니다. GRUs Chung et al., 2014. 레이어를 설명하는 방정식은 다음과 같습니다.
$$begin{align} r_t &= sigma(W_{a} x_t + b_a), & text{반복 게이트} \ i_t &= sigma(W_{x} x_t + b_x), & text{ 입력 게이트} \ a_t &= a^{cr_t}, & text{} \ h_t &= a_t odot h_{t-1} + sqrt{1 - a_t^2} odot(i_t odot x_t). & text{} end{align}$$
레이어의 출력은 다음과 같습니다. $y_t=h_t$ , 그리고 비선형성 $시그마$ 방정식에는 시그모이드 함수가 있습니다. 반복 가중치 $a$ 식 (4)에서 대각선이다. 따라서 모든 작업은 요소별로 수행됩니다. 우리는 매개변수화합니다 $a$ 식 (3)에서 다음과 같이 $a=시그마(람다)$ , 어디 $람다$ 학습 가능한 매개변수입니다. 이는 다음을 보장합니다. $0 <= a <= 1$ , 재발이 안정적인지 확인합니다. 변수 $c$ 는 8로 설정된 스칼라 값 상수입니다. 수치적 안정성을 위해 실제로는 다음을 계산합니다. $a^{cr_t}$ 로그 공간에서(부록 A 참조) 레이어에는 두 입력 모두에 게이트가 있습니다. $x$ 그리고 반복 체중 $a$ . 그러나 두 게이트 모두 반복 상태에 의존하지 않습니다. $h_{t-1}$ 이는 장치에서 계산이 효율적으로 실행될 수 있도록 보장합니다. 우리는 둘 다 초기화합니다 $W_{a}$ 그리고 $W_{b}$ LeCun init LeCun et al., 2002를 사용하여 초기화합니다. $람다$ 그렇게 $a^c$ 사이에 균일하게 분포되어 있습니다. $0.9$ 그리고 $0.999$ 훈련 시작 시 (Orvieto et al., 2023b.)와 유사합니다. SSM 문헌의 많은 최근 연구와 달리 RG-LRU는 직교 다항식 Gu et al.(2020) 이론에서 영감을 받은 초기화를 사용하지 않으며 기본 연속 시스템 Gu et al.의 이산화로 정의되지도 않습니다. 2021a. 원래 LRU 레이어와 달리 반복에서는 복잡한 대수를 사용하지 않습니다. 복잡한 반복을 사용하면 더 표현적인 레이어가 생성되지만 Orvieto et al., 2023a에서는 Gu 및 Dao(2023)에서도 관찰한 것처럼 복잡한 반복이 실제로 언어 모델링에 유익하지 않다는 것을 발견했습니다. (부록 B 참조)
게이트 동작
입력 게이트 $i_t$ 입력을 필터링(또는 축소)할 수 있는 LSTM의 것과 유사합니다. $x_t$ . 그러나 우리가 아는 한, 재발 게이트는 $r_t$ 문헌의 다른 게이팅 메커니즘과 다릅니다. 예를 들어, Mamba Gu and Dao, 2023에서 제안된 선택 메커니즘은 보간하는 GRU의 업데이트 게이트와 유사합니다. $x_t$ . 숨겨진 상태에 대한 효과는 LSTM의 망각 게이트와 유사하게 상태를 재설정하고 과거에 보유하고 있던 모든 정보를 잊어버릴 수 있게 해줍니다. 대조적으로, 우리의 재발 게이트는 Orvieto et al., 2023a의 표준 LRU 업데이트와 이전 숨겨진 상태 사이를 대략적으로 보간할 수 있으므로 입력을 효과적으로 삭제하고 이전 기록의 모든 정보를 보존할 수 있습니다(자세한 내용은 부록 A 참조). ). 우리는 이 게이트의 핵심 역할이 정보가 없는 입력의 영향을 줄여 모델이 초지수 기억을 달성할 수 있도록 하는 것이라고 믿습니다.
