WhiteMatter: All-to-All Cross-Layer Connections via KV Mixing

텍스트를 생성하는 현대의 인공지능은 정보가 일련의 쌓인 층(layers)을 통해 처리되는 트랜스포머(Transformer)라는 구조에 의존합니다. 이를 공장의 조립 라인에 비유해 본다면, 원재료가 1번 스테이션을 지나 2번 스테이션으로 이동하며 각 스테이션에서 제품을 정교하게 다듬는 과정과 같습니다. 표준적인 설정에서는 시스템이 다음 텍스트로 넘어갈 때, 각 스테이션은 이전 단계의 바로 위 스테이션이 수행한 작업만을 되돌아볼 수 있습니다. 이는 스택의 상단에서 형성된 더 깊고 복잡한 통찰력이 새로운 정보를 처리할 때 하위의 초기 스테이션에는 접근할 수 없음을 의미합니다. 시스템은 풍부한 이해의 역사를 만들어냈음에도 불구하고, 그 전체적인 깊이를 효율적으로 사용하는 데 어려움을 겪으며, 이로 인해 장기적인 문맥을 추적하거나 복잡한 문제를 해결하는 능력이 제한됩니다.

서던 캘리포니아 대학교의 연구진은 이러한 병목 현상을 해결하기 위해 화이트매터(WhiteMatter)라고 불리는 새로운 아키텍처를 제안했습니다. 화이트매터는 각 층을 고유한 깊이에 국한시키는 대신, 모든 부분이 과거의 모든 층으로부터 얻은 혼합된 요약 정보에 접근할 수 있도록 허용합니다. 연구팀은 동적 라우터(dynamic router) 역할을 하는 메커니즘을 구축하여, 이전 단어의 모든 층에서 나온 은닉 상태(hidden states)를 가져와 이를 더 작은 규모의 공유 채널들로 혼합합니다. 그러면 현재 단어의 각 층은 이 채널들 중 하나를 선택하여 읽어 들입니다. 이 설계는 인간의 뇌가 백질(white matter) 섬유를 통해 먼 영역들을 연결하는 방식을 모방한 것으로, 이를 통해 네트워크의 얕은 부분들이 과거의 깊고 정제된 정보를 받을 수 있게 하는 동시에, 깊은 부분들은 여전히 가공되지 않은 즉각적인 문맥에 접근할 수 있게 합니다.

연구진은 80억 개의 텍스트 토큰으로 구성된 방대한 데이터셋을 통해 이 시스템을 처음부터 학습시켜 테스트했습니다. 그들은 새로운 화이트매터 모델을 동일한 크기의 표준 모델 및 훨씬 더 큰 표준 모델들과 비교했습니다. 결과에 따르면, 16개의 층을 가진 화이트매터 모델은 24개의 층을 가진 표준 모델보다 성능이 뛰어났으며, 다음 단어를 예측하는 데 있어 유의미하게 낮은 오차율을 달성했습니다. 이러한 개선은 연구진이 시스템에 필요한 메모리 저장 공간을 절반으로 압축했을 때도 달성되었으며, 이는 모델이 높은 성능을 유지하면서도 더 적은 메모리를 사용할 수 있음을 입증했습니다. 또한 연구는 이 시스템이 많은 단어를 한꺼번에 처리할 수 있는 특정 반복적 방법을 사용하여 효율적으로 학습될 수 있음을 보여주었는데, 이는 이러한 깊은 연결 구조를 위해 보통 요구되는 '하나씩 처리하는 방식'을 피할 수 있게 해줍니다.

핵-혁신은 시스템이 과거 단어의 메모리를 처리하는 방식에 있습니다. 일반적인 설정에서는 단어의 메모리가 층별로 저장되며, 다음 단어는 동일한 깊이의 메모리에만 접근할 수 있습니다. 하지만 화이트매터는 이 규칙을 깨고, 단어 자체의 내용에 따라 업데이트되는 메모리 채널 풀을 생성합니다. 라우터는 과거 단어의 모든 층에 대한 은닉 상태를 분석하여 이를 몇 개의 채널로 혼합합니다. 이 채널의 수는 조정 가능하며, 연구진은 채널 수가 전체 층의 수보다 적을 때 모델의 복잡한 텍스트 이해 능력을 저해하지 않으면서도 메모리 점유율을 줄일 수 있다는 것을 발견했습니다. 이러한 유연성 덕분에 시스템은 모든 층의 이력을 상세히 저장할 필요 없이 가장 관련성 높은 혼합물만을 저장함으로써 더 효율적으로 작동할 수 있습니다.

