로봇에게 이야기 쓰기를 가르치려 한다고 상상해 보세요. 오랫동안 이를 수행하는 두 가지 주요 방법이 있었으며, 둘 다 치명적인 결함을 지니고 있습니다.
두 가지 구식 방법
"한 사람 쇼" (Decoder-Only): 이는 책을 읽고 즉시 다음 문장을 쓰려 하는 학생과 같습니다. 글쓰기 속도는 매우 빠르지만, 글을 쓰면서 지금까지 읽은 모든 내용을 머릿속에 기억해야 합니다. 이야기가 길어지면 뇌 (기억) 가 과부하가 걸려 속도가 느려집니다.
"두 사람 팀" (Encoder-Decoder): 이는 독자와 작가로 구성된 팀과 같습니다. 독자는 먼저 책 전체를 읽고 메모를 한 뒤, 그 메모를 작가에게 건네줍니다. 이는 작가의 뇌 에너지를 절약해 주지만, 독자는 매우 까다롭습니다. 그들은 약 15% 의 단어 ( "손상된" 부분) 에 대해서만 메모를 하고 나머지는 무시합니다.这意味着 팀이 대부분의 이야기를 무시하기 때문에 학습 속도가 느립니다.
새로운 해결책: "블록 기반 이중 디코더"
이 논문의 저자들은 두 세계의 장점을 모두 얻으려 하는 새로운 팀 구조를 제안합니다. 이를 블록 기반 이중 디코더라고 부릅니다.
간단한 비유를 들어 작동 방식을 설명해 보겠습니다.
긴 소설을 읽는다고 상상해 보세요. 하지만 한 글자씩 읽는 대신, 이를 블록 (장이나 장면과 같은) 으로 나눕니다.
컨텍스트 디코더 (독자): 이 부분은 특정 지점까지의 전체 이야기를 읽습니다. 배경과 등장인물을 이해하기 위해 앞장들을 빠르게 훑어보는 빠른 독자와 같습니다. "과거"만 읽기 때문에 한 번에 책 전체를 기억할 필요가 없습니다. 요약본을 생성합니다.
생성 디코더 (작가): 이 부분은 독자로부터 받은 요약본과 현재 텍스트 블록을 받아 이야기의 다음 부분을 씁니다.
마법의 트릭: "이중 인과적" 블록
비밀은 이야기를 어떻게 나누는지에 있습니다. 그들은 텍스트를 조각 (블록) 으로 잘라냅니다.
단일 조각 내부에서는 작가가 그 조각 안의 모든 단어를 볼 수 있습니다 (그룹 토론과 같습니다).
하지만 이전 조각을 볼 때는 작가가 원본 단어가 아닌 독자가 제공한 요약본만 볼 수 있습니다.
왜 이것이 중요한가요?
낭비되는 학습의 종식: 구식 "두 사람 팀"에서는 독자가 85% 의 단어를 무시했습니다. 이 새로운 시스템에서는 단 하나의 단어도 배울 기회를 놓치지 않습니다. 모델이 모든 토큰에 대해 "성적"을 받기 때문에 훨씬 더 빠르고 똑똑하게 학습합니다.
초월적인 메모리 효율성: 로봇이 실제로 이야기를 쓸 때 (추론 시) 책 전체를 기억할 필요가 없습니다. 독자의 요약본과 현재 블록만 기억하면 됩니다. 이는 필요한 메모리를 약 3 분의 2만큼 줄여줍니다. 이는 책가방에 도서관 전체를 싣는 것에서 단일 인덱스 카드만 들고 다니는 것으로 바꾸는 것과 같습니다.
속도: "독자" 부분은 시작 시 한 번만 실행된 후 유휴 상태로 남기 때문에, "작가" 부분은 훨씬 가볍고 빠릅니다. 이는 경주 시작을 위한 무거운 엔진을 가지되, 스프린트를 위한 경량화된 공기역학적 차체를 가진 것과 같습니다.
그들은 무엇을 발견했나요?
연구자들은 이 새로운 아키텍처를 구식 방법들과 비교하여 테스트했습니다.
학습: 새로운 모델은 속도면에서 금표준인 "한 사람 쇼"와 거의 비슷하게 잘 학습했으며, 구식 "두 사람 팀"보다 훨씬 더 잘 학습했습니다.
글쓰기 (추론): 실제로 텍스트를 생성할 때가 되었을 때, 새로운 모델은 스타였습니다. 구식 "두 사람 팀"보다 훨씬 적은 컴퓨터 메모리를 사용하면서도 매우 똑똑했습니다.
결론
이 논문은 작업을 두 개의 디코더로 나누고 텍스트를 블록으로 분할함으로써, 모든 단어에서 학습하면서도 (구식 효율 모델과 달리) 글을 쓸 때 메모리 문제에 휘말리지 않는 (구식 빠른 모델과 달리) 모델을 만들었다고 주장합니다. 이는 거대하고 비싼 컴퓨터 없이도 실행할 수 있는 고성능 로봇을 만드는 방법입니다.
기술 요약: 블록 기반 이중 디코더
문제 제기 이 논문은 트랜스포머 아키텍처 설계에서 훈련 효율성과 추론 효율성 간의 근본적인 트레이드오프를 다룹니다. 디코더 전용 모델은 뛰어난 훈련 효율성과 확장성을 제공하지만, KV 캐시 메모리 및 토큰당 계산 비용과 관련하여 높은 추론 비용을 겪습니다. 반면, 인코더 - 디코더 모델은 KV 캐시 감소로 인한 상당한 추론 시간 절감 효과를 제공하지만, 스패너 손상 (span corruption) 과 같은 목적 함수로 인해 ~15% 의 토큰만이 손실 신호를 생성하는 희소 감독과 효율적인 배치 처리를 방해하는 동적 시퀀스 길이로 인해 사전 훈련 중 희소 감독을 겪습니다. 프리픽스 LM 과 같은 기존 대안들은 이 격차를 메우려 시도하지만 여전히 많은 토큰이 훈련되지 않으며 인코더 - 디코더의 완전한 효율성을 달성하지 못합니다.
방법론: 블록 기반 이중 디코더 저자들은 완전한 손실 감독과 추론 효율성을 모두 달성하기 위해 두 개의 디코더 스택을 결합한 새로운 아키텍처인**블록 기반 이중 디코더 (Block-Based Double Decoders)**를 제안합니다.
아키텍처 구조:
컨텍스트 디코더: 전체 입력 시퀀스를 처리하고 각 토큰에 대한 인과적 잠재값 (ht) 을 출력하는 표준 인과적 디코더 전용 트랜스포머입니다.
생성 디코더: 컨텍스트 디코더의 인과적 잠재값, 원시 토큰 시퀀스, 그리고 블록 분할을 세 가지 입력으로 받는 두 번째 디코더 스택입니다.
이중 인과적 블록 기반 마스킹: 입력 시퀀스는 연속적인 블록으로 나뉩니다. 블록 내부에서는 생성 디코더가 인과적 자기 어텐션을 사용하고, 블록 간에는 이전 블록의 컨텍스트 디코더 잠재값에 대한 전체 교차 어텐션을 사용합니다. 이는 블록 내에서 그리고 이전 블록으로부터 어텐션이 흐르는 "이중 인과적" 구조를 생성합니다.
훈련 목적:
토큰을 마스킹하는 스패너 손상과 달리, 이 방법은 입력 시퀀스의모든 토큰이 순전파당 정확히 한 번씩 손실에 기여하도록 보장합니다.
패킹 후 시퀀스 길이가 정적으로 유지되어 동적 배치 처리와 토큰 패딩이 불필요해지므로 처리량이 향상됩니다.
저자들은 통합된 이중 키 어텐션 구현이 아직 제공되지 않았으므로, PyTorch 의 FlexAttention 을 사용하여 별도의 자기 어텐션 및 교차 어텐션 로지트를 계산한 후 사후적으로 로그 - 합 - 지수 (log-sum-exp) 정규화를 통해 결합하는 방식으로 어텐션 메커니즘을 구현합니다.
추론 메커니즘:
추론 시 컨텍스트 디코더는 프롬프트 (프리필) 에 대해 한 번 실행됩니다. 이는 표준 인과적 디코더이므로 자체 출력에 대한 KV 캐시가 필요하지 않습니다 (생성 디코더만 캐시가 필요합니다).
생성 디코더는 캐시된 컨텍스트 잠재값과 자체 이전 출력에 어텐션하며 토큰을 autoregressively 생성합니다.
이 분리는 디코더 전용 모델에서 일반적인 프리픽스 레벨 KV 캐싱의 재사용을 가능하게 하면서도 인코더 - 디코더 아키텍처의 메모리 이점을 유지합니다.
주요 기여
새로운 아키텍처: 컨텍스트 처리와 토큰 생성을 분리하기 위해 두 개의 디코더 스택을 활용하는 블록 기반 이중 디코더의 도입.
이중 인과적 마스킹: 효율적인 배치를 위한 정적 시퀀스 길이를 유지하면서 완전한 토큰 감독 (100% 손실 밀도) 을 가능하게 하는 새로운 어텐션 마스크.
추론 최적화: 이 아키텍처는 프리필 캐싱 능력을 희생하지 않으면서 (2/3 컨텍스트 디코더, 1/3 생성 디코더 분할 가정 시) 디코더 전용 모델 대비 KV 캐시 메모리 및 토큰당 계산을 약 2/3만큼 감소시킵니다.
훈련 효율성: 이 방법은 스패너 손상 훈련을 받는 인코더 - 디코더에 내재된 희소 감독 및 동적 배치 처리 문제를 피합니다.
실험 결과 저자들은 파라미터 수 (6.25M–100M) 및 토큰 예산 (62.5M–1B) 에 걸쳐 디코더 전용, 표준 인코더 - 디코더 (SED), 그리고 블록 기반 이중 디코더를 비교하는 확장 법칙 실험을 수행했습니다.
성능: 이중 디코더는 SED 베이스라인을 강력하게 능가하며 디코더 전용 모델과 밀접하게 추세를 따릅니다. 가장 큰 규모 (100M 파라미터, 1B 토큰) 에서 이중 디코더는 일치된 디코더 전용 모델보다 평가 손실이 **약 0.2 나트 (nats)**만큼만 나쁘며, 반면 SED 는 둘 다보다 약 0.7 나트 뒤처집니다.
확장 행동: SED 모델은 희소 감독으로 인해 파라미터를 늘려도 성능이 향상되지 않는 "데이터 제한 영역"에 있는 것으로 나타났습니다. 반면, 이중 디코더와 디코더 전용 모델은 용량 한계에 도달함에 따라 평탄화되는 표준 확장 법칙을 따랐습니다.
계산 트레이드오프: 추가 KV 프로젝션으로 인해 이중 디코더는 훈련 FLOPs 가 약간 증가합니다 (특정 구성에서 디코더 전용 대비 약 2.4% 더 많음). 그러나 이는 추론 이점에 비해 미미합니다.
의의 및 주장 이 논문은 블록 기반 이중 디코더가 훈련과 추론 효율성 간의 역사적인 이분법을 성공적으로 해결했다고 주장합니다. 인코더 - 디코더와 유사한 구조에 완전한 손실 감독을 복원함으로써, 저자들은 디코더 전용 모델의훈련 효율성 및 확장 특성을 유지하면서 인코더 - 디코더의 추론 효율성(특히 감소된 KV 캐시 및 토큰당 지연 시간) 을 달성할 수 있음을 입증합니다.
저자들은 modest 한 훈련 시간 손실 차이 (약 0.2 나트) 는 상당한 추론 시간 절감을 가능하게 하는 아키텍처 분리를 위한 필수 비용이라고 강조합니다. 이는 아키텍처가 메모리 제약이 있는 엣지 장치 및 시나리오에 특히 적합하며, 사전 훈련 중 잠재적 체인 - 오브 - 씽킹 추론과 컨텍스트 디코더에 특화된 루프형 트랜스포머 적용과 같은 새로운 연구 방향을 가능하게 한다고 주장합니다.