Orvieto, Antonio, et al. "Resurrecting recurrent neural networks for long sequences." International Conference on Machine Learning. PMLR, 2023.

Resurrecting Recurrent Neural Networks for Long Sequences

Abstract

Recurrent Neural Network (RNN)는 긴 시퀀스에 대해 빠른 추론 속도를 제공하지만, 최적화하기 어렵고 학습 속도가 느리다. 최근 Deep State-Space Model (SSM)은 긴 시퀀스 모델링 task에서 놀라운 성능을 보여주었으며, 빠른 병렬 학습과 RNN과 유사한 빠른 추론이라는 추가적인 이점을 가지고 있다. 그러나 SSM은 겉보기에는 RNN과 유사하지만, RNN보다 성능이 향상되는 이유가 불분명한 중요한 차이점들이 존재한다.

본 논문에서는 표준 신호 전파(signal propagation) 논증을 사용하여 Deep RNN을 신중하게 설계하면, 긴 범위 추론 task에서 Deep SSM의 인상적인 성능을 회복할 수 있으며, 학습 속도 또한 일치시킬 수 있음을 보여준다. 이를 위해 우리는 표준 RNN에 대한 일련의 변경 사항들을 분석하고 ablation 연구를 수행한다. 이러한 변경 사항에는 선형화(linearizing) 및 재귀(recurrence)의 대각화(diagonalizing), 더 나은 파라미터화(parameterizations) 및 초기화(initializations) 사용, 그리고 forward pass의 적절한 정규화(normalization) 보장 등이 포함된다.

우리의 결과는 Deep SSM의 인상적인 성능의 기원에 대한 새로운 통찰력을 제공하며, 동시에 Long Range Arena 벤치마크에서 Deep SSM의 성능과 계산 효율성을 모두 일치시키는 Linear Recurrent Unit이라는 RNN 블록을 소개한다.

1. Introduction

Recurrent Neural Network (RNN)는 딥러닝 초기부터 핵심적인 역할을 해왔으며, 시퀀스 데이터를 모델링할 때 자연스러운 선택이다 (Elman, 1990; Hopfield, 1982; McCulloch and Pitts, 1943; Rumelhart et al., 1985). 그러나 이러한 네트워크는 Turing completeness와 같은 강력한 이론적 특성(Chung and Siegelmann, 2021; Kilian and Siegelmann, 1996)을 가지고 있음에도 불구하고, 실제 학습에서는 어려움이 많다는 것이 잘 알려져 있다. 특히 RNN은 vanishing 및 exploding gradient 문제 (Bengio et al., 1994; Hochreiter, 1991; Pascanu et al., 2013)를 겪으며, 이로 인해 모델이 데이터 내의 **장거리 의존성(long-range dependencies)**을 학습하기 어렵게 만든다. 이러한 문제를 완화하기 위해 직교/유니터리 RNN (Arjovsky et al., 2016; Helfrich et al., 2018), 그리고 long short-term memory (LSTM) (Hochreiter and Schmidhuber, 1997) 및 gated recurrent units (GRUs) (Cho et al., 2014a)와 같은 gating mechanism을 포함한 여러 기술이 개발되었다. 그럼에도 불구하고, 이러한 모델들은 본질적으로 순차적인 계산 특성 때문에 최적화 속도가 느리며 (Kalchbrenner et al., 2016), 따라서 확장하기 어렵다.

최근 몇 년 동안 Transformer (Vaswani et al., 2017)는 시퀀스 모델링 task에서 점점 더 큰 주목을 받으며, 광범위한 응용 분야에서 놀라운 성공을 거두었다 (Brown et al., 2020; Dosovitskiy et al., 2020; Jumper et al., 2021). RNN과 비교하여 attention layer는 학습 중에 확장 및 병렬화하기 더 쉽고, 결정적으로 vanishing gradient 문제를 겪지 않는다. 이는 시퀀스 내의 어떤 두 토큰 간의 상호작용도 네트워크 내의 직접적인 엣지로 모델링되기 때문이다. 그러나 attention layer의 주요 문제는 계산 및 메모리 비용이 시퀀스 길이 $L$에 대해 $O(L^2)$로 제곱에 비례하여 증가한다는 점이다. 따라서 Transformer는 긴 시퀀스에 배포할 때 특히 비용이 많이 들 수 있다. 시퀀스 길이에 선형적으로 비례하는 RNN은 중간 정도의 시퀀스 길이에서도 추론 시 Transformer보다 일반적으로 빠르다 (Liu et al., 2019).

이러한 문제들에 동기를 받아, Gu et al. (2021a)은 최근 S4 모델을 도입했다. S4는 신중하게 설계된 **deep state-space model (SSM)**로, 매우 장거리 추론을 명시적으로 요구하도록 설계된 벤치마크인 Long Range Arena (LRA) (Tay et al., 2020)의 task에서 놀라운 성능을 달성했다. S4는 이론적으로 원칙적이며, 현대 제어 시스템의 잘 확립된 구성 요소인 연속 시간 선형 SSM에서 영감을 받았다. 더 중요한 것은, S4 layer와 그 변형들 (DSS, S4D, S5 등) (Gu et al., 2022a; Gupta et al., 2022a; Smith et al., 2022)이 적절한 이산화(discretization) 기술을 사용하여 hidden state (RNN처럼)를 통해 토큰 간의 상호작용을 모델링함으로써 attention layer의 $O(L^2)$ 병목 현상을 극복한다는 점이다. 이러한 모델은 layer를 RNN처럼 단순히 unroll함으로써 추론 시 매우 효율적으로 만들 수 있다. 또한, SSM은 시간 차원에서 선형이므로, 일반적인 RNN 학습의 느린 순차적 특성과 달리 학습 중에 쉽게 병렬화할 수 있다. 이는 긴 시퀀스에서 매우 계산 효율적이다.

Figure 1 | (왼쪽) 본 논문에서 소개하는 Deep Linear Recurrent Unit (LRU) 아키텍처는 S4 (Gu et al., 2021a)에서 영감을 받았다. 이 모델은 LRU 블록들의 스택으로 구성되며, 그 사이에 비선형 projection이 있고, skip connection과 batch/layer normalization과 같은 정규화 방법도 사용한다. 자세한 내용은 §D에서 설명하고, pseudocode는 §A에 제공한다. 또한 본 연구의 모든 recurrent 모듈 변형 (tanh dense, linear dense 등)에 대해 동일한 아키텍처 구조 (Norm-Recurrence-GLU-Skip)를 사용한다. (오른쪽) 서론에서 설명한 tanh RNN에서 LRU를 설계하기 위한 주요 단계들의 효과 요약. 각 단계에서 recurrent 모듈의 Long Range Arena (LRA) 평균 성능 (3개 seed)을 deep SSM의 평균 성능과 비교하여 보여준다. 모든 LRA task에서 LRU는 S4/S4D/S5와 같은 deep SSM의 성능과 일치한다. 자세한 결과는 §3에 있다.

S4 모델은 추론 시 RNN과 동등하지만, 학습 시에는 여러 독특한 특성을 가지고 있다. 예를 들어, S4는 잠재 연속 시간 미분 방정식 시스템의 이산화(discretization)로 매개변수화된다. 또한 S4는 다항식 projection 이론 (Gu et al., 2020)에서 영감을 받은 state matrix의 특정 초기화를 사용한다. 이러한 특성들이 모델의 인상적인 성능을 설명하는 것처럼 보일 수 있지만, 후속 연구들 (Gu et al., 2022a; Gupta et al., 2022a,b; Smith et al., 2022)은 S4가 사용하는 특정 초기화가 성능에 항상 결정적인 것은 아니며, 최고 성능을 달성하는 이산화 규칙이 이론과 다를 수 있다 (Smith et al., 2022)고 제안했다. 따라서 deep SSM의 이러한 독특한 특성들이 기계적으로 어떤 역할을 하는지, 그리고 어떻게 단순화될 수 있는지 불분명하다.

RNN과 deep SSM 간의 놀라운 유사성에 동기를 받아, 그리고 이러한 모델의 성능을 이끄는 근본적인 메커니즘을 더 잘 이해하기 위해, 우리는 장거리 추론을 위한 deep 아키텍처의 핵심 구성 요소로 사용될 때 RNN의 능력과 한계를 연구한다. 우리의 주요 목표는 다음 질문에 답하는 것이다:

"Deep RNN을 사용하여 deep 연속 시간 SSM의 성능과 효율성을 일치시킬 수 있는가?"

