Retrofitting Linear Attention into Diffusion Language Models

인공지능의 세계를 아주 똑똑한 로봇 사서가 이야기를 쓰고 있는 거대하고 북적이는 도서관이라고 상상해 보세요. 수년 동안 이 사서가 일하는 표준적인 방식은 "자기회귀(autoregressive)" 방식이었습니다. 즉, 단어 하나를 쓰고, 멈춰서 방금 쓴 모든 것을 생각하고, 다음 단어를 쓰고, 다시 멈춰서 생각하는 식이었죠. 이는 마치 벽돌을 하나씩 쌓아 올리며 다음 벽돌을 놓기 전에 접착제가 마르기를 기다리는 것과 같았습니다. 이 방식은 신뢰할 만했지만, 긴 이야기를 쓸 때는 매우 느렸습니다.

최근에는 "디퓨전(Diffusion)"이라 불리는 새로운 스타일이 등장했습니다. 디퓨전 사서는 벽돌을 하나씩 쌓는 대신, 물음표로 가득 찬 빈 페이지에서 시작하여 전체 이야기를 한꺼번에 추측하고, 단어들이 의미를 갖출 때까지 그 추측을 반복해서 정교하게 다듬습니다. 이는 마치 팀 단위의 예술가들이 동시에 벽화를 그리는 것과 같습니다. 그들은 이전 벽돌이 마르기를 기다릴 필요가 없기 때문에 훨씬 더 빠르게 작업할 수 있습니다. 하지만 이 빠른 방식에도 문제점이 있습니다. 벽화가 커질수록, 예술가들은 새로운 단어가 기존의 것들과 잘 어울리는지 확인하기 위해 이미 결정된 모든 단어를 끊임없이 되돌아봐야 합니다. 이는 마치 101페이지를 쓰면서 10,000페이지짜리 책 전체를 기억하려고 애쓰는 것과 같으며, 결국 모두의 속도를 늦추고 도서관의 메모리를 가득 채우게 됩니다.

이것이 바로 연구진이 해결하고자 했던 퍼즐입니다. 그들은 이렇게 질문했습니다. "사서가 과거를 기억하는 방식을 바꿈으로써 이 빠른 디퓨전 방식을 더 빠르게 만들 수 있을까?" 그들의 해답은 "블록-하이브리드 어텐션(Block-Hybrid Attention)"이라는 영리한 기술이었습니다. 그들은 사서가 현재의 문단이 매끄럽게 흐르도록 하기 위해서는 현재의 블록을 완벽하게 기억해야 하지만, 이전의 장들을 모든 단어 하나하나 고해상도로 재독할 필요는 없다는 사실을 깨달았습니다. 대신, 이전의 장들을 작고 고정된 크기의 "요약 시트"로 요약할 수 있습니다. 이를 통해 사서는 전체 역사의 무게에 짓눌리지 않고도 빛의 속도로 새로운 부분을 써 내려갈 수 있습니다.

핵심 아이디어: 이중 속도 메모리 시스템

이 논문은 기존의 강력한 디퓨전 언어 모델(구체적으로는 160억 개의 파라미터를 가진 LLaDA 2.1 모델)을 처음부터 다시 학습시키지 않고도 이 "요약 시트" 전략을 사용할 수 있도록 개조(또는 업그레이드)하는 방법을 소개합니다.

모델의 두뇌를 20개의 서로 다른 사고 계층(layer)을 가진 것으로 생각해 보세요. 원래 모델에서는 모든 계층이 매우 주의 깊은 사서처럼 작동하여, 현재의 단어를 이해하기 위해 이전의 모든 단어를 읽습니다. 이는 정확하지만 무거운 방식입니다. 연구진의 혁신인 블록-하이브리드 어텐션은 작업을 두 가지 모드로 나눕니다:

1. "현재" 모드 (정밀 어텐션): 모델이 현재 생성 중인 단어 블록을 작업할 때는 "매우 주의 깊은" 모드를 유지합니다. 문장이 완벽하게 의미를 갖도록 하기 위해 현재 문단의 모든 단어를 살피는 표준적이고 무거운 메모리를 사용합니다.
2. "과거" 모드 (선형 어텐션): 모델이 이전 블록들(완료된 장들)을 기억해야 할 때는 "선형 어텐션"으로 전환합니다. 책 전체를 다시 읽는 대신, 고정된 크기의 요약 상태를 참조합니다. 이는 책 전체를 다시 읽는 대신 책의 색인이나 한 페이지 분량의 요약본을 보는 것과 같습니다. 이 요약본은 책이 아무리 길어져도 크기가 일정하게 유지됩니다.

구현 방법: 2단계 업그레이드

연구진은 단순히 부품을 교체하고 결과가 좋기를 바란 것이 아니라, 업그레이드 과정에서 모델의 지능을 잃지 않도록 세심한 2단계 과정을 거쳤습니다.

