행렬 직교화(Matrix Orthogonalization)가 순환 모델(Recurrent Models)의 기억력을 향상시킨다

2026-06-30

이 연구는 Paradigm의 지원을 받았습니다.

Transformer는 놀라운 연상 회상 (Associative Recall, AR) 능력을 보여줍니다. 어텐션 (Attention)은 각 토큰이 그 이전의 토큰들에 직접 접근할 수 있도록 하며, 이는 순환 신경망 (Recurrent Neural Networks, RNNs)과 같은 다른 아키텍처들이 따라잡기 어려웠던 메커니즘입니다.

하지만 일부 도메인에서는 Transformer의 이차 복잡도 어텐션 (Quadratic-attention) 오버헤드를 감당할 수 없습니다. 한 가지 예로 Dreamer 스타일의 장기 지평 강화학습 (Long-horizon RL)이 있습니다. 이러한 종류의 애플리케이션을 위해서는 순환 신경망 (RNNs)이 제대로 작동하게 만들어야 하지만, 연상 회상 (Associative Recall) 능력을 포기하고 싶지는 않습니다.

연상 회상 (Associative Recall)을 위해 가장 잘 알려진 RNN은 행렬 메모리 (Matrix memory)를 유지하는 LSTM의 변형인 mLSTM입니다. mLSTM은 하나의 벤치마크인 MQAR에서 베이스라인 대비 실질적으로 향상된 회상 능력을 보여줍니다. 하지만 순수한 회상 (Recall) 능력만으로는 순환 성능을 측정하기에 충분하지 않을 수 있습니다. 환경 전이 (Environment transitions)가 노이즈가 심할 수 있는 분야에서는 노이즈가 있는 연상 회상 (Noisy Associative Recall, NAR)이 유용한 대리 테스트 (Proxy test)가 됩니다.

MQAR는 NAR을 측정하지 않으므로, MAD의 노이즈가 있는 AR 작업 제품군을 살펴볼 수 있습니다. 다음은 작업이 어떻게 구성되는지에 대한 예시입니다:

0 9 3 10 12 13 15 14 0 9 5 8 2 9

여기서 키(Key) 0은 값(Value) 9에 매핑되고, 키 3은 값 10에 매핑되는 식입니다. MAD 생성기는 키(Keys), 값(Values), 그리고 방해 요소(Distractors)를 위해 서로 다른 토큰 범위를 사용합니다. 따라서 키가 0-5라면, 토큰 12-15는 방해 요소가 됩니다. NAR에 능숙한 모델은 시작 부분에서 0 -> 9를 보았을 때, 중간에 섞인 방해 요소 토큰들을 무시하면서 10번째 위치에서 9를 예측할 수 있어야 합니다.

그렇다면 어떻게 순환 NAR (Non-Associative Recurrence)을 개선할 수 있을까요? 우리는 언어 모델링 (Language Modelling)에서 매우 성공적이었던 옵티마이저 (Optimizer)인 Muon의 아이디어를 빌려올 수 있습니다. Muon은 모멘텀 (Momenta)을 직교화 (Orthogonalize)하여, 표현된 방향들의 평형 장치 (Equalizer) 역할을 합니다. 이는 몇몇 강력한 방향이 업데이트를 지배하는 것을 방지하고, 더 약한 방향들을 끌어올립니다. 특히 최근 연구에 따르면 Muon이 꼬리 부분의 연상 기억 (Tail-end Associative Memory) 학습에서 Adam보다 뛰어난 성능을 보인다는 점이 매우 관련이 깊습니다. 핵심 아이디어는 이러한 평형화가 약한 기억들이 밀려나는 것을 방지한다는 것입니다.

이에 영감을 받아, 우리는 읽기 (Read) 과정에서 mLSTM 메모리 행렬을 직교화하고 이 추가적인 프로세스로 학습하는 것이 NAR 성능을 향상시키는지 테스트하기로 했습니다.

우리는 MAD 노이즈 AR 샘플을 사용한 다음 토큰 예측 (Next-token prediction) 작업에서 mLSTM 베이스라인을 직교화된 변형 모델과 비교합니다. 학습 및 평가를 위해 우리는 다양한 어휘 크기 (Vocab sizes)와 시퀀스 길이 (Sequence lengths)에 대해 frac_noise를 0.8로 설정한 MAD 노이즈-회상 (MAD noisy-recall)을 사용합니다. 모든 모델은 배치 크기 (Batch size) 64에서 2k 스텝 동안 AdamW (betas = 0.9, 0.999, weight_decay = 0.01)를 사용하여 학습되었습니다. 학습률 (Learning rate)은 각 작업 설정에 대해 3e-4, 1e-3, 3e-3, 1e-2를 스윕 (Sweeping)하여 선택되었습니다.

우리는 매 스텝마다 학습을 위한 새로운 배치를 생성하며, 실험마다 별도의 고정된 검증 세트 (Validation set)를 유지합니다. 직교화를 위해, 우리는 프로베니우스 노름 (Frobenius norm, eps = 1e-6)으로 정규화하고 5회의 뉴턴-슐츠 (Newton-Schulz) 반복을 적용합니다. 우리는 이 프로세스를 통해 그래디언트 (Gradients)가 흐르도록 허용합니다. 결정적으로, 우리는 직교화된 메모리를 다시 쓰지 않는데, 이는 성능을 저하시킨다는 것을 발견했기 때문입니다. 우리는 오직 읽기 (Readouts) 용도로만 이를 사용합니다. 우리 실험에 대한 완전하게 재현 가능한 코드는 여기에서 찾을 수 있습니다.

우리는 직교화 (orthogonalization)가 전반적으로 성공률과 평균 정확도를 향상시킨다는 것을 발견했습니다. 흥미로운 점은 vocab-96 영역으로 진입함에 따라 그 격차가 더 벌어지는 것처럼 보인다는 것이며, 이는 직교화가 일반적인 mLSTMs가 어려움을 겪는 까다로운 NAR (Non-Autoregressive) 작업에서 가장 큰 도움을 준다는 것을 시사합니다. 후자의 두 사례 (vocab 96, seq len 768/1024)에서 직교화는 mLSTMs를 실패 직전의 상태 (24개 시드 중 4개 해결)에서 상당히 더 신뢰할 수 있는 성능 (14-16개 시드 해결)으로 끌어올립니다. 이는 우리가 의도했던 작은 개입(intervention)치고는 매우 놀라운 결과입니다. Newton-Schulz 방식은 추가적인 FLOPs와 실제 실행 시간 (wall-clock time)을 대가로 지불하는 대신, 고정된 파라미터 수 내에서 추가적인 이득을 얻게 해줍니다.

우리는 이러한 결과에 대해 과도하게 해석하지 않도록 주의해야 합니다. 이 결과들은 작은 모델 영역에서 유효하며, NAR은 합성 작업 (synthetic task)입니다. NAR에서의 이득이 더 큰 모델을 위한 실제 벤치마크 전반의 이득으로 이어지는지 조사해 볼 가치가 있을 것입니다.

이 포스트를 작성하는 동안 피드백과 제안을 준 Dan Robinson, Alpin Yukseloglu, 그리고 Glen Taggart에게 감사드립니다.