우리는 이 질문에 긍정적인 답변을 제시한다. 우리는 S4와 같은 deep SSM이 제공하는 성능 향상이 vanilla deep RNN에 일련의 작은 변화를 줌으로써 달성될 수 있음을 보여준다. 이러한 변화를 통해 우리는 Long Range Arena (LRA) 벤치마크 (Tay et al., 2020)에서 이러한 deep SSM의 성능과 효율성을 회복할 수 있다. 우리는 이 새로운 RNN 모델을 **Linear Recurrent Unit (줄여서 LRU)**이라고 부른다.

주요 단계 (Main Steps). 우리는 여기에서 성능이 우수하고 효율적인 RNN 모델을 만들기 위한 주요 단계들을 설명한다. 이러한 관찰 중 일부는 이전 연구들에서 이루어졌지만 (§B 참조), 우리는 새로운 통찰력을 이끌어내는 새로운 관점과 신중한 ablation을 제공한다. 본 논문에서 제시된 각 단계는 recurrent network의 특정 속성을 밝히고, deep RNN을 학습하고 초기화하는 데 있어 도전 과제와 모범 사례를 보여준다.

• 선형 Recurrence (Linear Recurrences). deep 아키텍처에서 SSM layer를 tanh 또는 ReLU 활성화 함수를 사용하는 vanilla RNN layer로 대체하면 Long Range Arena (LRA)에서의 성능이 크게 떨어진다. 놀랍게도, §3.1에서 우리는 RNN의 recurrence에서 비선형성을 단순히 제거하는 것 (즉, 선형 recurrence를 사용하는 것)이 테스트 정확도를 크게 향상시킨다는 것을 발견한다. 우리는 §E.1에서 선형 RNN layer와 비선형 MLP 블록을 쌓는 것 (Fig.1)이 recurrence에 비선형성이 없어도 복잡한 sequence-to-sequence 맵을 모델링할 수 있음을 보여줌으로써 이러한 효과를 설명한다. 비선형성을 제거하는 것이 표현력을 해치지 않는 것처럼 보이지만, 이는 gradient가 얼마나 빨리 vanishing 또는 exploding될 수 있는지를 직접 제어할 수 있는 능력부터 학습을 병렬화할 수 있는 능력까지 여러 장점을 가져온다. 우리의 발견은 recurrence 또한 선형인 deep SSM의 성공을 부분적으로 설명하기도 한다.
• 복소 대각 Recurrent Matrix (Complex Diagonal Recurrent Matrices). dense 선형 RNN layer는 네트워크의 표현력이나 초기화 시 feature에 영향을 주지 않으면서 복소 대각 형태로 재매개변수화될 수 있다 (§3.2). 대각선형 RNN layer는 또한 parallel scan을 사용하여 recurrence의 고도로 병렬화 가능한 unrolling을 허용하여 학습 속도를 크게 향상시킨다 (Martin and Cundy, 2017). 우리는 이러한 관찰이 이전 SSM (Gupta et al., 2022a; Smith et al., 2022)에서 활용되었으며, 선형 RNN layer에도 중요한 효율성 향상을 제공한다는 것을 검증한다.
• 안정적인 지수 매개변수화 (Stable Exponential Parameterization). §3.3에서 우리는 대각 recurrent matrix에 지수 매개변수화를 사용하는 것이 중요한 이점을 가진다는 것을 보여준다. 결정적으로, 이는 학습 중 안정성을 쉽게 강제할 수 있게 하며, 이는 다시 초기화 분포를 수정하여 장거리 추론을 용이하게 하고 성능을 향상시킨다. 우리의 결과는 여러 최신 SSM이 사용하는 특정 결정론적 초기화보다는, 초기화 시 recurrent layer의 고유값 분포가 모델이 장거리 추론을 포착할 수 있는지 여부를 결정한다는 것을 시사한다.
• 정규화 (Normalization). §3.4에서 우리는 매우 장거리 의존성을 가진 task를 학습할 때 forward pass에서 hidden activation을 정규화하는 것이 중요하다는 것을 보여준다. 이 최종 수정을 통해 우리의 RNN은 LRA 벤치마크의 모든 task에서 deep SSM의 성능과 일치할 수 있다. state-space model과 다시 연결하여, 우리는 §4에서 우리의 정규화가 S4의 이산화 구조와 어떻게 연결될 수 있는지 보여준다.

본 논문에서 사용된 deep Linear Recurrent Unit (LRU) 아키텍처와 위 각 단계가 성능에 미치는 영향을 Fig.1에 요약하였다. 우리는 본 연구의 주요 목적이 S4 기반 모델의 성능을 능가하는 것이 아니라, 적절하게 초기화되고 매개변수화될 때 단순한 RNN도 장거리 추론 task에서 강력한 성능을 달성할 수 있음을 입증하는 것임을 강조한다. 우리는 본 논문에서 도출된 통찰력이 미래 아키텍처를 설계하고 기존 아키텍처를 단순화하는 데 유용할 수 있다고 믿는다.

2. Preliminaries

이 섹션에서는 본 연구에서 다룬 핵심 아키텍처 구성 요소들(RNN 및 SSM)을 비교하고, 우리의 방법론 및 실험 설정을 설명한다. 관련 아키텍처에 대한 더 자세한 논의는 관련 연구 섹션 §B를 참조할 수 있다.

2.1. Recap of recurrent block structures

본 논문에서 다루는 주요 아키텍처 구성 요소들을 개괄적으로 설명하며, Vanilla RNN과 최근의 S4-like deep SSMs (Gu et al., 2021a, 2022a; Gupta et al., 2022a; Smith et al., 2022) 간의 주요 차이점에 초점을 맞춘다.

RNN Layer.
$H_{\text {in }}$-차원 입력 시퀀스 $(u_{1}, u_{2}, \ldots, u_{L})$는 중간 레이어 계산의 결과(순차적 구조를 유지하는) 또는 초기 입력으로 간주될 수 있다. $N$-차원 hidden state를 가진 RNN layer는 학습 가능한 파라미터 $A \in \mathbb{R}^{N \times N}, B \in \mathbb{R}^{N \times H_{\text {in }}}, C \in \mathbb{R}^{H_{\text {out }} \times N}, D \in \mathbb{R}^{H_{\text {out }} \times H_{\text {in }}}$를 사용하여 recurrent 계산을 통해 $H_{\text {out }}$-차원 출력 시퀀스 $(y_{1}, y_{2}, \ldots, y_{L})$를 계산한다.

$$x_{k}=\sigma\left(A x_{k-1}+B u_{k}\right), \quad y_{k}=C x_{k}+D u_{k}$$

여기서 $x_{0}=0 \in \mathbb{R}^{N}$에서 시작한다. $\sigma$는 비선형 함수를 나타내며, 종종 tanh 또는 sigmoid 활성화 함수로 선택된다. 만약 $\sigma$가 항등 함수라면, RNN layer는 **선형(linear)**이라고 한다.

S4-like recurrent layer.
Gu et al. (2021a)에서 소개된 S4 recurrence의 간소화된 버전을 제시한다. 입력 $(u_{0}, u_{1}, \ldots, u_{L-1})$은 이제 latent continuous-time signal $u_{\mathrm{ct}}: \mathbb{R}_{\geq 0} \rightarrow \mathbb{R}^{H_{\mathrm{in}}}$을 stepsize $\Delta>0$의 배수로 샘플링한 결과로 간주된다. 즉, 모든 $k \in 0, \ldots, L-1$에 대해 $u_{\mathrm{ct}}(\Delta k):=u_{k}$이다. 출력 시퀀스 $(y_{0}, y_{1}, \ldots, y_{L-1})$는 다시 stepsize $\Delta$로, $x_{\mathrm{ct}}(0)=0$으로 초기화된 다음 continuous-time state-space model에 의해 계산된 signal $y_{\mathrm{ct}}: \mathbb{R}_{\geq 0} \rightarrow \mathbb{R}^{H_{\text {out }}}$에서 샘플링된다.

$$\begin{aligned} & \frac{d}{d t} x_{\mathrm{ct}}(t)=\tilde{A} x_{\mathrm{ct}}(t)+\tilde{B} u_{\mathrm{ct}}(t) \\ & y_{\mathrm{ct}}(t)=\Re\left[\tilde{C} x_{\mathrm{ct}}(t)\right]+\tilde{D} u_{\mathrm{ct}}(t) \end{aligned}$$

