[논문리뷰] Understanding the differences in Foundation Models

Dynamic Systems Framework for Architecture Comparison

Dynamic Systems Framework(DSF)

• Recurrent state space representation
• Linear Structured time-varying dynamical system
  $h_i = \Lambda_ih_{i-1} + B_iu_i \\y_i = C_ih_i + D_iu_i$
• Convolutional Form
  • diagonal은 lambda 항이 빠지고 D_i 항이 더해짐 (recurrent 하기 전 가장 처음의 hidden state와 계산되는 부분이라서)
  • 해당 커널은 attention matrix와 같은 사이즈를 가짐
  • recurrent view에 따른 커널은 causal convolution 커널이 됨. 만약 non-causal kernel을 사용한다면 state update 과정을 바꿔줘야함

Architecture reformulation

• Attention

  • key, query, value에 대한 Mapping으로 정의할 수 있음

    $y_i = \sum_{j=0}^i {\phi(q_i)^T\psi(k_j) \over \eta(q_i, \mathbf{k}_i)}W_vu_j$
    $\zeta(q_i^T, k_j) = {\phi(q_i)^T\psi(k_j) \over \eta(q_i, \mathbf{k}_i)}$

  • 이를 dynamic system of dimension N=nd로 정의하면 아래와 같이 정의할 수 있음

    $\Lambda_i = {\eta(q_{i-1},k_{i-1} \over \eta(q_i, k_i)} \mathbb{I}_{nd} \\ B_i = \left({1 \over \eta(q_{i-1}, k_{i-1})} \mathbb{I}_d \otimes \psi(k_j) \right) W_v \\ C_i = \mathbb{I}_d \otimes \phi(q_i)^T$

  • linear attention이라면 각 mapping이 linear하게 정의됨

  • softmax라면 infinite dimensional 하게 정의됨 (softmax의 차원이 무한하므로) ⇒ 근사해서 finite dimension에서 관리하게됨→ seperable attention으로 정의가 가능함

• State Space Model

  • S6

    $\Lambda_i = e^{-(\Delta_i \otimes \mathbb{I}_n) \odot A} \\ B_i = \Delta_i \otimes b_i \\ C_i = \mathbb{I_d} \otimes c_i^T \\ \Delta_i = \mathrm{diag(softplus(}W_\Delta(W_uu_i)+b_\Delta))\\ b_i = W_Bu_i \\ c_i = W_Cu_i \\ b_\Delta : bias$

  • Delta_i 는 우리가 알고 있는 step size, learnable with input

  • b_i, c_i 역시 input-dependent한 state로 정의, B_i는 여기에 더불어 delta에 따라 변화함

  • 기존 식에 적용하면 우리가 알고 있는 mamba 모델로 정의됨

  • 또한 여기서 Lambda_i 는 만약 SSD처럼 아키텍처를 가지고 있다면 reversed sigmoid로 정리될 수 있음! → elementwise

• RNN

  • qLSTM : input processing의 tanh를 제거하고 다시 작성하면 아래와 같이 작성할 수 있음

  • hidden dimension은 input channel의 수와 같음 (즉 차원이 nd 일 때 n=1인 경우)

    $\Lambda_i = \mathrm{diag}(\sigma(W_fu_i)) \\ B_i = \mathrm{diag}(\sigma(W_iu_i))\\ C_i = \mathrm{diag}(\sigma(W_ou_i))$

  • RG-LRU : SSM과 유사한 형태로 정의됨. state expansion은 정의하지 않으므로 n=1 (state expansion : n > 1)

    $\Lambda_i = e ^{-cr_i \odot softplus(A)}\\ B_i = \sqrt{1 - \Lambda_i^2} \odot \mathrm{diag}(\sigma(W_Bu_i))\\ C_i = \mathbb{I}_d$

Architecture comparison : Theoretical and Experimental Results

Softmax attention vs Seperable attenion

• seperable attention : query-key 연산 전에 key-value 연산을 먼저해서 계산 복잡도를 줄이고자 함
• DSF에서 표현되는걸 보면 왜 seperable attention이 finite kernel 로 표현되어야하는지를 알 수 있음
• softmax가 당연히 성능은 더 좋음 (n = infinite)로 가니까 가질 수 있는 정보량이 많음.
• 여기서 왜 SSM과 RNN이 따라잡을 수 있는지 설명이 되는데, SSM과 RNN은 linear하게 구성이 되어 n을 증가시키면 점차 softmax attention의 성능을 따라잡을 수 있게 됨. < 진짜 그냥 state size가 문제인 것일 수도 있음
  • 결국 trade off가 존재한다는거고 거기서의 최적화가 필요하다는 걸로 귀결됨… 어떻게 정보량을 최대화할 것인가?

Generalized Linear Attention vs S6

• S6의 b_i, c_i = k_i, q_i
• state expansion of S6 = Hidden dimension in attention → 동일한 output matrix C_i를 도출하게 된다.
• 차이점
  • number of parameters
    • Transition matrix Lambda (Mamba 에서는 A) 가 가지는 파라미터 개수가 S6은 채널 수와 동일 ($\Delta_i$에 의해 결정되는데 이것이 d개의 파라미터를 가짐, d)이고 attention 은 1. ($\eta(q_i, \mathbf{k})$가 유일한 파라미터이며 정의 상 스칼라임)
    • 다만 Mamba paper에서 d=1로 해도 성능이 어느정도 보장된다고 주장한 바 있음
  • normalization strategy vs discretization step
    • attention에서는 normalizing map이 존재하나 연산 과정에서 스칼라이므로 cancel 되어 사라짐 (convolution 커널에 이러한 사항이 반영됨)
    • 반면 S6에서는 discretization parameter Delta만이 공유되므로 Lambda_i와 B_i에서 사라지지 않음 → selectivity를 가능케하나, input-dependent한 요소에 대해 각 구조에서 normalizing이 다르게 반영될 수 있음
  • Multi-head

    • attention은 multihead로 바꿀 시 헤드 개수에 따라 파라미터가 증가하며 성능 향상에 도움이 되나 S6은 그런 경향성이 미미한 것으로 보인다.

    • 따라서 S6에는 recursive structure가 더 중요한 역할을 하거나 eta에서의 정규화 파라미터 영향이 더 클 수 있따. 따라서 SSD, Linear attention, normalized attention 을 비교한다.

    • normalized attention ← mamba에서 영감을 받아서 각 요소마다 다르게 normalize 할 수 있도록 파라미터를 추가

      $\phi(q_i) = q_i,\ \psi(k_j) = k_j), \ \eta(u_i) = e ^{W_nu_i}$

    • 결과를 보면 softmax가 제일 좋은데 SSD도 좋은 성능을 보이는걸 볼 수 있음!