Matrix Recurrent Units(MRU)에 관한 업데이트: Attention의 대안 [R]

최근 저는 제가 Attention의 대안으로 개발한 새로운 선형 시간 시퀀스 아키텍처(linear-time sequence architecture)인 Matrix Recurrent Units(MRU) 알고리즘을 다시 살펴보았습니다. 리포지토리(repo)에서 자세히 설명하고 있지만, 핵심은 MRU가 임베딩(embedding)을 입력 상태 행렬(input state matrix)로 변환하고, 시퀀스 차원(sequence dimension)을 따라 행렬들을 누적적으로 곱하여 출력 상태 행렬(output state matrix)을 얻은 다음, 다시 행렬을 벡터로 변환하는 방식으로 작동한다는 것입니다. MRU를 딥러닝(DL) 하드웨어에서 효율적으로 만들기 위해, 저는 해당 연산의 결합 법칙(associativity)을 활용하여 병렬 스캔(parallel scan)을 구현했습니다.

약 1년 전, 저는 Reddit에 제 프로젝트를 공유했으며(그 이후 계정 이름을 변경했습니다), toy dataset인 shakespeare-char에서 좋은 결과를 얻었습니다. 한 댓글 작성자는 행렬 상태(matrix states)를 제한하기 위해 취해진 단계들에 대해 질문했고, 또 다른 작성자는 더 포괄적인 데이터셋으로 학습할 때 학습이 본질적으로 불안정하다는 것을 발견했습니다. 저는 입력 상태 행렬을 생성하는 다양한 방법들을 실험함으로써 이 문제들을 해결했습니다. 원래 저는 단순히 입력 벡터를 행렬로 재구성(reshape)하고 항등 행렬(identity)을 더하는 방식을 사용했습니다. 그 이후로 저는 다음과 같은 방법들을 구현했습니다:

• 벡터의 요소들을 사용하여 왜대칭 행렬(skew-symmetric matrix)을 채우고, 행렬 지수(matrix exponential) 또는 케일리 맵(Cayley Map)을 사용하여 직교 행렬(orthogonal matrix)을 생성하는 방식
• LDU 인자(LDU factors)를 벡터의 요소들로 채우고, 행렬식(determinant)을 1로 강제하기 위해 D에 활성화 함수(activation function)를 사용하는 방식
• 행렬 지수 또는 케일리 맵을 사용하여 직교 행렬 Q를 만들고 상삼각 행렬(upper-triangular matrix) R을 채움으로써 QR 분해를 생성하는 방식
• 행렬식을 구하여 얻은 행렬식 교정 스칼라 인자(determinant-correcting scalar factor)로 나누는 방식

저는 이러한 수정 사항들이 다양한 트레이드오프(tradeoffs)를 가지며 손실 급증(loss spikes)을 방지한다는 것을 발견했습니다. 흥미롭게도, 스칼라 인자 방식은 더 나쁜 결과를 초래했습니다. 입력 상태를 나누는 것은 출력 상태를 스케일링(scaling)하는 방식으로만 영향을 미쳐야 하는데, 이는 스케일링되지 않은 모델이 더 복잡한 관계를 학습하는 대신 단순한 스칼라 감쇠(scalar decay) 패턴을 학습함으로써 toy dataset에서 "속임수(cheating)"를 쓰고 있었음을 나타냅니다.