여기서 $\Re(p)$는 복소수 벡터 $p$의 실수부를 나타내며, $\tilde{A}=\operatorname{diag}(\tilde{a})$ (여기서 $\tilde{a} \in \mathbb{C}^{N}$), $\tilde{B} \in \mathbb{C}^{N \times H_{\text {in }}}, \tilde{C} \in \mathbb{C}^{H_{\text {out }} \times N}, \tilde{D} \in \mathbb{R}^{H_{\text {out }} \times H_{\text {in }}}$이다. 이 모델의 continuous-time 특성을 무시하면, Eq.(1)과 비교했을 때 가장 두드러진 차이점은 (a) 우변의 계산이 hidden state와 입력에 대해 **선형(linear)**이라는 점, 그리고 (b) 대부분의 파라미터가 복소수 값을 가지며, $\tilde{A}$가 대각 행렬이라는 점이다. $\tilde{B}, \tilde{C}, \tilde{D}$는 복소수 무작위 또는 균일 초기화를 따르는 반면, 전이 행렬 $\tilde{A}$는 구조화되어 있으며, 즉 HiPPO 이론 (Gu et al., 2020)을 통해 대각 형태로 결정론적으로 초기화된다. 일반적인 선택 (Gu et al., 2022a)은 $n=1,2, \ldots, N$에 대해 $\tilde{a}_{n}=-\frac{1}{2}+i \pi n$ (S4D-Lin) 및 $\tilde{a}_{n}=-\frac{1}{2}+i \frac{N}{\pi}\left(\frac{N}{n+1}-1\right)$ (S4D-Inv)이다.
학습 및 추론을 위해 Eq.(2)의 continuous-time 시스템은 high-accuracy Zero-Order-Hold (ZOH) 또는 Bilinear method를 통해 stepsize $\Delta$로 이산화된다. ZOH 방법은 다음을 제공한다.

$$x_{k}=A x_{k-1}+B u_{k}, \quad y_{k}=C x_{k}+D u_{k},$$

여기서 $x_{-1}=0, A=\exp (\Delta \tilde{A}), B=(A-I) \tilde{A}^{-1} \tilde{B}, C=\tilde{C}, D=\tilde{D}$이며, $\exp$는 행렬 지수(matrix exponential)를 나타낸다. $u_{\mathrm{ct}}$가 타임스탬프 사이에서 상수라는 가정 하에 (이는 모델링 가정으로 간주될 수 있음), 이 수치 적분은 정확하다 (Jacquot, 2019). 또한, $\tilde{A}$가 대각 행렬이므로 이러한 모든 이산화 연산은 요소별로 빠르게 수행될 수 있다.

몇 가지 주요 차이점.
RNN과 SSM 간의 몇 가지 중요한 차이점을 강조하기 위해 몇 가지 구조적 및 계산적 특성을 지적할 가치가 있다:

• Eq.(3)은 선형이므로, 계산이 순차적으로 수행되어야 하는 비선형 RNN과 달리, 병렬 스캔(parallel scans) (Martin and Cundy, 2017; Smith et al., 2022)을 사용하여 $k=L-1$까지 효율적으로 병렬화될 수 있다.
• Eq.(3)은 선형 RNN 계산과 유사하지만, (a) $A$와 $B$가 이산화에 의해 규정된 특이한 방식으로 파라미터화된다는 점, 그리고 (b) 이 행렬들이 파라미터를 공유한다는 점을 주목하는 것이 중요하다. 특히 $\Delta$는 $A$와 $B$ 모두에 영향을 미친다. 이러한 차이점은 SSM에서 학습이 continuous-time 파라미터 $\tilde{A}, \tilde{B}, \tilde{C}, \tilde{D}, \Delta$에 대해 수행되므로, 파라미터화 선택이 최적화에 직접적인 영향을 미친다는 점에서 중요하다.
• Vanilla RNN과 달리, 대부분의 SSM은 HiPPO 이론을 사용하여 결정론적으로 초기화되는 복소수 값의 대각 recurrent 행렬을 사용하며, 문헌에서는 SSM의 성공의 상당 부분을 사용된 특정 초기화에 기인한다 (Gu et al., 2021a, 2022b; Gupta et al., 2022a).

위의 요점들은 우리의 연구를 동기 부여한다: 본 논문에서는 Gu et al. (2021a, 2022a); Smith et al. (2022)와 동일한 아키텍처를 고려하지만, recurrent core의 SSM layer를 RNN으로 대체한다. 그런 다음 LRA (Tay et al., 2020) task에서 S4-like 성능을 점진적으로 회복하기 위해 어떤 단계를 밟아야 하는지 연구한다. 각 단계의 효과는 경험적 증거와 이론적 고려 사항에 의해 뒷받침되며, Fig.1에 제시된 아키텍처로 이어진다.

2.2. Experimental setup

본 논문에서는 Long Range Arena (LRA) 벤치마크 (Tay et al., 2020)를 고려한다. 이 벤치마크는 모델의 장거리 시퀀스 모델링 능력을 테스트하기 위해 설계된 task들로 구성되어 있다 (단, sequential CIFAR-10 분류 task에서는 흑백 이미지 대신 컬러 이미지를 사용한다). Transformer는 이 task들 대부분에서 좋은 성능을 보이지 못하는 반면, 깊은 SSM(State Space Model)은 이 task들에서 뛰어난 성능을 보여주었다 (Dao et al., 2022a; Gu et al., 2021a). 따라서 LRA 벤치마크는 깊은 RNN의 장거리 모델링 능력을 탐구하기에 적합하다.

모든 실험에서 우리는 Gu et al. (2021a) (Fig.1)와 유사하게 residual connection과 layer/batch normalization (Ba et al., 2016; Ioffe and Szegedy, 2015)을 갖춘 6개 layer 네트워크를 사용한다. 우리는 SSM layer를 RNN layer로 대체하고, 일련의 단계를 거쳐 LRU recurrence를 구축한다 (자세한 내용은 §3 참조). 모든 실험은 3회 반복되며, 평균과 표준 오차를 보고한다. 네트워크는 AdamW optimizer (Loshchilov and Hutter, 2017)를 사용하여 학습된다. Gu et al. (2021a); Steil (2004)의 제안에 따라, recurrent 파라미터에는 더 작은 learning rate와 weight decay를 적용하지 않는다. 모든 모델의 learning rate와 같은 하이퍼파라미터는 최고 정확도를 위해 logarithmic grid에서 튜닝된다. 실험 설정에 대한 더 자세한 내용은 §D를 참조하라.

3. Designing Performant Deep RNNs

이 섹션에서는 RNN이 LRA 벤치마크에서 deep SSMs의 인상적인 성능에 도달하기 위해 필요한 근본적인 설계 단계에 대해 논의한다. 서론에서 이미 개략적으로 설명한 이 단계들을 논리적인 순서로 제시하고, 각 주장은 실험적 증거와 이론적 고찰로 뒷받침하며, 이는 §E에서 더 자세히 다룬다.

우리는 Figure 1의 아키텍처를 고려한다. 여기서 순환 계산(recurrent computation)은 vanilla RNN부터 시작하여 점진적으로 수정된다. 먼저 §3.1에서 선형 순환(linear recurrences) 사용의 이점을 보여주고, 이어서 §3.2에서는 표현력(expressivity)과 초기화 분포에 영향을 주지 않으면서 학습 및 추론 속도를 높이는 방법을 제시한다. §3.3에서는 **파라미터화(parameterization) 및 초기화 분포를 변경하는 것이 RNN을 안정화하고 장거리 모델링(long-range modeling)을 개선하는 방법(및 이유)**에 대해 논의한다. 마지막으로 §3.4에서는 hidden activation에 대한 정규화 전략을 제안하여 LRU 아키텍처를 완성하며, 이는 deep SSMs와 성능 면에서 거의 일치하는 결과를 가져온다.

3.1. Linear RNN layers are performant

본 연구의 주요 발견 중 하나는 선형 RNN layer가 비선형 MLP 또는 GLU (Dauphin et al., 2017) 블록과 결합될 때 놀랍도록 표현력이 뛰어나며, 동일한 아키텍처 내에서 튜닝된 비선형 RNN 변형보다 우수한 성능을 보인다는 점이다. Tb.1에서 우리는 단순히 비선형성을 제거하여 다음 상태를 $x_{k}=A x_{k-1}+B u_{k}$로 계산하는 것이 대부분의 LRA task에서 test accuracy를 향상시킬 수 있음을 보여준다. vanilla 선형 RNN 블록이 제공하는 성능 향상이 일부 task (sCIFAR, PathFinder 및 PathX)에서 S4에 비해 여전히 뒤처지지만, 이 첫 번째 발견은 본 논문의 나머지 부분에서 recurrence 내의 비선형성을 제거하는 동기가 된다. 이후 섹션에서는 recurrence의 선형성을 활용하여 학습 속도를 크게 높이고, 장거리 의존성을 학습하기 위한 원칙적인 초기화 및 정규화 원리를 도출한다. Text 및 Retrieval task에서는 vanilla RNN을 사용한 성능이 이미 deep SSM의 성능과 일치함을 주목한다 (이러한 task에 대한 S4D/S5의 성능은 Tb. 3 참조).

