Hugging Face 저지연 텍스트 생성을 위한 Assisted Generation

TL;DR

Hugging Face는 Assisted Generation이라는 새로운 디코딩 방식을 도입했으며, 이는 자동 회귀 텍스트 생성의 지연 시간을 줄이기 위해 설계되었습니다. 작은, 빠른 어시스턴트 모델이 후보 토큰을 제안하고, 큰 모델이 이를 단일 전방 패스로 검증함으로써, 특히 메모리 오프로드를 사용할 때 일반 하드웨어에서 지연 시간을 최대 10배까지 감소시킬 수 있습니다.

텍스트 생성 지연 시간의 병목 현상

텍스트 생성 지연 시간은 주로 연산 문제라기보다 메모리 대역폭 문제입니다. 표준 자동 회귀 전방 패스에서는 모델 레이어 가중치를 GPU RAM에서 GPU 연산 코어로 로드할 때 병목이 발생합니다. 큰 모델은 전체 응답을 생성하기 위해 수백 번의 순차 전방 패스가 필요하므로, 이 메모리 전송 오버헤드가 주요 비용이 됩니다.

Flash Attention, INT8 quantization, 배치 처리(throughput 증가), Tensor Parallelism(메모리 대역폭 분산)과 같은 기존 최적화 기법도 존재하지만, 이들은 종종 높은 비용이나 지연 시간 측면에서의 트레이드오프를 동반합니다. Assisted Generation은 기본 모델이 수행해야 하는 전방 패스 횟수를 줄임으로써 이를 해결합니다.

Assisted Generation 작동 방식

Assisted Generation은 언어 디코더의 특정 특성을 활용합니다: 단일 전방 패스를 사용해 다음 토큰을 예측하는 것이 아니라 토큰 시퀀스를 검증할 수 있다는 점입니다. 캐시 없이 토큰 시퀀스를 모델에 전달하면, 모델은 해당 시퀀스의 모든 위치에 대한 logits를 반환합니다.

Assisted Generation 루프

이 과정은 작은 어시스턴트 모델과 큰 기본 모델 사이의 순환으로 진행됩니다:

  1. Candidate Generation: 작은 어시스턴트 모델이 greedy decoding을 사용해 짧은 후보 토큰 시퀀스를 생성합니다(초기에는 5 토큰 윈도우).
  2. Validation: 기본 모델이 이 후보들에 대해 단일 전방 패스를 수행해 각 위치의 logits를 얻습니다.
  3. Comparison: 기본 모델이 예측한 토큰을 어시스턴트의 후보와 좌측부터 순차적으로 비교합니다.
  4. Correction: 첫 번째 불일치가 어시스턴트가 실패한 지점을 나타냅니다. 해당 불일치 지점에서 기본 모델의 예측 토큰을 유지하고, 이후의 어시스턴트 후보는 모두 버립니다.
  5. Adjustment: 휴리스틱을 통해 다음 반복에 요청할 후보 수를 조정합니다: 모든 토큰이 일치하면 윈도우를 2씩 늘리고, 불일치가 발생하면 1씩 감소시킵니다.

어시스턴트 모델에 대한 요구 사항

효율성을 유지하려면 어시스턴트 모델은 두 가지 기준을 만족해야 합니다:

  • Shared Tokenizer: 어시스턴트는 기본 모델과 동일한 토크나이저를 사용해야 하며, 이는 비용이 많이 드는 CPU 기반 디코딩 및 재인코딩 단계를 피하기 위함입니다.
  • Size Differential: 어시스턴트는 기본 모델보다 최소 한 단계(10배) 이상 작아야 하며, 그래야 어시스턴트의 생성 시간이 기본 모델의 전방 패스에 비해 무시할 수 있을 정도가 됩니다.

성능 및 활용 가능성

Assisted Generation은 하드웨어와 모델 구성에 따라 다양한 수준의 속도 향상을 제공합니다:

  • Memory Offloading: 모델이 GPU 메모리에 들어가지 않아 오프로드를 사용할 때 가장 큰 이득을 보이며, 속도 향상이 10배까지 도달합니다.
  • GPU-Resident Models: 모델이 GPU 메모리에 적재된 경우, 속도 향상은 2배(표준)에서 3배(INT8 quantization 적용) 정도입니다.
  • Task Suitability: 입력에 기반한 작업, 예를 들어 자동 음성 인식(ASR), 번역, 요약 등에 가장 효과적입니다.

샘플링과의 호환성

Assisted Generation은 greedy decoding을 위해 설계되었지만, multinomial sampling과도 함께 사용할 수 있습니다. 다만 어시스턴트의 효과는 temperature 설정에 크게 좌우됩니다. 낮은 temperature(0에 가깝게)에서는 greedy decoding에 가까워 대부분의 지연 시간 이점을 유지하지만, 높은 temperature에서는 무작위성이 증가해 어시스턴트가 실패할 확률이 높아지고 속도 향상이 감소합니다.

향후 방향 및 관련 연구

Hugging Face는 텍스트 생성의 미래가 토큰당 고정된 연산 비용에서 벗어날 수 있다고 전망합니다. 대신 텍스트 복잡도에 따라 서로 다른 크기의 모델이 시퀀스의 다양한 부분을 생성하도록 설계될 수 있습니다.

이 접근 방식은 Blockwise Parallel Decoding(Google Brain)과 Speculative Sampling(DeepMind)과 같은 연구와 핵심 원리를 공유합니다. 두 방법 모두 전방 패스를 활용해 더 긴 연속을 검증합니다.

Sources