대규모 언어 모델에서 '무한 입력'을 가능하게 해주는 기법 「StreamingLLM」이 등장

 
 
계산 비용과 성능을 유지한 채 무한한 입력을 처리할 수 있는 대규모 언어 모델 기법 「StreamingLLM」의 논문이 2023년 9월 29일에 공개되었습니다. 입력이 길어짐에 따라 첫 번째 정보는 손실되기 때문에 책 요약 등의 작업에는 적합하지 않지만 대화가 길어도 성능을 유지하면서 원활하게 답변할 수 있다고 합니다.

mit-han-lab/streaming-llm: Efficient Streaming Language Models with Attention Sinks
https://github.com/mit-han-lab/streaming-llm

GitHub - mit-han-lab/streaming-llm: Efficient Streaming Language Models with Attention Sinks

[2309.17453] Efficient Streaming Language Models with Attention Sinks
https://arxiv. org/abs/2309.17453

모델을 로드할 때 둘 다 같은 동작을 합니다.

대량의 질문을 준비해, 「질문 → AI에 의한 회답 → 질문 → AI에 의한 회답 → …」을 반복해 갑니다.

잠시 후 StreamingLLM을 사용하지 않는 모델은 응답이 지리멸렬해졌습니다만, StreamingLLM을 사용한 모델은 성공적으로 응답할 수 있었습니다.

그 후, StreamingLLM을 사용하지 않는 모델은 메모리를 다 사용해 정지해 버렸습니다만 StreamingLLM를 사용한 모델은 계속 움직입니다.

기존 방법과의 비교 이미지는 아래 그림과 같습니다. 4개의 기법의 동작 이미지와 함께 알고리즘의 시간 계산량을 O-기법으로 표시한 것, 예측 정확성을 나타내는 PPL 값 등이 표시되어 있습니다. L은 사전 학습 텍스트의 길이이고 T는 몇 번째 토큰을 예측할 것인지를 나타냅니다.

(a)의 Dense Attention은 시간 계산량이 T의 제곱 시간이 되는 것 외에 메모리를 대량으로 소비합니다. 또한 텍스트 길이가 사전 학습 텍스트 길이를 초과하면 성능이 크게 저하됩니다.

(b)의 Window Attention은 가까운 L 토큰 분의 캐시만을 유지하는 구조로, 시간 계산량은 억제되었지만 최초의 토큰이 캐시 범위로부터 벗어나면 단번에 퍼포먼스가 저하되어 버린다는 것.

(c)의 Sliding Window w/ Re-computation은 토큰마다 가까운 L 토큰분의 캐시를 갱신하는 구조로, 긴 문장에서도 예측 정확성을 유지할 수 있지만 계산량은 증가합니다.

(d)의 StreamingLLM이 이번에 개발된 구조입니다. 논문 저자들은 (b)의 Window Attention에서 첫 번째 토큰 캐시를 유지하면 성능이 회복되는 현상을 발견하고 "attention sink"라는 이름을 붙여 어텐션 계산에 통합했습니다. Window Attention의 시간 계산량을 유지한 채 긴 문장으로 좋은 퍼포먼스를 발휘할 수 있다는 것.

StreamingLLM을 통합하면 Llama-2, MPT, Falcon, Pythia와 같은 대규모 언어 모델에서 400만 토큰 이상의 길이의 문장을 처리할 수 있습니다. 또, 사전 학습 시에 attention sink로서 플레이스홀더 토큰을 준비하면 새로운 속도 향상을 확인할 수 있었다고 합니다.

덧붙여 장대한 문장 입력이 가능해진다고는 해도, 콘텍스트 길이의 제한은 변경되지 않습니다. 예를 들어, Lama-2를 4096 토큰의 콘텍스트 길이로 사전 학습한 경우, StreamingLLM을 사용한 경우에도 최대 캐시 크기는 4096 토큰으로 남아 있으며, 부족한 부분은 중간 토큰을 파기함으로써 대응합니다.

StreamingLLM은 장기 기억을 강화하는 구조가 아니기 때문에, 책과 같은 긴 텍스트를 요약하는 태스크의 경우, 마지막 부분만을 요약할 가능성이 높다는 것. 한편, 매일의 어시스턴트 등, 과거의 데이터 없이 모델을 계속적으로 동작시키고 싶은 경우에는 적절한 기법입니다. 대화 길이가 콘텍스트 길이를 초과하는 경우에도 계속 작동할 수 있습니다.

실제로 대규모 언어 모델에 StreamingLLM을 포함한 코드는 GitHub에 공개되어 있습니다.
https://github.com/mit-han-lab/streaming-llm

GitHub - mit-han-lab/streaming-llm: Efficient Streaming Language Models with Attention Sinks