Table 1 | recurrent unit에서 비선형성을 제거한 효과가 test accuracy에 미치는 영향 (§3.1). LRA의 sCIFAR, ListOps, Text 및 Retrieval task에 대한 결과만 보여준다. 이 모델들은 PathFinder/PathX에서 무작위 추측을 넘어서지 못했다 (Tb. 2 및 3에서 추가 개선 사항). deep SSM의 성능은 Tb.3에 나와 있다.

Tb.1의 경험적 결과는 놀랍다. recurrent 비선형성은 이론적으로나 실제적으로나 RNN 성공의 핵심 구성 요소로 여겨지기 때문이다 (Erichson et al., 2021; Pascanu et al., 2013; Siegelmann, 2012). 실제로 단일 layer sigmoidal 및 tanh RNN의 강력한 속성 중 하나는 Turing completeness이며, 이는 선형 변형으로는 달성할 수 없다 (Chung and Siegelmann, 2021). 그러나 우리가 사용하는 아키텍처(Fig.1)는 표준 RNN보다 깊고, 각 RNN 블록 뒤에 position-wise로 배치된 비선형성을 포함한다. §E.1에서는 recurrent 비선형성이 deep 모델의 표현력과 학습 가능성에 미치는 영향을 조사한다. spectral analysis와 Koopman operator theory (Koopman and Neumann, 1932)를 활용하여, 선형 RNN layer와 비선형 feedforward 블록을 교차 배치하는 것이 고도로 비선형적인 시스템을 근사하는 데 충분하다는 점을 논의한다. 우리 분석의 핵심 관찰은 position-wise 비선형성이 신호 정보를 고주파로 효과적으로 전달하여, 시스템이 spectral domain에서 선형성을 넘어설 수 있게 하고 layer capacity를 증가시킨다는 것이다. 선형 recurrences의 장점에 대한 우리의 주장을 더욱 강화하기 위해, §E.2에서는 선형 및 비선형 RNN이 중요한 종류의 근사 함수(선형 연산자, Wang et al. (2022) 참조)를 공유하지만, 비선형 활성화가 학습 속도를 잠재적으로 늦출 수 있음을 보여준다.

3.2. Using complex diagonal recurrent matrices is efficient

우리는 이제 복소수 대각선 recurrent matrix를 사용함으로써 deep linear RNN의 학습 및 추론 속도를 성능 저하 없이 크게 향상시킬 수 있음을 보인다. 계산 효율성을 위해 선형 시스템을 대각화하는 아이디어는 Gupta et al. (2022a)의 DSS 도입 이후 모든 deep SSM의 지배적인 특징이었지만, 이 섹션에서는 Tb.1의 Glorot-initialized deep linear RNN의 초기화 스펙트럼(§3.2.1 참조)과 정확히 일치하도록 대각화된 버전을 구성한다. 이 접근 방식의 주된 목적은 초기화와 대각화가 성능에 미치는 영향을 분리하는 것이다 (Tb.2 및 Tb.3 참조). 우리는 §3.2.1에서 몇 가지 유용한 선형 대수 요소를 상기하는 것으로 시작하여, §3.2.2에서 초기화 시 고유값 스펙트럼을 보존하면서 recurrence를 대각화하는 방법에 대해 논의한다.

3.2.1. Linear RNN eigendecomposition

재귀식 $x_{k}=A x_{k-1}+B u_{k}$는 $x_{-1}=0 \in \mathbb{R}^{N}$이라는 가정을 사용하여 쉽게 전개할 수 있다:

$$x_{0}=B u_{0}, \quad x_{1}=A B u_{0}+B u_{1}, \quad x_{2}=A^{2} B u_{0}+A B u_{1}+B u_{2}, \quad \ldots \quad \Longrightarrow \quad x_{k}=\sum_{j=0}^{k-1} A^{j} B u_{k-j} .$$

위 방정식에서 행렬 $A$의 지수 연산은 RNN에서 잘 알려진 vanishing/exploding gradient 문제의 원인이다 (Bengio et al., 1994; Pascanu et al., 2013). 비선형 RNN에서는 상태 $x_{k}$가 활성화 함수의 compact image 위에 존재하도록 강제되지만, 우리의 선형 변형 모델의 hidden-state는 $k$가 증가함에 따라 잠재적으로 기하급수적으로 폭발하거나 소멸할 수 있다. 이 현상은 고유값(eigenvalue) (a.k.a. spectral) 분석을 통해 더 잘 이해될 수 있다: 임의로 작은 항목 섭동(perturbation)까지, 모든 행렬 $A \in \mathbb{R}^{N \times N}$은 **대각화 가능(diagonalizable)**하다 (Axler, 1997). 즉, $A=P \Lambda P^{-1}$로 쓸 수 있으며, 여기서 $P \in \mathbb{C}^{N \times N}$는 가역 행렬이고 $\Lambda=\operatorname{diag}\left(\lambda_{1}, \lambda_{2}, \ldots, \lambda_{N}\right) \in \mathbb{C}^{N \times N}$이다. 고유값과 고유벡터가 실수인 대칭(symmetric) 설정과 달리, 비대칭(non-symmetric) 설정에서는 완전한 등가성을 달성하기 위해 복소수 항목을 허용해야 한다는 점이 중요하다. 분해식 $A=P \Lambda P^{-1}$를 Eq.(4)에 대입하고 양변에 $P^{-1}$를 곱하면 $\bar{x}_{k}=\sum_{j=0}^{k-1} \Lambda^{j} \bar{B} u_{k-j}$를 얻는다. 여기서 $\bar{x}_{k}:=P^{-1} x_{k}, \bar{B}:=P^{-1} B$이다. 출력은 $y_{k}=\Re\left[\bar{C} \bar{x}_{k}\right]+D u_{k} \in \mathbb{R}^{H}$로 계산될 수 있으며, 여기서 $\bar{C}=C P^{-1}$이고, 우리는 $\bar{C} \bar{x}_{k}$의 실수 부분을 취한다. 따라서 $(A, B, C, D)$를 학습하는 대신, 동등하게 $(\Lambda, \bar{B}, \bar{C}, D)$를 학습할 수 있으며, 여기서 $\Lambda, \bar{B}, \bar{C}$는 복소수 값을 가지고 $\Lambda$는 대각 행렬이다.

복소수가 정말 필요한가? 우리는 복소수를 사용하는데, 이는 비대칭 행렬을 대각 형태로 편리하고 간결하게 표현할 수 있기 때문이다. 그러나 이것이 유일한 옵션은 아니다. 실수만으로도 (거의) 효율적으로 작업할 수 있다. 이에 대한 방법은 §E.3에서 논의한다.

안정성(Stability). $\bar{x}_{k}=\sum_{j=0}^{k-1} \Lambda^{j} \bar{B} u_{k-j}$이므로, 타임스탬프 $k$에서 $\bar{x}$의 $j$번째 구성 요소의 norm은 $\left|x_{k, j}\right|= O\left(\left|\bar{x}_{k, j}\right|\right)=O\left(\left|\lambda_{j}\right|^{k}\right)$와 같이 진화한다. 따라서 안정성(즉, $x_{k}$가 폭발하지 않음)을 보장하기 위한 충분 조건은 모든 $j$에 대해 $\left|\lambda_{j}\right|<1$이다 (Gu et al., 2021a).

3.2.2. Learning in the diagonalized space

diagonal form으로 recurrent linear system을 학습하는 것은 학습 및 추론 모두에서 상당한 계산 속도 향상을 제공한다. 예를 들어, sCIFAR 구현에서 diagonal linear RNN은 ReLU가 있는 dense RNN보다 학습 속도가 약 8배 빠르며, S4D 및 S5 구현 속도와 일치함을 확인했다. 이러한 계산상의 이점의 주요 원인은 다음과 같다: (a) diagonal matrix의 거듭제곱은 사소한 반면(학습 및 추론 모두 속도 향상), dense matrix의 지수화는 계산 비용이 많이 든다. (b) nonlinear recurrence는 순차적으로 계산되어야 하는 반면, linear recurrence를 펼치는 것은 associative scan을 사용하여 병렬화할 수 있어 학습 속도가 빨라진다 (Gupta et al., 2022a; Smith et al., 2022).

