[논문 리뷰] MEGABYTE: Predicting Million-byte Sequences with Multiscale Transformers

Lili Yu, Dániel Simig외 4명|arXiv (Cornell University)|2023. 05. 12.

Speech Recognition and Synthesis인용 수 15

한 줄 요약

MegaByte는 전역 패치-수준 트랜스포머와 패치별 로컬 자기회귀 모델링을 결합하여 긴 바이트 시퀀스를 효율적으로 모델링하는 두 단계의 패치 기반 디코더 아키텍처를 제시하고, 대규모에서 토큰화 없는 시퀀스 모델링을 가능하게 한다.

ABSTRACT

Autoregressive transformers are spectacular models for short sequences but scale poorly to long sequences such as high-resolution images, podcasts, code, or books. We proposed Megabyte, a multi-scale decoder architecture that enables end-to-end differentiable modeling of sequences of over one million bytes. Megabyte segments sequences into patches and uses a local submodel within patches and a global model between patches. This enables sub-quadratic self-attention, much larger feedforward layers for the same compute, and improved parallelism during decoding -- unlocking better performance at reduced cost for both training and generation. Extensive experiments show that Megabyte allows byte-level models to perform competitively with subword models on long context language modeling, achieve state-of-the-art density estimation on ImageNet, and model audio from raw files. Together, these results establish the viability of tokenization-free autoregressive sequence modeling at scale.

연구 동기 및 목표

• 매우 긴 시퀀스(바이트)에 대한 토큰화 없는 시퀀스 모델링의 필요성을 제시하고 대형 디코더의 비효율성을 규명한다.
• 전역 컨텍스트를 가진 패치로 시퀀스를 분할하고 패치 내 로컬 모델을 갖는 다중 스케일 디코더 아키텍처를 제안한다.
• 해당 아키텍처가 서브-제곱 어텐션, 더 큰 유효 피드포워드 용량, 그리고 더 빠른 생성 속도를 제공함을 보인다.
• 경쟁력 있는 언어 모델링, ImageNet에서의 최첨단 밀도 추정, 그리고 원시 오디오 모델링을 시연한다.
• 텍스트, 이미지, 오디오 모달리티에서 계산 및 데이터 제어 하에 성능을 평가한다.

제안 방법

• 입력 시퀀스를 길이 P의 고정 크기 패치로 패치화하여 전체 길이 T에서 K개의 패치를 형성한다.
• 패치 임베더를 사용해 바이트를 위치 정보를 갖는 패치 표현으로 변환한다.
• 패치 간 컨텍스트를 형성하기 위해 대형 글로벌 트랜스포머를 사용하여 패치 표현을 순서적으로 처리한다(인과적, 패치-수준 자기-주의).
• 각 패치 내에서 더 작은 로컬 트랜스포머를 적용해 글로벌 출력과 패치 내 바이트 임베딩을 이용해 해당 패치 내부의 바이트를 자기회귀적으로 예측한다.
• 로컬 모델의 로컬 임베딩 표현 위에 소프트맥스(ssoftmax)로 토큰 확률을 계산한다.
• 효율성과 맥락 활용을 개선하기 위한 확장으로 합성곱 패치 인코딩, 패치 간 교차 어텐션, 스트라이드 추론 등의 확장을 선택적으로 보강한다.

실험 결과

연구 질문

• RQ1 토큰화 없는 자동회귀 모델이 백만 바이트를 넘는 시퀀스를 효율적으로 처리할 수 있는가?
• RQ2 두 수준(전역 패치-수준 및 로컬 패치 내) 트랜스포머 아키텍처가 표준 디코더 및 기존의 긴 시퀀스 모델과 비교하여 경쟁력 있는 성능과 향상된 효율성을 달성하는가?
• RQ3 텍스트, 이미지, 오디오 작업에서 MegaByte를 사용할 때 자기-어텐션 복잡도, 토큰당 계산량, 생성 속도에서 얻는 이점은 무엇인가?
• RQ4 고정된 컴퓨트 예산 하에서 패치 크기와 글로벌/로컬 모델 용량 배정이 성능에 어떤 영향을 미치는가?

주요 결과

Dataset	Transformer (bpb)	PerceiverAR (bpb)	MegaByte (bpb)
PG-19	1.057	1.104	1.000
Stories	1.064	1.070	0.978
Books	1.097	1.104	1.007
arXiv	0.816	0.791	0.678
Code	0.575	0.546	0.411

• MegaByte는 적절한 패치 크기에서 O(T^{4/3}) 복잡도의 서브-제곱 어텐션을 달성하여 백만 바이트 시퀀스 모델링을 가능하게 한다.
• 패치당 로컬 자기회귀는 패치당 큰 피드포워드 용량을 가능하게 하면서 전체 계산량을 관리해 모델 표현력을 향상시킨다.
• 계산/데이터 제어 실험에서 MegaByte는 긴 컨텍스트 언어 모델링(PG-19, Stories, Books, arXiv, Code)에서 표준 디코더 트랜스포머 및 PerceiverAR를 능가한다.
• ImageNet 밀도 추정(64x64, 128x128, 640x640 변형)에서 MegaByte는 최첨단과 일치하거나 이를 능가하며, 가장 강력한 기준선 대비 약 절반의 계산을 사용한다.
• 원시 바이트의 오디오 모델링에서 MegaByte는 바이트 수준 기준선보다 더 낮은 bpb(bits per byte)를 달성하여 토큰화 없는 자기회귀 오디오 모델링을 효과적으로 보여준다.
• 스트라이드 추론과 패치 간 확장은 추가로 성능을 향상시키며 긴 컨텍스트 활용과 생성 속도를 개선한다.