Matrix Orthogonalization 도입으로 mLSTM Noisy AR 성능 최대 45.4%p 향상

Matrix Orthogonalization Improves Memory in Recurrent Models

2026년 7월 1일4분advanced

AI 요약

Context

Transformer의 Quadratic-attention 오버헤드로 인해 Long-horizon RL 등 특정 도메인에서 Recurrent Neural Network(RNN) 채택이 불가피한 상황. 기존 mLSTM은 Matrix Memory를 통해 Associative Recall(AR) 성능을 높였으나, 노이즈가 포함된 Noisy Associative Recall(NAR) 환경에서 성능 저하 문제 발생.

Technical Solution

Muon Optimizer의 Momentum Orthogonalization 개념을 mLSTM Memory Matrix Read 과정에 이식
Newton-Schulz iteration 5회를 적용하여 Memory Matrix의 Represented Directions을 균등하게 조정하는 구조 설계
특정 강한 방향성이 업데이트를 지배하여 약한 기억(Weak Memories)이 Crowd-out 되는 현상을 방지하는 Equalization 로직 구현
Frobenius norm(eps = 1e-6) 기반의 정규화를 통해 수치적 안정성 확보
Orthogonalized Memory를 Read-out에만 사용하고 Write-back을 배제하여 성능 저하 방지 및 Gradient flow 유지
AdamW Optimizer 기반의 학습 환경에서 Newton-Schulz 과정을 통한 추가 FLOPs 및 연산 시간 트레이드오프 수용

실천 포인트

1. Recurrent 모델의 Memory Matrix에서 특정 패턴의 지배 현상이 의심될 경우 Orthogonalization 검토

2. Newton-Schulz iteration 적용 시 Write-back 여부에 따른 성능 변동성 테스트 필수

3. 고정 파라미터 환경에서 추가 FLOPs를 통한 성능 이득의 비용 효율성 분석 필요

태그

#RNN #mLSTM #Matrix Orthogonalization #Newton-Schulz #Associative Recall

원문 읽기