Figure $2 \mid$ Glorot 초기화를 따르는 $A \in \mathbb{R}^{N \times N}$의 고유값: $A$의 각 항목은 평균 0, 분산 $1/N$의 가우시안 분포에서 독립적으로 샘플링된다. 고유값은 복소수이며 ($A$는 대칭이 아님) 복소 평면에 표시된다. 검은색 원은 단위 원판 $\{|z|=1\} \subseteq \mathbb{C}$이다. 극한 동작(균일 초기화)은 Thm. 3.1에 의해 예측된다.

Figure 3 | Lemma 3.2를 사용하여 샘플링된 항목을 가진 diagonal matrix $A$의 고유값. $r_{\text {min }}=0$, $r_{\text {max }}=1$일 때, 분포는 극한에서 Glorot 초기화와 일치한다.

동등한 초기화 (Equivalent initialization). diagonal linear system의 이점을 초기화의 역할과 분리하기 위해, 우리는 §3.1의 dense linear RNN과 우리의 diagonal system을 비교할 때 recurrence의 고유값 스펙트럼을 변경하지 않는 diagonal system에 대한 초기화를 찾는다. 여기서 $A$는 Glorot 초기화를 따랐다. 다행히도, 우리는 random matrix theory의 고전적인 결과(Ginibre, 1965)를 사용할 수 있다.

Theorem 3.1 (Strong circular law). $A_{N}$의 empirical spectral measure를 $\mu_{N}$이라고 하자. 여기서 $A_{N}$은 i.i.d. Gaussian entries를 가진 실수 $N \times N$ 행렬이며, 각 항목은 평균 0, 분산 $1/N$을 가진다. 그러면 $\mu_{N}$은 $N \rightarrow \infty$일 때 $\{|z| \leq 1\} \subseteq \mathbb{C}$ 상의 균일 확률 측도로 weakly almost surely 수렴한다.

위 정리는 Fig.2에 설명되어 있으며, Glorot 초기화에서 $A$의 스펙트럼이 사실상 $\mathbb{C}$의 단위 원판에서 샘플링됨을 보여준다. 이 결과는 Glorot 초기화가 대략적으로 안정적인 초기화를 제공한다는 것을 의미하므로 (§3.2.1의 정의 참조), §3.1에서 linear RNN의 강력한 성능을 뒷받침한다. 더욱이, Theorem 3.1로부터 diagonal system에 대한 동등한 스펙트럼 초기화가 도출되며, 이는 큰 너비(large width) 극한에서 정확하게 성립한다: $\Lambda$는 단위 원판에 균일하게 샘플링된 항목을 가진 diagonal matrix여야 한다. 복소수의 지수 정의 $\exp (-\nu+i \theta):=e^{-\nu}(\cos (\theta)+i \sin (\theta))$를 사용하여, 우리는 $\mathbb{C}$에서 반지름 $r_{\text {min }}$과 $r_{\text {max }}$ 사이의 링에 균일하게 샘플링하는 간단한 방식을 채택한다.

Lemma 3.2. $u_{1}, u_{2}$를 구간 $[0,1]$에서 독립적인 균일 확률 변수라고 하자. $0 \leq r_{\min } \leq r_{\max } \leq 1$이라고 하자. $\nu=-\frac{1}{2} \log \left(u_{1}\left(r_{\max }^{2}-r_{\min }^{2}\right)+r_{\min }^{2}\right)$ 및 $\theta=2 \pi u_{2}$를 계산한다. 그러면 $\exp (-\nu+i \theta)$는 $\mathbb{C}$에서 반지름 $r_{\text {min }}$과 $r_{\text {max }}$ 사이의 링에 균일하게 분포한다.

우리는 $r_{\text {min }}=0$ 및 $r_{\text {max }}=1$로 설정하여 Glorot 초기화의 스펙트럼을 복구한다 (이러한 hyper-parameter 튜닝은 §3.3에서 탐색할 것이다). Tb. 2 (첫 두 행)는 complex diagonal form으로 deep linear RNN을 학습한 결과를 보여준다. 여기서 $\Lambda$의 각 diagonal entry는 Lemma 3.2를 사용하여 $\left[r_{\min }, r_{\max }\right]=[0,1]$로 $\mathbb{C}$의 단위 원판에 균일하게 초기화된다. 우리 실험에서 $\bar{B}, \bar{C}$ (편의상 $B$와 $C$로 다시 명명)는 실수부와 허수부 모두 Glorot 초기화를 따르며 (각각 별도로 parameterize됨), 입력-출력 투영에서 길이를 보존하기 위해 각 구성 요소의 분산은 절반으로 줄어든다 (Glorot and Bengio, 2010). 마지막으로, SSM 계산 후 신호의 실수부는 유지되고 허수부는 버려진다 (Gu et al. (2022a); Gupta et al. (2022a)에서와 같이).

Tb. 2의 결과는 recurrence를 diagonalize하는 것이 ListOps 및 sCIFAR와 같은 task에서 정확도를 놀랍게도 향상시킨다는 것을 보여준다. 더 중요한 것은, 모든 LRA task에서 학습 및 추론 시간을 크게 줄이고 (학습 속도 비교는 §C. 1의 Tb. 4 참조), RNN을 S4D 및 S5와 같은 deep SSM만큼 빠르게 학습시킨다.

Table 2 | transition matrix의 다른 parametrization 하에서 linear diagonal complex RNN의 테스트 정확도 (§3.2 참조). 성능은 Tb.1의 결과를 직접적으로 향상시키며, $\Lambda$의 지수(극좌표) 표현의 이점을 보여준다. linear RNN block에 대한 최상의 parametrization option은 굵은 글씨로 표시되어 있다. Ring Init은 $r_{\text {min }}$과 $r_{\text {max }}$가 튜닝된 변경된 초기화를 나타낸다. Text 및 Retrieval task의 성능은 linear RNN이 이미 S4 결과와 일치하므로 표시되지 않는다 (Tb. 1과 Tb.3 비교). 이 모델들은 아직 PathX를 해결할 수 없으며, hidden activation을 정규화하고 $\Lambda$의 고유값을 작은 위상으로 초기화해야 한다 (Tb.3 참조).

3.3. Benefits of stable exponential parameterization

§3.2에서 우리는 복잡한 대각 재귀(diagonal recurrences)가 계산적으로 효율적임을 보였다. 그러나 일부 실험에서는 대각 모델 학습이 dense 모델 학습보다 불안정할 수 있음도 관찰했다. 장거리 의존성(long-range dependencies)을 학습하고 gradient vanishing을 피하기 위해, 재귀(recurrence)의 고유값(eigenvalues)은 크기가 1에 가까워야 한다 (Gu et al., 2022b; Gupta et al., 2022a). 하지만 이러한 고유값은 학습 중 시스템을 불안정하게 만들 가능성도 있다. 이 섹션에서는 RNN의 안정적인 파라미터화(parameterization)와 $r_{\text {min }}$ 및 $r_{\text {max }}$ 튜닝(Lemma 3.2 참조)의 이점을 보여준다.

지수 파라미터화(exponential parameterization) 하의 최적화. Lemma 3.2는 대각화된 RNN의 자연스러운 파라미터화를 $\Lambda=\operatorname{diag}(\exp (-\nu+i \theta))$로 제안한다. 여기서 $\nu \in \mathbb{R}^{N}$와 $\theta \in \mathbb{R}^{N}$는 학습 가능한 파라미터이다( $\Lambda$의 실수부와 허수부 대신). §E.2에서 시각화하기 쉬운 2차원 예시(Fig.8 참조)를 통해 설명하듯이, 이 선택은 크기(magnitude)와 진동 주파수(oscillation frequencies)를 분리하여 Adam을 사용한 최적화를 더 쉽게 만든다. ZOH 이산화(discretization)의 일부 특징과 유사하며(§2 및 §4 참조), 특히 PathFinder의 성능을 무작위 수준 이상으로 끌어올리는 이 지수 파라미터화의 긍정적인 효과는 Tb.2의 세 번째 행에서 관찰할 수 있다.

