LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
AI 모델의 새로운 가속화 기술 제시
메타(Meta)가 대규모 언어 모델(Large Language Model, LLM)의 추론 속도를 획기적으로 개선할 수 있는 ‘LayerSkip’ 기술을 공개했다. 이 기술은 기존 LLM 대비 최대 2.16배 빠른 추론 속도를 달성했으며, 정확도 저하를 최소화하면서도 컴퓨팅 자원을 효율적으로 활용할 수 있다는 장점이 있다.
현재 LLM은 높은 컴퓨팅 요구사항과 메모리 사용량으로 인해 GPU 서버 운영에 상당한 비용과 에너지가 소비되는 문제점을 안고 있다. 일반 GPU나 모바일, 엣지 디바이스에서 구동 가능한 가속화 솔루션들이 존재하지만, 대부분 정확도가 크게 감소하는 한계가 있었다.
LayerSkip의 핵심은 AI 모델의 계층(layer)을 선택적으로 건너뛰는 방식에 있다. 연구진은 모델 학습 시 초기 계층에서는 낮은 비율로, 후기 계층에서는 높은 비율로 레이어를 건너뛰도록 설계했다. 이를 통해 모든 계층을 순차적으로 처리하는 기존 방식 대비 처리 속도를 대폭 개선했다.
특히 기존의 모델 가속화 접근법들이 주로 가중치의 비제로(non-zero) 수를 줄이거나(희소성), 가중치당 비트 수를 줄이거나(양자화), 계층당 헤드 수를 줄이는(헤드 프루닝) 방식에 집중했던 것과 달리, LayerSkip은 각 토큰에 필요한 계층 수를 줄이는 방식을 채택했다. 이는 특별한 하드웨어나 소프트웨어 커널 없이도 구현 가능하다는 장점이 있다.
자체 추측 디코딩으로 정확도 유지
LayerSkip의 또 다른 혁신은 ‘자체 추측 디코딩(Self-Speculative Decoding)’ 기술의 도입이다. 이는 초기 계층에서 빠르게 결과를 예측한 후, 남은 계층들을 통해 결과를 검증하고 보정하는 방식이다. 기존의 추측 디코딩 방식들이 주 모델과 별도의 드래프트 모델을 필요로 했던 것과 달리, LayerSkip은 단일 모델 내에서 이를 구현했다.
이러한 접근은 메모리 사용량을 줄이고 구현 복잡성을 낮추는 동시에, 드래프트 단계와 검증 단계에서 계산과 활성화를 공유할 수 있게 한다. 특히 키-값(Key-Value) 캐시를 효율적으로 관리하여 추가적인 메모리 부담을 최소화했다.
연구진은 다양한 규모의 라마(Llama) 모델을 대상으로 광범위한 실험을 진행했다. 실험은 처음부터의 사전학습, 지속학습, 특정 도메인 파인튜닝, 특정 작업 파인튜닝 등 다양한 학습 시나리오를 포함했다.
실험 결과, CNN/DM 문서 요약 작업에서 2.16배, 코딩 작업에서 1.82배, TOPv2 의미 분석 작업에서 2.0배의 속도 향상을 달성했다. 특히 요약 작업에서는 정확도 저하 없이 이러한 속도 향상을 달성했다는 점이 주목할 만하다.
향후 발전 방향과 한계점
연구진은 LayerSkip이 가진 몇 가지 한계점도 함께 공개했다. 우선 이 기술을 적용하기 위해서는 모델을 재학습하거나 파인튜닝해야 한다는 점이다. 또한 레이어 드롭아웃을 위한 최대 확률(p_max), 조기 종료를 위한 스케일(e_scale) 등의 하이퍼파라미터 조정이 필요하며, 이는 정확도 유지를 위해 세심한 튜닝이 요구된다.
처음부터 사전학습을 진행할 경우에는 학습률(learning rate)을 증가시켜야 하는데, 이를 위한 최적값을 찾는 과정이 복잡하고 시간이 많이 소요될 수 있다는 점도 한계로 지적됐다.
연구진은 향후 연구 방향으로 조기 종료 계층의 정확도를 더욱 향상시키는 것을 제시했다. 이를 통해 자체 추측 디코딩 과정에서 더 높은 속도 향상을 달성할 수 있을 것으로 기대된다. 또한 각 토큰별로 다른 종료 계층을 동적으로 결정하는 조건을 탐구하여 토큰 수용률을 개선하는 방안도 연구될 예정이다.
기사에 인용된 리포트 원문은 링크에서 확인할 수 있다.
기사는 클로드 3.5 Sonnet과 챗GPT-4o를 활용해 작성되었습니다.
관련 콘텐츠 더보기
