압축된 디코딩 1: 플로우 매칭을 이용한 이미지 자기 회귀 모델의 일단계 샘플링
Distilled Decoding 1: One-step Sampling of Image Auto-regressive Models with Flow Matching
December 22, 2024
저자: Enshu Liu, Xuefei Ning, Yu Wang, Zinan Lin
cs.AI
초록
자기회귀(AR) 모델은 텍스트 및 이미지 생성에서 최첨단 성능을 달성했지만 토큰 단위 처리로 인해 생성 속도가 느립니다. 우리는 야심찬 질문을 제기합니다: 사전 훈련된 AR 모델을 적응하여 단 한 번 또는 두 번의 단계로 출력물을 생성할 수 있을까요? 성공할 경우, 이는 AR 모델의 개발 및 배포를 크게 발전시킬 것입니다. 우리는 기존의 작업들이 한 번에 여러 토큰을 생성하여 AR 생성 속도를 높이려는 시도가 출력물 간의 조건부 의존성으로 인해 출력 분포를 기본적으로 포착할 수 없다는 점을 알아차렸습니다. 이는 몇 단계 생성에 대한 효과를 제한합니다. 이를 해결하기 위해 우리는 Distilled Decoding (DD)을 제안합니다. 이는 흐름 일치를 사용하여 사전 훈련된 AR 모델의 출력 분포로부터 가우시안 분포로의 결정적 매핑을 생성합니다. 그런 다음 이 매핑을 증류시키는 네트워크를 훈련하여 몇 단계 생성을 가능하게 합니다. DD는 원래 AR 모델의 훈련 데이터가 필요하지 않아 더 실용적입니다. 우리는 ImageNet-256에서 최첨단 이미지 AR 모델에 DD를 평가하고 유망한 결과를 제시합니다. VAR의 경우, 10단계 생성이 필요한데 DD는 1단계 생성을 가능하게 합니다(6.3배 속도 향상), FID가 4.19에서 9.96으로 상당한 증가가 있습니다. LlamaGen의 경우, DD는 256단계 생성을 1단계로 줄여 217.8배 속도 향상을 달성하며 FID가 4.11에서 11.35로 비슷한 증가가 있습니다. 두 경우 모두 기준선 방법은 FID가 100을 초과하여 완전히 실패합니다. DD는 텍스트에서 이미지 생성에도 뛰어나며, LlamaGen의 경우 256단계 생성을 2단계로 줄여 FID가 25.70에서 28.95로 미미한 증가가 있습니다. 이미지 AR 모델의 단일 단계 생성 가능성을 시연한 최초의 작업으로, DD는 AR 모델이 본질적으로 느린 것으로 여겨지는 관념에 도전하며 효율적인 AR 생성을 위한 새로운 기회를 엽니다. 프로젝트 웹사이트는 https://imagination-research.github.io/distilled-decoding에서 확인할 수 있습니다.
English
Autoregressive (AR) models have achieved state-of-the-art performance in text
and image generation but suffer from slow generation due to the token-by-token
process. We ask an ambitious question: can a pre-trained AR model be adapted to
generate outputs in just one or two steps? If successful, this would
significantly advance the development and deployment of AR models. We notice
that existing works that try to speed up AR generation by generating multiple
tokens at once fundamentally cannot capture the output distribution due to the
conditional dependencies between tokens, limiting their effectiveness for
few-step generation. To address this, we propose Distilled Decoding (DD), which
uses flow matching to create a deterministic mapping from Gaussian distribution
to the output distribution of the pre-trained AR model. We then train a network
to distill this mapping, enabling few-step generation. DD doesn't need the
training data of the original AR model, making it more practical.We evaluate DD
on state-of-the-art image AR models and present promising results on
ImageNet-256. For VAR, which requires 10-step generation, DD enables one-step
generation (6.3times speed-up), with an acceptable increase in FID from 4.19
to 9.96. For LlamaGen, DD reduces generation from 256 steps to 1, achieving an
217.8times speed-up with a comparable FID increase from 4.11 to 11.35. In
both cases, baseline methods completely fail with FID>100. DD also excels on
text-to-image generation, reducing the generation from 256 steps to 2 for
LlamaGen with minimal FID increase from 25.70 to 28.95. As the first work to
demonstrate the possibility of one-step generation for image AR models, DD
challenges the prevailing notion that AR models are inherently slow, and opens
up new opportunities for efficient AR generation. The project website is at
https://imagination-research.github.io/distilled-decoding.