안정성 강제(Enforcing stability). 지수 파라미터화의 중요한 이점은 고유값에 안정성을 쉽게 강제할 수 있다는 점이다. 이를 이해하기 위해, 초기화 시 $\left|\lambda_{j}\right|=\left|\exp \left(-\nu_{j}\right)\right| \leq 1$임을 주목하라. 이는 $\nu_{j}>0$이기 때문이다. 따라서 학습 중 안정성을 유지하기 위해, 우리는 지수 함수 또는 다른 양의 비선형성을 사용할 수 있다: $\lambda_{j}:=\exp \left(-\exp \left(\nu_{j}^{\log }\right)+i \theta_{j}\right)$. 여기서 $\nu^{\log } \in \mathbb{R}^{N}$는 우리가 최적화하는 파라미터이며, 초기화 시 $\nu_{j}^{\log }:=\log (\nu)$로 설정한다. 유사한 아이디어가 deep SSMs (Gu et al., 2021a)에서 이산화(discretization) 맥락에서 사용된다는 점에 주목하라. 우리는 단순한 ReLU 비선형성 대신 지수 비선형성을 선택하여 $|\lambda|=1$ 주변에서 세분화(granularity)를 증가시킨다. 이는 $\nu^{\log }=-\infty$에서 달성된다(반면 $|\lambda|=0$은 $\nu^{\log }=\infty$에서 달성된다). 안정적인 파라미터화는 대부분의 LRA task에서 도움이 된다. Tb.2의 네 번째 행에서 우리는 sCIFAR, ListOps 및 Pathfinder에 대한 그 효과를 보여준다. 우리는 LRA에서 가장 어려운 장거리 의존성 task 중 하나인 Pathfinder에서 가장 극적인 개선을 관찰했으며, 여기서 성능은 이제 93% 이상에 도달한다.

안정적인 파라미터화의 이점은 우리가 $\Lambda$의 고유값을 단위 원판(unit disk)에 더 가까운 링(ring)에 초기화하는 아이디어를 탐구할 때 더욱 분명해진다(Lemma 3.2에서 $r_{\text {min }}$을 1에 가깝게 증가시켜 네트워크가 더 긴 범위의 상호작용에 편향되도록 하고 gradient vanishing을 피하는 것). 실제로 Gu et al. (2022b); Gupta et al. (2022a)에서 자세히 논의된 바와 같이, 멀리 떨어진 토큰 간의 상호작용을 고려해야 하는 추론의 경우, 재귀(recurrence)의 고유값은 크기가 1에 가까워야 한다. 그렇지 않으면, Eq.(4)의 대각 버전에서 명확하듯이, 원점(origin)에 가까운 고유값의 거듭제곱을 취할 때 과거 토큰의 신호는 빠르게 사라진다 (§3.2.1 참조). §C의 Tb.5의 마지막 행에서 보여주듯이, 안정성을 강제하지 않으면 sCIFAR task에서 $r_{\text {max }}$를 0.9 이상으로 증가시킬 때 성능이 저하되기 시작한다. 안정성을 강제하면 $r_{\text {max }}$를 0.99까지 증가시키고 성능을 향상시킬 수 있다. 우리는 $r_{\text {min }}$과 $r_{\text {max }}$의 다른 값들을 스윕하는 다른 task에서도 유사한 이점을 본다(Tbs. 7 & 8에 더 자세한 내용이 있다). 마지막으로, 여기서 우리는 $\Lambda$의 고유값 크기를 변경하는 것을 탐구했지만, §3.4에서는 고유값을 작은 위상(phase)으로 초기화하여 더 전역적인 패턴을 학습하는 이점도 보여준다. 이는 특히 장거리 추론 task에 유용하다.

Table 3 | 안정적인 지수 파라미터화와 링(ring) 초기화를 갖춘 대각 RNN에 $\gamma$ 정규화(normalization)를 추가한 후의 성능 (§3.4 참조). PathX의 경우, 초기화 시 **더 작은 고유값 위상(eigenvalue phase)**을 추가로 사용한다. 우리는 이 아키텍처를 LRU라고 명명한다. 우리는 초기화 분포와 학습률 설정을 위해 $r_{\text {min }}$과 $r_{\text {max }}$를 스윕한다. 또한 S4/S4D/S5의 결과(우리의 RNN 모델과 유사한 하이퍼파라미터 스윕을 사용한 자체 파이프라인에서의 재현 결과 포함)도 보고한다. LRU는 모든 LRA task에서 이러한 deep SSMs와 유사한 성능에 도달한다.

3.4. Additional considerations for long-range reasoning tasks

지금까지 우리의 모델은 벤치마크에서 가장 어려운 데이터셋인 PathX (시퀀스 길이 $16k$ 토큰)를 학습하는 데 성공하지 못했다. 이 섹션에서는 매우 긴 범위의 종속성(long-range dependencies)을 학습하는 모델의 능력을 향상시키기 위해 필요한 추가적인 수정 사항을 논의하고, LRU 모델을 완성한다.

정규화 (Normalization)
§3.3에서 우리는 장거리 task에서 더 나은 성능을 위해 $\Lambda$의 고유값(eigenvalues)을 단위 원(unit disk)에 가깝게 초기화했다. 그러나 $r_{\text{min}}$과 $r_{\text{max}}$를 1에 가깝게 이동시킬수록, 초기화 시 학습 손실(training loss)이 발산하기 시작하는 것을 관찰했다 (Fig.5 참조). 이 섹션에서는 먼저 이 현상을 설명하는 결과를 제시한 후, 이 문제를 해결하고 성능을 더욱 향상시키기 위한 은닉 활성화(hidden activations)에 대한 실용적인 정규화 기법을 도출한다.

명제 3.3 (순방향 전달 발산 - Forward-pass blow-up)
$\Lambda$가 대각 행렬이고, 고유값들이 $\mathbb{C}$에서 반지름 $r_{\min} < r_{\max} < 1$인 두 원 사이의 링(ring)에 균일하게 샘플링되었다고 가정하자. 이 경우, 상수 또는 백색 잡음(white-noise) 입력과 Glorot 입력 투영(input projection) 하에서, 상태 $x_k$의 제곱 노름(squared norm)은 $k \rightarrow \infty$일 때 다음 값으로 수렴한다.

$$\mathbb{E}\left[\left\|x_{\infty}\right\|_{2}^{2}\right]=\frac{1}{r_{\max }^{2}-r_{\min }^{2}} \log \left(\frac{1-r_{\min }^{2}}{1-r_{\max }^{2}}\right) \mathbb{E}\left[\|B u\|_{2}^{2}\right] .$$

이 결과는 $r_{\text{min}} = r_{\text{max}} = r$일 때 다음과 같은 직관적인 형태를 가진다: 만약 우리가 단위 원에 $\rho$-가깝게 초기화한다면, 순방향 전달(forward pass)은 $1/\rho$ 인자만큼 발산한다 (이전 상태의 기여가 더 오래 지속되기 때문): $\epsilon = r_{\text{max}}^2 - r_{\text{min}}^2$이고 $\rho = 1 - r_{\text{max}}^2$라고 하면:

$$\lim _{\epsilon \rightarrow 0} \frac{\mathbb{E}\left[\left\|x_{\infty}\right\|_{2}^{2}\right]}{\mathbb{E}\left[\|B u\|_{2}^{2}\right]}=\lim _{\epsilon \rightarrow 0}\left[\frac{1}{\epsilon} \log \left(1+\frac{\epsilon}{\rho}\right)\right]=\lim _{\epsilon \rightarrow 0}\left[\frac{1}{\epsilon}\left(\frac{\epsilon}{\rho}+O\left(\epsilon^{2}\right)\right)\right]=\frac{1}{\rho}=\frac{1}{1-r^{2}} .$$

순방향 전달을 위한 효과적인 정규화 기법을 도출하기 위해, 백색 잡음 입력 하에서 1차원 설정에 대한 $1/\rho$ 이득 공식의 간소화된 유도를 제시한다. ${ }^{8}$ $\Lambda = \lambda \in \mathbb{C}$이고 $B=1$이라고 하자. $p \in \mathbb{C}$의 켤레 복소수를 $p^*$라고 하면, $|p|^2 = p^* p$이고, 입력에 대한 기댓값에서 Eq.(4)와 $i \neq j$일 때 $\mathbb{E}[u_{k-i} u_{k-j}] = 0$이라는 사실을 사용하여 다음을 얻는다:

$$\mathbb{E}\left|x_{k}\right|^{2}=\left(\sum_{i=0}^{k-1} \lambda^{i} \mathbb{E}\left[u_{k-i}\right]\right)\left(\sum_{j=0}^{k-1} \lambda^{j} \mathbb{E}\left[u_{k-j}\right]\right)^{*}=\sum_{i, j=0}^{k-1} \lambda^{i}\left(\lambda^{j}\right)^{*} \mathbb{E}\left[u_{k-i} u_{k-j}\right]=\sum_{i=0}^{k-1}|\lambda|^{2 i} \xrightarrow{\infty} \frac{1}{1-|\lambda|^{2}} .$$

