Medusa (인공지능)

Medusa(영어: Medusa)는 대형 언어모델(LLM) 기반 생성 모델의 추론 속도를 가속화하기 위해 제안된 프레임워크이다.

개요

메두사 작동 개요

Medusa는 기존 언어모델이 토큰을 순차적으로 하나씩 생성하는 방식의 병목을 해결하기 위해 고안된 방식이다. 기존 방식에서는 출력할 토큰 K개에 대해 모델이 K번의 연산을 실행해야 하지만, Medusa는 여러 개의 디코딩 헤드를 추가해 후속 토큰을 병렬로 예측한다.

작동 원리

추가 디코딩 헤드

Medusa는 기존 모델의 출력층 위에 여러 개의 디코딩 헤드를 추가하고, 각 헤드는 다음 토큰뿐 아니라 뒤이어 생성될 여러 토큰을 예측하도록 설계된다. Medusa-1은 백본 모델을 고정한 채 헤드만 학습하고, Medusa-2는 모델 전체를 학습한다.

트리 기반 어텐션 및 수락 스킴

여러 후보 시퀀스는 트리 구조로 구성되며, 병렬 검증을 위한 어텐션 마스크가 사용된다. 또한 가능성이 높은 후보(prefix)를 수락하는 typical acceptance 방법을 통해 품질 저하 없이 속도 향상을 달성한다.

특징 및 장점

• 별도의 드래프트 모델이 필요하지 않다.
• 인프라 변경이 비교적 작다.
• 다양한 모델 규모에서 속도 향상을 입증했다.

제약 및 고려사항

• 여러 헤드가 추가되므로 메모리 사용량이 증가한다.
• 주로 배치 사이즈 1 환경에서 효과가 크며, 대량 처리 환경에서는 추가 연구가 필요하다.
• Medusa-2는 모델 전체를 학습하므로 학습 비용이 크다.

적용 사례 및 활용 가능성

• 실시간 대화형 어플리케이션 등 지연이 중요한 환경에 적합하다.
• 코드 생성, SQL 질의 생성 등 규칙성이 높은 출력에서 속도 향상 효과가 크다.
• 기존 LLM 시스템에 적은 구조 변경으로 적용 가능하다.

역사

• 2024년 1월 19일, Medusa 논문이 arXiv에 공개되었다.
• Amazon SageMaker 블로그에서 Medusa-1 적용 사례가 공유되었다.
• Medusa 오픈소스 저장소가 GitHub에 공개되었다.

관련 기술 및 비교

• Speculative Decoding: 별도 모델을 사용한 후보 생성 방식.
• Medusa는 해당 방식의 복잡성을 줄이는 대안으로 평가된다.
• 트리 기반 병렬 디코딩과도 밀접한 연관이 있다.

Medusa와 추측 디코딩의 차이

Medusa와 기존 추측 디코딩(speculative decoding)은 모두 토큰 생성 속도를 높이기 위한 기법이지만, 설계 철학과 구조적 접근 방식에서 명확한 차이를 가진다.

구조적 차이

• 기존 추측 디코딩
  • 별도의 드래프트 모델이 필요하다.
  • 드래프트 모델이 후보 토큰을 생성하고, 메인 모델이 이를 검증한다.
  • 두 모델 간 상호작용을 기반으로 한 이원 구조를 가진다.*Medusa
  • 메인 모델 내부에 여러 디코딩 헤드를 추가한다.
  • 하나의 모델이 병렬 예측과 검증을 동시에 수행한다.
  • 트리 기반 검증 구조를 사용하며 별도 모델이 필요 없다.

결과적 차이

• 기존 방식은 두 모델을 관리해야 하므로 구현 복잡성과 인프라 요구량이 증가한다.
• Medusa는 단일 모델 확장 방식이므로 구조가 단순하고 적용 비용이 낮다.
• 두 방식 모두 속도 향상을 목표로 하지만 접근 방식이 크게 다르기 때문에 Medusa는 전통적 speculative decoding의 변형이자 대안으로 평가된다.

같이 보기

• 대형 언어 모델
• 추측 디코딩
• Transformer (딥러닝)
• 추론 최적화

참고 문헌

• Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (arXiv)
• Achieve ~2x speed-up in LLM inference with Medusa-1 (AWS Blog)