또한, 입력 상태(input states)를 직교(orthogonal)하도록 강제하기 위해 Cayley Map 또는 행렬 지수(matrix exponential)를 사용하는 것은 놀랍게도 모델이 시퀀스에 대한 정보를 학습하는 것을 대부분 방해했으며, Cayley QR 방식보다 FFN(Feed-Forward Network)에 더 가까운 성능을 보였습니다. 직교 행렬(orthogonal matrices)의 저조한 성능은 전단 변환(shear transformations)을 학습하는 능력이 모델에게 매우 중요할 수 있음을 나타냅니다. 아마도 회전(rotations)은 이전 상태에 대한 의존성을 강제하는 반면, 전단(shearing)은 모델이 이전 상태와 더 독립적으로 상태를 조정할 수 있게 해주기 때문일 것입니다.
https://preview.redd.it/9ebh98q6uo8h1.png?width=2528&format=png&auto=webp&s=03ccef7f9b90762281aba31ab88af0368e273f69
https://preview.redd.it/fkkud7q6uo8h1.png?width=2528&format=png&auto=webp&s=5e9a2ef2b0e4319990950f16aa0648adebc2c360
위 그래프는 임베딩(embedding), 상태(state), 키(key), 값(value) 크기를 256으로 설정한 소규모 MRU LM, transformer, 그리고 FFN에 대한 shakespeare-char 데이터셋에서의 훈련 손실(train loss)과 검증 손실(validation loss)입니다. MRU LM은 단일 MRU 레이어와 4개의 MLP를 가지고, transformer는 단일 어텐션(attention) 레이어와 4개의 MLP를 가지며, FFN은 오직 4개의 MLP만을 가집니다. 저는 MRU의 효과만을 분리하기 위해 단일 시퀀스 믹싱(sequence-mixing) 레이어만을 사용했습니다.
마지막으로, 저는 더 큰 데이터셋으로 이동하여, 베이스라인 GPT-2 모델과 어텐션을 MRU로 대체한 모델을 훈련함으로써 https://huggingface.co/roneneldan/TinyStories-33M 을 재현하려고 시도했습니다. 결국 훈련 실행을 조기에 중단하게 되었지만, 손실 곡선(loss curves)은 이미 MRU가 이 작업에서 더 낮은 성능을 보인다는 것을 결정적으로 보여주는 것으로 보입니다. MRU의 입력 상태 행렬(input state matrices)을 생성하기 위해, 저는 가장 좋은 성능을 보이는 LDU 인수 분해(LDU factors) 방식을 사용했습니다.
https://preview.redd.it/p2uh1pyfuo8h1.png?width=2528&format=png&auto=webp&s=d6406574e0275f1aad52e89cca6462fd55116fcd
위 그래프는 huggingface 모델 카드와 동일한 하이퍼파라미터 및 차원을 사용하는 transformer와 MRU를 사용하는 LM의 검증 손실(validation loss)입니다. 공식 TinyStories 모델은 20 에포크(epochs) 동안 훈련되었으며, 이는 약 200k 스텝(steps)에 해당합니다.

다른 선형 시간 (linear-time) 모델들과 비교하기 위해, "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention"에서 설명된 알고리즘을 사용하여 선형 트랜스포머 (linear transformer)를 짧게 훈련해 보기도 했습니다. 제 연구 결과에 따르면, MRU가 생성적 언어 모델링 (generative language modeling)을 위한 어텐션 (attention)의 직접적인 대체재로 작동하기는 어려울 것으로 보이지만, 저는 이미 이 알고리즘을 위한 기초를 마련해 두었습니다. MRU는 어텐션 (attention), 상태 공간 모델 (state space models, SSM), 전통적인 RNN, 그리고 빠른 가중치 프로그래머 (fast weight programmers)와 같은 다른 알고리즘들과 비교했을 때 매우 다른 강점과 약점을 가지고 있습니다. MRU는 시퀀스(sequence)를 따라 훨씬 더 많은 누적 계산을 수행하며 (각 토큰에 대한 계산이 독립적인 것과는 대조적임), 훨씬 더 가볍기 때문에 더 빠르지만, 저장 용량 (storage capacity)은 훨씬 낮습니다. 저는 MRU의 대안적인 용도들이 여전히 탐구되어야 한다고 믿습니다.

MRU의 한 가지 활용 방안은 어텐션 (attention)의 쿼리 (query) 및 키 (key) 벡터에 적용하는 것입니다. RoPE와 유사하게 벡터의 청크 (chunks)를 회전시키겠지만, 2차원보다 더 큰 차원에서 동적이고 비가환적 (non-commutative)인 각도로 청크를 회전시킬 수 있을 것입니다. 이는 제가 계속해서 연구할 이 알고리즘의 많은 응용 분야 중 하나이며, 다른 분들도 그 응용 분야에 관심을 가지시길 바랍니다. 관심이 있으시다면 mikayahlevi@gmail.com, Reddit, GitHub, 또는 제가 활동하는 다른 플랫폼을 통해 연락해 주세요.
submitted by /u/mikayahlevi to r/MachineLearning
[link] [comments]