위 공식은 우리의 재귀(recurrence)에서 모든 유클리드 방향에 대해 성립하므로 ($\Lambda$는 대각 행렬), 요소별로 초기화되는 정규화 파라미터를 추가할 수 있다. 또한, $\lambda$가 1에 가까워질수록 $1-|\lambda|^2$는

Figure 4 | 임펄스 입력 $u=(1,0,0, \ldots, 0) \in \mathbb{R}^{16 k}$ 하에서의 $x \in \mathbb{R}^{3}$의 변화. $x$의 3가지 구성 요소가 다른 색상으로 표시되어 있다. $\Lambda$는 파라미터 $\nu_j = 0.00005$를 가지며, $\theta_j$는 $[0, 2\pi]$ 또는 작은 위상 $[0, \pi/50]$에서 균일하게 샘플링된다. $L=1024$와 같은 짧은 시퀀스 (PathFinder, sCIFAR)의 경우, $[0, 2\pi]$는 허용 가능한 전체 진동 횟수를 가진 커널을 생성한다: $u_0$에 대한 정보는 전체 상태 기록에서 몇 번만 호출된다. 반면, $L$이 높은 경우, 유사한 효과를 얻기 위해서는 초기화 시 허수부의 범위가 더 작아야 한다.

Figure 5 | PathX task에서 정규화 및 초기화 시 작은 위상 사용의 효과. 각 설정에 대해 100k 반복 동안 세 번의 독립적인 실행에 대한 평균 및 표준 오차를 보여준다. 정규화가 없으면, 모델은 초기화 시 더 높은 손실 값을 보이고 빠르게 준최적 값(suboptimal value)으로 수렴하며, 훈련 및 테스트 정확도는 모두 무작위 확률 수준이다. 정규화를 추가하면 도움이 된다: 훈련 손실은 초기화 시 더 낮고, 옵티마이저는 준최적 영역을 벗어나 훈련 정확도도 증가한다. 흥미롭게도, 이 모델은 여전히 전혀 일반화되지 못한다. 마지막으로, 초기화 위상을 줄이면 (즉, $\theta$의 범위를 조정하면) 훈련 세트에서의 수렴이 극적으로 향상되며, 테스트 세트에도 일반화된다.

0에 가까워지므로, 이 파라미터의 SGD를 통한 추가적인 적응이 어려워진다. 따라서 우리는 **정규화 파라미터 $\gamma^{\log} \in \mathbb{R}^N$**를 사용하며, 이는 요소별로 $\gamma_i^{\log} \leftarrow \log(\sqrt{1-|\lambda_i|^2})$로 초기화된다. ${ }^{9}$ 그리고 재귀식을 다음과 같이 수정한다:

$$x_{k}=\Lambda x_{k-1}+\exp \left(\gamma^{\log }\right) \odot\left(B u_{k}\right)$$

여기서 $\odot$는 요소별 곱셈을 나타낸다. $\gamma$ 파라미터는 RNN이 해당 고유 방향으로 공급되는 입력을 적응적으로 스케일링할 수 있도록 한다. 우리는 $\gamma$ 정규화가 sCIFAR 및 Pathfinder와 같이 단위 원에 가깝게 초기화하는 것이 유리한 task에서 성능을 일관되게 향상시킨다는 것을 Tb.3에서 보여주었다.

초기화 시 고유값 위상 감소 (Reducing Eigenvalue Phase at Initialization)
대각화된 재귀식의 맥락에서, 우리는 $\Lambda = \operatorname{diag}(\exp(-\exp(\nu^{\log}) + \theta))$를 가진다. 여기서 $\nu^{\log} \in \mathbb{R}^N$는 로그 고유값 크기(log eigenvalue magnitudes)의 벡터이고, $\theta \in \mathbb{R}^N$는 고유값 위상(eigenvalue phases)의 벡터이다. $\nu^{\log}$가 원점으로부터의 거리를 인코딩하는 반면, $\theta$는 벡터 $1+0i$로부터의 각도이다. 긴 시퀀스의 경우, $\theta \sim [0, 2\pi]$로 균일하게 초기화하면 대부분의 상태 엔트리가 초기화 시 전반적으로 많은 수의 진동을 보일 것이다 (Fig.4의 상단 패널 참조). 동등하게, 이 설정에서 대부분의 상태 차원은 국부적인 진동 패턴의 평균을 포착하는 컨볼루션 ${ }^{10}$의 결과이다. 이러한 동작은 장거리 종속성 포착 능력 ($\nu^{\log}$에 의해 제어됨)과는 독립적이지만, RNN이 저장하는 정보의 본질과 관련이 있다. 따라서 우리는 긴 시퀀스 데이터에서 균일한 위상으로 $\Lambda$를 초기화하는 것이 네트워크를 입력 시퀀스에서 가짜 특징(spurious features)을 학습하도록 본질적으로 편향시킨다고 주장한다. 모델은 이러한 준최적 초기화에서 회복할 수 없다: 실제로 우리는 PathX에서 지금까지 가장 좋은 모델의 경우, 몇 번의 반복 후 훈련 손실이 매우 준최적의 최소값으로 수렴하여 무작위 확률 테스트 성능으로 이어진다는 것을 관찰했다 (Fig.5 참조). 이 문제를 해결하기 위해, 우리는 $\theta$의 범위를 0 주변의 얇은 조각으로 제한하여 모델이 더 전역적인 특징(global features)을 학습하도록 편향시키는 것으로 충분하다는 것을 발견했다. 최적의 $\theta$ 값이 작기 때문에, 최적화를 돕기 위해 위상을 로그로 매개변수화한다: $\theta = \exp(\theta^{\log})$, 여기서 $\theta^{\log}$가 최적화된다.

초기화 시 위상 범위를 $[0, \pi/10]$로 제한함으로써, 우리의 LRU는 PathX에서 94.2%의 성능을 달성했으며, 이는 state-of-the-art deep SSMs와 일치한다. 우리는 다른 LRA task에 대해 초기화 시 더 작은 위상을 사용하는 것을 탐색하지 않았지만, 이것이 다른 task에서도 성능을 더욱 향상시킬 수 있다고 믿는다. $\gamma$ 정규화와 초기화 시 고유값 위상 제한 모두 PathX를 해결하는 데 결정적이었다는 점에 유의해야 한다. $\gamma$ 정규화를 도입하지 않고 초기화 시 위상만 제한했을 때는 학습이 불가능했다.

§3의 모든 구성 요소를 종합하여, 우리는 이 새로운 모델을 **Linear Recurrent Unit (줄여서 LRU)**이라고 명명한다. 이는 깊은 RNN을 효율적으로 초기화하고 학습하기 위한 유연하고 해석 가능하며 원칙적인 프레임워크를 제공하며, Tb.3에서 보여주듯이 모든 LRA task에서 deep SSMs의 성능 및 효율성과 일치한다.

4. Insights on S4 and Variants

우리는 §3의 ablation 연구가 deep SSM의 성공을 이끄는 근본적인 메커니즘을 설명한다고 믿는다. 따라서 본 논문을 마무리하며, 이 섹션에서는 LRU 모델과 diagonal SSM 간의 주요 유사점과 차이점을 자세히 살펴보고 몇 가지 통찰력을 제시한다. §2에서와 마찬가지로, 기술적인 복잡성을 피하기 위해 원래 S4 논문에서 파생된 모델들의 주요 특징을 포착하는 단순화된 논의를 제공한다. 다른 모델들과의 비교는 §B를 참조하라.

§2에서 자세히 설명했듯이, **diagonal SSM (DSS, S4D, S5)**은 잠재 연속 시간 모델 $\dot{x}_{\mathrm{ct}}(t)=\tilde{A} x_{\mathrm{ct}}(t)+\tilde{B} u_{\mathrm{ct}}(t)$의 이산화(discretization)를 통해 인스턴스화되고 파라미터화된다. 여기서 $A=\operatorname{diag}(\tilde{a})$는 복소수 엔트리로 초기화되며, 종종 HiPPO 이론(Gu et al., 2020)에 의해 규정되거나 영감을 받는다. 스텝 사이즈 $\Delta$를 갖는 Zero-Order-Hold (ZOH) 이산화는 **점화식 $x_{k}=\exp (\Delta \tilde{A}) x_{k-1}+(\exp (\Delta \tilde{A})-I) \tilde{A}^{-1} \tilde{B} u_{k}$**로 이어진다. 이 공식은 우리의 Eq.(7)에 비해 복잡하다고 볼 수 있지만, 다음 단락에서 설명하듯이 우리의 공식과 관련이 있다.