이 과정을 학습 단계에서 구현하기 위해 연구진은 순환 의존성(circular dependency) 문제를 극복해야 했습니다. 단어의 메모리는 그 단어의 처리에 달려 있고, 처리는 메모리에 달려 있기 때문에, 시스템은 단순히 한 번의 패스(pass)로 모든 것을 계산할 수 없습니다. 연구팀은 텍s트를 그룹 단위로 처리하는 가우스-세이델 반복법(cyclic Gauss–Seidel iteration)이라는 방법을 개발했습니다. 이를 통해 시스템은 텍스트를 단계별로 업데이트할 수 있으며, 시퀀스 내의 후속 그룹들이 동일한 패스 내에서 이전 그룹으로부터 업데이트된 정보를 즉시 사용할 수 있게 됩니다. 이 접근 방식은 동일한 정확도에 도달하기 위해 훨씬 더 많은 패스를 요구하는 기존 방법들보다 유의미하게 빠르다는 것이 밝혀졌으며, 이를 통해 이러한 깊고 상호 연결된 모델의 학습을 표준 하드웨어에서도 가능하게 만들었습니다.

실험 결과, 이 아키텍처는 생성되는 텍스트의 품질을 개선할 뿐만 아니라 다양한 추론 작업에서의 성능도 향상시킨다는 것이 확인되었습니다. 언어 이해를 위한 표준 벤치마크 테스트에서 화이트매터 모델은 동일한 깊이의 표준 모델들을 지속적으로 능가했으며, 심지어 더 큰 표준 모델들보다도 뛰어난 성능을 보였습니다. 16개 층 모델에 대해 8개의 채널만을 사용하는 절반 메모리 구성은 전체 캐시를 사용할 때보다 훨씬 적은 메모리를 사용하면서도 대부분의 성능 이점을 유지했습니다. 이는 모든 깊이로부터 정보를 동적으로 혼합하는 능력이 단순히 메모리 저장 공간의 크기를 키우는 것보다 더 가치 있다는 점을 시사합니다.

또한 연구는 시스템이 다양한 학습 조건 하에서 어떻게 작동하는지 탐구했습니다. 연구진은 학습 중에 데이터를 순회하는 특정 방식이 최종 성능에 큰 영향을 미친다는 것을 발견했습니다. 최종적인 단계별 디코딩 과정을 밀접하게 모방하는 스케줄로 학습된 모델이 더 우수하고 안정적이었습니다. 그러나 더 단순한 학습 스케줄을 사용하더라도 화이트매터 모델은 표준 모델보다 뛰어난 성능을 보였는데, 이는 아키텍처 자체의 변화가 개선의 주요 동력임을 나타냅니다. 연구진은 학습 과정이 표준 모델보다 더 많은 계산 능력을 요구하지만, 실제 텍스트를 생성하는 디코딩 과정은 거의 동일한 속도로 이루어져 실무 적용이 가능하다는 점을 언급했습니다.

광범위한 인공지능 맥락에서 이 연구는 딥러닝 모델이 정보를 시간에 따라 처리할 때 겪는 근본적인 한계를 다룹니다. 모든 층이 네트워크 전체의 역사를 합성된 뷰(view)로 접근할 수 있게 함으로써, 화이트매터는 깊은 아키텍처를 괴롭히는 고립 문제를 극복합니다. 이번 연구 결과는 시스템의 효율성이 단순히 크기나 깊이를 키우는 데서 오는 것이 아니라, 내부 구조 간에 정보를 어떻게 연결하고 공유하느냐에 달려 있음을 시사합니다. 과거의 상호작용에 대한 메모리를 데이터의 풍부함을 잃지 않으면서도 압축할 수 있는 능력은 더욱 유능하고 효율적인 언어 모델을 구축하는 유망한 길을 제시합니다.

연구진은 이 접근 방식이 심층-얕은 피드백(deep-to-shallow feedback)을 성공적으로 통합하여, 시스템이 표현의 전체 깊이를 활용할 수 있게 한다고 결론지었습니다. 또한 이것이 줄어든 메모리 점유율로도 가능하다는 것을 입증함으로써, 더 나은 성능을 위해서는 반드시 더 많은 저장 공간이 필요하다는 가설에 도전했습니다. 학습 과정에는 일부 추가적인 복잡성이 따르지만, 모델의 실제 사용 시 발생하는 정확도와 효율성의 이점은 상당합니다. 이 연구는 신경망의 서로 다른 부분들 사이의 연결 방식을 재고하는 것이 어떻게 성능의 획기적인 향상으로 이어질 수 있는지를 보여주는 구체적인 사례이며, 미래의 인공지능 시스템 발전을 위한 새로운 방향을 제시합니다.