RIFT - 트랜스포머의 읽기 쓰기 경쟁을 멈추다

Transformer엔 여태까지 당연하게 받아들이고 지나갔지만 자세히 들여다보다 보면 조금 이질적으로 느껴지는 부분이 하나 있다. 각 레이어는 token별 hidden representation을 받아 delta representation을 생성하고, 이 갱신은 대체로 하나의 residual stream 위에서 누적된다. 마지막 hidden state는 vocabulary embedding과 비교되어 다음 토큰의 logit이 된다.

이 구조는 아름답고 단순하다. 그래서 더 오래 살아남았을 것이다. 하지만 나는 예전부터 단일 residual stream 위의 hidden representation이 서로 다른 두 개의 일을 너무 자연스럽게 떠안고 있는 것 아닌가 하는 의문이 있었다.

첫 번째 역할은 다음 토큰을 고르기 위한 답안을 만드는 일에 가깝다. 마지막 hidden state는 다음 토큰의 정답 embedding과 잘 맞아야 한다. 모델이 "다음에 올 단어는 이것"이라고 말하려면, 학습이 진행될수록 각 레이어의 delta representation이 마지막 logit 계산에 유리하게 작동하도록 조금씩 압력을 받는다.

두 번째 역할은 참조 가능한 문맥을 남기는 일이다. 같은 hidden representation은 이후 레이어의 attention에서 key와 value가 되고, 이후 token position들은 query를 통해 그 정보를 읽는다. 즉 이 표현은 나중에 다시 읽힐 문맥을 잃지 않아야 한다.

여기서 중요한 것은 "정답 token embedding에 가까워지는 것"과 "문맥 전달에 사용하기 좋은 상태로 남는 것"이 항상 같은 목표는 아닐 수 있다는 점이다. 기존 Transformer는 이 두 요구를 하나의 residual stream 안에서 동시에 만족시키도록 훈련한다. 물론 hidden dimension이 충분히 크면 두 역할이 공존할 수 있다. 실제로 대형 모델들은 그 방식으로 놀라운 성능을 보여줬다. 하지만 그 성공이 곧 구조적으로 낭비가 없다는 뜻은 아니다.

나는 이 지점에서 두 가지 가능성을 의심했다. 하나는 표현 병목이다. 하나의 hidden representation이 logit을 잘 만들기 위한 방향과 attention에 좋은 방향을 동시에 담아야 한다면, 모델은 같은 차원 안에서 서로 다른 목적의 정보를 타협시켜야 한다. 다른 하나는 레이어 깊이 방향에서 생기는 계산 역할의 비효율적인 분화다. 깊은 층들이 마지막 logit에 맞춰 표현을 수습하는 데 치우친다면, 깊이를 아무리 늘려도 문맥을 수집·전파하는 역할에 일부 손해를 볼 수 밖에 없다.

RIFT(Read–Write Interleaved Flow Transformer)는 이 의문에서 출발했다. 핵심 아이디어는 단일 residual stream 위에서 함께 섞이던 두 역할을 attention state에 영향을 주는 흐름과, attention state를 읽기만 하는 흐름 둘로 분리하는 것이다.

write stream은 앞선 position들에서 문맥을 수집하고, 이후 position들이 attention으로 참조할 수 있는 상태를 만든다. 어떤 정보가 뒤에서 다시 읽혀야 하는지 정리해 key와 value로 남기는 흐름이다.

read stream은 그 상태를 읽어 최종 예측을 만들어 간다. 앞선 position들에서 모인 정보를 attention으로 가져오되, 자신이 하는 일은 다음 토큰을 고르기 좋은 hidden state로 점점 정렬되는 쪽에 가깝다.

조금 다르게 말하면, RIFT는 "기억할 표현"과 "말할 표현"을 분리하려는 시도다. 기존 Transformer가 한 장의 포스트잇에 나중에 참고할 메모와 지금 제출할 답안의 풀이를 서로 경합하듯 덮어썼다면, RIFT는 두 장을 준비해 한쪽에는 뒤에서 다시 읽을 맥락을 남기고 다른 한쪽에는 지금 답을 쓰기 위한 표현을 계속 고쳐 나간다.

물론 역할이 다르다고 무조건 분리하는 것이 답은 아니다. 예를 들어 검색 임베딩에서 같은 파라미터 예산이라면 dual encoder보다 single encoder가 더 나은 균형을 만들 수 있듯, Transformer의 단일 residual stream도 이미 꽤 효율적인 타협일 수 있다.

문제는 이런 실험이 가볍게 돌려 볼 수 있는 종류가 아니라는 점이었다. 한동안은 아이디어만 적어 두고 있었다. AI 붐 이후 GPU 가격과 대여 비용이 너무 많이 올라서, 개인 프로젝트로 밀어붙이기에는 현실적인 벽이 컸다.

그런데 6월에 기회가 생겼다. 한국창업진흥원에서 B200을 약 6개월간 대여해주는 지원을 받을 수 있게 됐다. 오랫동안 머릿속에서만 굴리던 구조를 이제는 실제 코드와 실험 로그로 검증할 수 있게 된 것이다.

이 프로젝트에서 내가 확인하고 싶은 것은 단순히 "RIFT가 성능이 더 좋다"가 아니다. 지금의 Transformer가 이미 충분히 좋은 해답이라면, RIFT는 실패를 통해 그 이유를 알려줄 것이다. 반대로 두 흐름의 분리가 작은 모델에서라도 일관된 이득을 보인다면, 우리는 Transformer 내부의 계산을 조금 다른 방식으로 설계할 여지를 얻게 된다.