행렬 지수(Matrix exponentials)는 학습을 더 쉽게 만든다. ZOH 공식의 지수(exponential)는 $\dot{x}_{\mathrm{ct}}(t)=\tilde{A} x_{\mathrm{ct}}(t)$의 정확한 적분에서 비롯되며, 이는 $x_{\mathrm{ct}}(\Delta k)=\exp (\Delta \tilde{A}) x_{\mathrm{ct}}(\Delta(k-1))$로 이어진다. 또한, 안정성을 강화하기 위해 S4에서 영감을 받은 모델에서는 $A$의 실수 부분이 종종 양의 비선형성(positive nonlinearity)에 입력되는데, 이는 §3.3에서도 동일하게 적용된다. §3.3의 결과와 최적화 이점에 대한 논의(§E.2 참조)를 통해, 우리는 지수 파라미터화의 힘이 반드시 정확한 적분(우리 시스템에는 없음)에 기인하는 것이 아니라, 점화식에서의 크기-위상 분리(magnitude-phase decoupling)에 더 근본적으로 뿌리를 두고 있으며(이는 Adam으로 학습을 더 쉽게 만든다, Fig.8 참조), 대각화된 공간에서 학습하는 전반적인 이점(Tb.2 참조)에도 기인한다고 주장한다. 또한, 비선형성을 추가하여 점화식을 안정화하는 것이 우리 실험에서도 유익했지만, 이는 S4의 기본 이론에서 규정하는 바는 아니다.

구조화된 초기화는 필요하지 않다. Gu et al. (2022a); Gupta et al. (2022b); Smith et al. (2022)도 HiPPO 구조에서 벗어난 $A$에 대한 초기화를 논의하지만 (§2 및 § 참조), 우리가 아는 한, 단위 원판(unit disk)의 한 조각에 대한 간단한 균일 초기화가 적절한 정규화와 결합될 때 LRA에서 가장 어려운 task인 PathX를 해결할 수 있음을 보여준 것은 우리가 처음이다. 우리는 또한 (Tb.2) 원판에 대한 균일 초기화(이는 Glorot 초기화의 대각화된 버전이다, Thm. 3.1)가 나머지 LRA task에서 더 복잡한 deep state-space model에 가까운 성능을 달성하기에 충분함을 보여준다. 우리의 결과는 궁극적으로 HiPPO 이론이 이 분야의 발전에 근본적이지만, S4 성공의 주요 원천으로 간주되어서는 안 된다는 것을 시사한다.

이산화는 초기화 스펙트럼을 변경한다. 단순화를 위해 S4D-Lin에 초점을 맞추자. S4D-Lin의 경우 $A=\operatorname{diag}(\tilde{a})$이며 $\tilde{a}_{n}=-\frac{1}{2}+i \pi n$이다. 이는 $\exp (-\Delta / 2+i \pi \Delta n)$로 초기화된 요소(즉, 고유값)를 갖는 대각 전이 행렬을 생성한다. 일반적인 선택, 예를 들어 $\Delta=1e-3, N=128$에서 SSM 고유값은 **크기 $\exp (-\Delta / 2) \approx 0.9995$와 위상 $\theta=\pi \Delta n \tilde{\epsilon}[0, \pi / 8]$**를 갖는다. 즉, 초기화는 $\mathbb{C}$의 단위 원에 가까운 링(ring)에서 수행되며, 고유값 크기와 연결된 제한된 위상을 갖는다. §3.3 및 §3.4의 결과에서 분명하듯이, 고유값 위상과 크기를 연결하는 것이 좋은 성능을 달성하는 데 필수적이지 않다. 실제로 Tb.3에서 볼 수 있듯이, Long Range Arena (PathX 제외)의 테스트 정확도는 완전한 링에서 더 자연스러운 크기 독립적 초기화를 사용하여 복구될 수 있다. §3.4에서 논의했듯이, 초기화 위상을 0 주변의 작은 범위로 변경하는 것은 첫 번째 원리(first principles)에 의해 동기 부여될 수 있지만, 극도로 긴 시퀀스에만 필요하다. 이 수정은 S4에 이미 하드코딩되어 있으며, 작은 $\Delta$를 선택하면 위상도 축소된다. 그러나 우리의 결과는 $\Delta$ 파라미터를 통해 학습 중에 실수부와 허수부를 연결하는 것이 PathX에서도 좋은 성능을 달성하는 데 필요하지 않음을 분명히 보여준다.

이산화는 정규화를 수행한다. 우리 모델과 ZOH 이산화된 S4 점화식 간의 가장 눈에 띄는 시각적 차이는 **$u_{k}$에 대한 행렬 승수: $(\exp (\Delta \tilde{A})-I) \tilde{A}^{-1} \tilde{B}$**에 있다. S4D에 대한 실험을 수행한 후, 우리는 이 승수를 $\Delta$에 대한 1차 전개(first-order expansion), 즉 $\Delta \tilde{B}$로 단순히 대체하는 것이 성능에서 거의 일치함을 발견했다. 입력 차원 $H=1$ 및 단위 $B \in \mathbb{R}^{N \times 1}$ (추론을 단순하게 유지하기 위해)의 경우, 해당 점화식은 $x_{k}=\exp (\Delta \tilde{a})+\Delta 1_{N} u_{k}$이다. 이 점화식의 요소별 전개(unrolling)는 ( $u$ 앞의 $\Delta$ 없이) $\left|x_{k, i}\right| \leq \sum_{j=0}^{k-1}\left|\exp \left(\Delta \tilde{a}_{i}\right)\right|^{j} u_{k-j, i}$를 생성하며, $k \rightarrow \infty$ 극한에서 $O\left(\Delta^{-1}\right)$를 제공한다. 따라서 $B$ 앞의 $\Delta$ 승수는 폭주(blow-up)를 피하기 위해 점화식을 효과적으로 스케일링한다. 이는 우리의 $\gamma$ 정규화 계수와 유사하다.

파라미터 공유는 필요하지 않다. 이산화의 결과로, $\tilde{A}$와 $\tilde{B}$를 모두 곱하는 $\Delta$ 파라미터는 학습 중에 점화식과 입력 투영을 결합한다. S4 ablation 연구에서, 우리는 이들을 두 개의 별도 파라미터로 분리하는 것(폭주를 보장하기 위해 동일한 초기화를 유지)이 성능을 저하시키지 않음을 발견했다. 이는 ODE 이산화 관점(파라미터 공유를 유도하는)이 S4 성능을 달성하는 데 필요하지 않음을 시사한다.

이 논의를 통해 우리는 (diagonal) state-space model의 성공이 선형 점화식과 복소수 대각 지수 행렬의 사용, 그리고 이산화에 의해 유도되는 정규화 및 초기화에 기인한다고 결론 내린다. 반면에, 파라미터 공유 또는 연속 시간 해석과 같은 이산화의 다른 인공물은 반드시 성능에 기여하지는 않는다.

5. Conclusion

본 논문에서는 **Linear Recurrent Unit (LRU)**이라는 새로운 RNN layer를 소개하고, 이를 깊은 시퀀스 모델의 핵심 layer로 효과적이고 효율적으로 사용하는 방법을 제시한다. 우리는 **바닐라 RNN의 선형화(linearization), 대각화(diagonalization), 안정적인 지수 파라미터화(stable exponential parameterization), 그리고 정규화(normalization)**에 이르는 일련의 단계별 수정 과정에 대한 이론적 통찰력과 광범위한 ablation 연구를 제공한다. 이러한 수정들은 특히 장거리 추론(long range reasoning)을 요구하는 task에서 성능을 크게 향상시킨다.
우리의 recurrence는 최신 **deep SSM(State Space Model)**과 유사점을 공유하지만, 잠재적인 연속 시간 시스템의 이산화(discretization)나 구조화된 전이 행렬(structured transition matrices)에 의존하지 않는다. 대신, 우리의 개선 사항은 Glorot-initialized RNN에서 시작하여 딥러닝 커뮤니티에서 표준적으로 사용되는 초기화 및 forward pass 분석 논증에서 직접적으로 도출된다. 우리의 최종 모델은 모든 LRA task에서 최신 deep state-space model (예: S4 또는 S5)의 성능과 일치한다.