LayerSkip : 더 빠른 추론을 위한 레이어 건너뛰기와 자가 추론 방식

 

1.서론

1.2.기존 LLM의 문제점

 

ChatGPT와 같은 기존 대형 언어 모델(LLM) 추론은 텍스트 생성 등의 작업에서 모든 레이어를 거치면서 순차적 계산을 합니다. 당연히 처리 속도가 느리고 많은 메모리가 필요하게 되며 일반 컴퓨터에서는 효율적이지 않습니다. 특히 모바일 기기나 엣지 디바이스 등에서 실행하기에는 무리가 있죠. 기존의 속도를 높이기 위한 방법으로 다른 방식들도 초안 기본 모델과 검증 모델 두개를 사용해 예측하기 위해 복잡한 구현물과 많은 메모리 사용이 역시 문제였습니다. 

 

1.2.LayerSkip

LayerSkip은 Meta에서 공개한, 대형 언어 모델의 추론 속도를 높이기 위해 레이어를 건너뛰는 기법과 자가 추론 방식을 결합하여 연산 비용을 줄이면서도 정확도를 유지할 수 있습니다. 그리고, LayerSkip은 추론에만 초점을 맞춘 다른 기술과는 달리 학습 및 추론최적화를 모두 포함하는 포괄적 접근방식을 택하고 있습니다. 그리고, 추론가속을  위해 특수한 하드웨어나 소프트웨어 커널이 필요치 않습니다. 

 

2.LayerSkip에 적용된 기술

2.1. 핵심 기술 요소

그림: LayerSkip 개요

• 훈련 중 Layer Dropout과 Early Exit Loss: 초기 레이어에는 낮은 드롭아웃, 후반부에는 높은 드롭아웃을 적용하여 훈련
• 조기 종료 추론: 각 토큰 예측 시 초기 레이어만 사용하여 빠르게 추론 가능
• 자가 추론 디코딩: 초안을 작성한 후 남은 레이어를 사용하여 예측 검증을 수행. 이 과정에서 초안 단계와 검증 단계의 캐시(KV 캐시)를 재사용하여 속도를 높임

 

 

2.2. 주요 구성요소

 

LayerSkip의 주요 구성 요소는 다음과 같습니다.

훈련 전략

• 레이어 드롭아웃 : 훈련 중 일부 레이어를 무작위로 건너뛰는 방식으로, 특히 초기 레이어에 드롭아웃을 낮추고 후반부 레이어에 높은 드롭아웃을 적용하여 모델이 더 빨리 예측을 수행하도록 학습합니다.
• 조기 종료 손실 : 모든 레이어가 동일한 언어 모델(로지스틱) 헤드로 연결되도록 학습하여, 중간 레이어에서 정확히 예측할 수 있도록 만듭니다.

추론 최적화

• 자가 추론 디코딩(Self-Speculative Decoding) : 추론 중 초기 레이어에서 초안을 생성하고 후반부 레이어를 이용해 이를 검증 및 수정하는 방식으로 메모리 사용량을 줄이고, 속도를 높입니다.

 

3.LayerSkip의 장점

• 효율성 : 전체 레이어를 거치지 않고 일부 레이어에서 예측을 조기 종료하여 빠른 추론을 실현합니다.
• 유연성 : 동일한 모델 내에 서로 다른 크기의 하위 모델을 생성하여 계산 요구사항에 따라 적용 가능한 추론을 허용합니다.
• 메모리 최적화 : 자가 추론 디코딩은 다른 추측 디코딩 방식에 비해 메모리 사용량이 더 적습니다. 
• 공유 컴퓨팅 : 초안 및 검증 단계 간의 공유 컴퓨팅 및 활성화의 이점이 있습니다. 
• 오픈소스 : 코드가 공개되어 있습니다. CC4.0

 

 

---

참고정보

 

1.논문

https://arxiv.org/pdf/2404.16710

 

 

2.코드(깃헙)

GitHub - facebookresearch/LayerSkip: Code for "LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding", ACL 2024