애플 연구진 'ParaRNN' 공개… 비선형 RNN 병렬 학습으로 665배 가속, 7B 클래식 RNN으로 트랜스포머 경쟁력 확보 (ICLR 2026 Oral)
애플 연구진이 비선형 순환신경망(RNN) 학습을 병렬화하는 'ParaRNN' 프레임워크를 공개했다. ICLR 2026에 오럴(Oral)로 채택된 논문 'ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models'에서 발표됐으며, 코드베이스는 오픈소스로 공개됐다.
핵심 성과는 시퀀셜 학습 대비 665배의 학습 속도 향상이다. 이 효율 향상으로 7B 파라미터 규모의 클래식 RNN을 처음으로 학습할 수 있게 됐고, 언어 모델링에서 트랜스포머와 경쟁할 만한 성능을 달성했다고 애플은 밝혔다.
RNN은 추론 시 토큰 생성 비용이 컨텍스트 길이와 무관해 효율적이지만, 학습 시 시퀀스 차원에 대한 병렬화가 어렵다는 한계가 있었다. 트랜스포머의 어텐션은 시퀀스 길이에 대해 제곱으로 연산이 증가하지만 병렬 처리가 가능한 반면, RNN은 단계별 펼침(unroll)을 거쳐야 한다.
기존 우회책은 Mamba 같은 선택적 상태공간 모델(SSM)처럼 은닉 상태에 대한 순환 관계를 선형으로 단순화하는 것이었다. 선형 연산은 결합법칙이 성립해 병렬 환원(parallel scan)으로 시퀀셜 단계를 로그 단계로 줄일 수 있지만, 모델링 가능한 은닉 상태 진화의 종류가 줄어 표현력이 제한된다.
ParaRNN은 비선형 방정식을 푸는 뉴턴법을 차용했다. RNN 전 시퀀스를 모든 단계의 은닉 상태가 미지수인 단일 방정식 시스템으로 재구성하고, 뉴턴법이 자코비안(국소 도함수)으로 비선형성을 선형 근사해 반복적으로 해를 구한다. 선형화된 RNN 시스템은 선형 SSM과 동일한 형태가 되어 자코비안이 상태 행렬 역할을 하며 병렬로 풀 수 있다.
실험에서는 GRU·LSTM 셀에 적용해 단 3회 반복으로 수렴하는 것을 확인했다. 즉, 잘 설계된 3회의 병렬 SSM 적용으로 비선형 RNN의 시퀀셜 적용과 동일한 은닉 상태 진화를 복원할 수 있어 학습 시 wall-clock 시간을 크게 단축한다.
대규모 학습을 위해 자코비안 구조도 단순화했다. ParaGRU는 대각(diagonal), ParaLSTM은 블록 대각(block-diagonal) 자코비안을 갖도록 셀 정의의 행렬을 주대각 성분만 살리는 방식으로 적응했다. 이로써 원래는 은닉 상태 크기에 대해 저장은 제곱, 곱셈은 세제곱으로 늘어나는 비용을 통제했다.
이 구조에서 효율적 병렬 환원을 수행할 맞춤 CUDA 커널을 구현했고, GPU 메모리 계층을 따라 데이터 지역성을 최대화했다. 완전 융합(fully-fused) 구현은 뉴턴 반복, 시스템 어셈블리, 병렬 환원을 하나의 커널에서 처리한다. 코드베이스에는 순수 PyTorch, CUDA 가속 환원이 추가된 PyTorch, 완전 융합 CUDA 세 가지 구현이 제공된다.
검증을 위해 400M에서 7B 파라미터 범위의 모델을 언어 모델링 과제로 학습했다. 7B 규모에서 클래식 RNN이 LLM으로서 경쟁력을 보였고, 퍼플렉서티와 다운스트림 태스크 성능에서 ParaGRU와 ParaLSTM이 경쟁력 있는 결과를 기록했다고 애플은 밝혔다.