• 1단계: 섀도우 트레이닝 (Shadow Training). 그들은 원래의 똑똑한 모델(스승)을 가져와서 고정(freeze)했습니다. 그런 다음, 20개의 어텐션 계층 중 6개를 새로운 "블록-하이브리드" 버전(학생)으로 교체했습니다. 학생은 "오염된(corrupted)" 버전의 텍란 입력을 사용하여 스승의 출력을 정확히 흉내 내도록 훈련되었습니다. 이는 마치 학생이 마스터 셰프의 그림자처럼 따라 하며, 셰프의 정확한 움직임을 복제하고 셰프가 만든 것과 똑같은 맛을 경험하되, 오직 자신이 맡은 특정 요리에 대해서만 수행하는 것과 같았습니다. 이 단계는 약 24시간이 소었습니다.
• 2단계: 미세 조정 (Fine-Tuning). 학생이 스승을 흉내 내는 법을 배우자, 연구진은 전체 모델이 다시 함께 작동하도록 했습니다. 그들은 LoRA(Low-Rank Adaptation)라는 기술을 사용하여 모델의 연결 부위를 아주 정밀하게 조정했습니다. 이는 새로운 부품이 기존 부품과 대화할 때 발생할 수 있는 작은 결함을 수정하기 위해 업그레이드된 모델에 빠른 "광택"을 내는 것과 같았습니다. 이 단계는 약 6시간이 걸렸습니다.

결과: 더 빠르고, 가볍고, 여전히 똑똑함

연구진은 업그레이드된 모델인 LLaDA-HYBRID를 테스트하여, 이 모델이 여전히 글을 쓰고 문제를 해결하는 데 유능한지, 그리고 실제로 더 빠른지를 확인했습니다.

여전히 사고할 줄 아는가?
놀랍게도 그렇습니다. 20개 중 6개의 계층을 이 "요약 시트" 방식으로 교체했음에도 불구하고, 모델의 성능은 원래 모델과 매우 유사하게 유지되었습니다.

• 코딩 테스트인 HumanEval에서 원래 모델은 **75.6%**를 기록했고, 하이브리드 모델은 **72.0%**를 기록했습니다.
• 수학 테스트인 CMATH에서 원래 모델은 88.3%, 하이브리드 모델은 **86.7%**를 기록했습니다.
• 흥미롭게도, 또 다른 코딩 테스트인 **MBPP+**에서 하이브리드 모델은 57.7%에서 63.0%로 오히려 성능이 향상되었습니다.
  논문은 모델이 가장 어려운 추론 과제(GPQA-Diamond라는 어려운 과학 퀴즈)에서는 약간의 성능 저하(38.9%에서 30.8%로 하락)가 있었지만, 코드 작성 및 수학 문제 해결 능력은 대체로 보존했다고 설명합니다.

실제로 더 빠른가?
여기서 마법 같은 일이 일elle납니다. 모델이 더 이상의 단어를 생성하기 위해 대화의 전체 역사를 다시 읽을 필요가 없기 때문에, 더 많은 요청을 동시에 처리할 수 있습니다.

• 모델이 초당 얼마나 많은 단어를 생성할 수 있는지 테스트했을 때, 하이브리드 모델은 최고 성능(128개의 동시 요청 처리 시)에서 원래 모델보다 1.7배 더 빨랐습니다.
• 원래 모델은 이야기가 길어짐에 따라 메모리 사용량이 증가하여 결국 더 이상 사용자를 수용하지 못하는 한계점에 도달했습니다. 반면, 고정된 크기의 "요약 시트" 덕분에 하이브리드 모델은 메모리가 부족해지기 전까지 더 많은 동시 요청을 처리할 수 있었습니다.

시사점

이 논문은 우리가 항상 빠르고 새로운 AI를 처음부터 새로 만들 필요는 없다는 점을 시사합니다. 대신, 기존의 강력한 AI를 가져와 "하이브리드 메모리" 시스템을 부여할 수 있습니다. 현재의 순간에 날카롭게 집중하면서 동시에 과거를 압축된 형태로 요약함으로써, 우리는 이러한 모델을 훨씬 더 빠르고 효율적으로 만들 수 있습니다. 연구진은 이 업그레이드가 공공 데이터를 사용하여 약 60시간의 훈련만으로 가능하다는 것을 발견했으며, 이는 지능을 크게 희생하지 않으면서도 AI 서비스를 더 빠르고 저렴하게 만들 수 있는 유망한 경로를 제시합니다. 다만, 저자들은 이것이 시작일 뿐이라고 언급했습니다. 그들은 6개의 계층만을 업그레이드했으며, 향-후 연구를 통해 더 많은 계층을 업그레이드하거나 더 긴 이야기를 처리하는 방법을 찾을 수 있을 것입니다.