QuadMLP - Transformer FFN를 새로운 관점에서 바라보다

Transformer를 설명할 때 대개 attention가 가장 먼저 스포트라이트를 받는다. 하지만 attention이 가져온 문맥을 보고 어떤 지식을 꺼내 residual stream에 보탤지 결정하는 FFN(Feed-Forward Network)도 못지않은 관심이 필요하다고 느낀다.

FFN은 거대한 연상 사전이라고 생각하면 쉽다. token의 현재 표현이 이 레이어에 들어오면, FFN은 그 안에서 여러 단서를 찾는다. 어떤 단서는 “이 문맥에서는 복수형을 기대하라”는 규칙일 수 있고, 어떤 단서는 “이 이름 뒤에는 특정 사실이 잘 이어진다”는 기억일 수 있다. 단서를 통해 꺼낸 결과는 다시 hidden state에 더해져, 다음 레이어와 최종 예측에 사용된다.

물론 실제 FFN 안에 사람이 읽을 수 있는 단어별 사전 항목이 줄지어 있는 것은 아니다. 넓은 intermediate 공간과 down projection에 분산된 연속적인 연상 기억에 가깝다. 그래도 이 관점은 유용하다. attention이 “어디서 단서를 가져올까”라면, FFN은 “그 단서와 현재 상태를 보고 사전의 어느 항목들을 얼마나 꺼낼까”에 가깝다.

Llama 3.1 계열의 FFN에는 단순한 한 갈래 MLP 대신 gated MLP, 특히 SwiGLU가 쓰인다. SwiGLU를 수식으로 표현하면 다음과 같다.

up   = W_up x
gate = W_gate x
FFN(x) = W_down ( up ⊙ SiLU(gate) )

여기서 흔한 직관은 이렇다. up이 꺼낼 후보를 만들고, gate가 그 신호를 특정 세기로 통과시키거나 아예 막는다. 충분히 유용한 설명이다. 하지만 이 설명에만 매몰되면 gate를 그저 “사전 항목의 세기를 제한하는 문지기”로만 보게 된다.

나는 두 projection의 상호작용을 다른 관점에서 바라보고자 했다. up과 gate값을 각각 하나의 축 위에 표현한다면, 둘의 상호작용은 좌표평면으로 시각화 할 수 있지 않을까?

전투기, 화물선, 벌새, 달팽이

FFN을 잠시 떠나 좌표평면을 생각해보자. 두 축만 있으면 우리는 자연스럽게 네 사분면을 구분할 수 있다. 가로축을 크기(오른쪽으로 갈수록 큼), 세로축을 속도(위로 갈수록 빠름)로 잡아 보자.

전투기와 화물선은 모두 크지만, 속도라는 축에서 완전히 다르고 벌새와 달팽이는 모두 작지만, 마찬가지로 다른 존재다. “크기”만으로 속도를 설명할 수 없고, “속도”만으로 크기를 설명할 수 없다. 따라서 두 특성이 조합하는 네 칸은 각기 다른 새 정보를 파생시킬 여지가 있다.

이 비유를 FFN에 적용하면 Transformer가 왜 SwiGLU를 활용하는지 알게된다. 가장 단순한 bilinear interaction은 up ⊗ gate다. 이 곱은 두 축의 상호작용을 하나의 값으로 표현한다.

위의 그림에서 확인할 수 있듯 SiLU 없는 up ⊗ gate는 1·3사분면과 2·4사분면이 경쟁하는 구도를 만든다. 전투기와 달팽이는 서로 정반대의 특성을 갖지만 동일한 부호를 공유하고, 벌새와 화물선도 마찬가지다. 이는 우리의 직관과 반대된다.

반면 SwiGLU는 gate가 음수인 2·3사분면을 강하게 억제하는 방식으로 이 충돌을 회피한다. 그 결과 주된 경쟁 구도가 1·4사분면의 특성 차이에 집중되고, 정반대 특성과의 불필요한 신호 공유는 끊어진다.

다시 말해, SwiGLU는 up과 gate가 파생한 정보 중 일부를 배제함으로써 두 projection의 상호작용을 down값 하나의 세기로 변환한다. 현 토큰 표현에 맞춰 새 정보를 충돌 없이 residual stream에 주입하는 데에 SiLU가 열쇠 역할을 하는 것이다.

하지만 동시에, 여기에 SwiGLU의 빈틈이 있다. 각 네 사분면은 모두 유용한 정보를 생산하지만 우리는 그 중 절반을 포기한다.

한 평면을 여러 채널로 읽는 방법

QuadMLP는 이 문제를 해결하기 위해 up, gate 쌍이 만든 평면을 하나의 값으로 읽는 대신, 여러 채널로 나누어 읽는다. 각 채널은 자신이 맡은 사분면에서만 반응하고, 그 안에서는 |up|과 |gate|가 함께 클수록 신호가 강해진다. 그러면 같은 두 projection으로도 어떤 구분을 보존할지 채널마다 다르게 정할 수 있다.

아래 그림은 이 방식을 이해하기 위한 세 예시다. 한 쌍의 projection에서 어느 사분면을 같은 출력 채널로 묶을지에 따라, 같은 입력 평면도 서로 다른 관점으로 읽을 수 있다.

Type A는 대각선 쌍(1·3, 2·4)을 두 채널로 읽는다. 여기서는 두 축의 부호가 같은지, 다른지가 채널을 가르는 기준이다. Type B도 두 채널을 쓰지만 gate의 부호를 기준으로 1·4와 2·3을 나눈다. 기존 SwiGLU가 주로 읽는 쪽과 억제하던 쪽을 각각 독립된 채널로 바라보는 예다. Type C는 네 사분면을 모두 별도 채널로 두어, 가장 세밀하게 구분을 보존한다.

따라서 QuadMLP는 더 복잡한 비선형성을 위한 제안이라기보다, 이미 계산한 up과 gate의 결합을 여러 관점에서 읽어 residual stream에 전달하는 효율성을 개선하려는 시도라고 볼 수 있다.