소식

Apple은 대형 모델이 게으른 방법을 학습할 수 있도록 합니다. 첫 번째 토큰을 더 빠르게 생성하고 정확성을 유지합니다.

2024-08-02

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina



기계 심장 보고서

머신하트 편집부

게으르면 일을 더 잘하게 됩니다.

Llama 3.1이 방금 출시되었습니다. 아직 사용해 보셨나요? PC가 최신 최고 사양이더라도 가장 작은 8B 버전을 실행하더라도 상당한 지연이 발생할 수 있습니다. 모델의 추론 효율성을 향상시키기 위해 연구자들은 다양한 방법을 고안해 냈지만 그 중 다수는 모델의 정확도를 어느 정도 희생하게 만들었습니다.

최근 Apple과 Meta AI의 연구팀은 Llama 2의 사전 채우기 단계의 추론 속도를 2배 이상으로 높이면서 정확도가 크게 떨어지지 않는 새로운 방법을 제안했습니다. 3.1의 속도 향상은 어느 정도 영감을 줍니다. 그들은 이 접근 방식을 Lazy LLM(Lazy Large Language Model)이라고 부릅니다.



논문 제목: LazyLLM: 효율적인 장기 컨텍스트 LLM 추론을 위한 동적 토큰 정리

논문 주소: https://arxiv.org/abs/2407.14057

그렇다면 어떻게 LLM을 게으르게 만들 수 있을까요? 그들의 방법을 이해하려면 먼저 표준 프롬프트 기반 LLM 추론 프로세스가 무엇인지 알아야 합니다. 간단히 말해서 프로세스는 그림 1과 같이 사전 채우기와 디코딩의 두 단계로 나뉩니다.



사전 채우기 단계에서 모델은 프롬프트에서 각 토큰의 KV 캐시를 계산 및 저장하고 첫 번째 토큰을 예측합니다. 우리는 채우기 전 단계에서 소요된 시간을 "첫 번째 토큰까지의 시간(TTFT)"이라고 부릅니다.

사전 채우기 단계 다음에는 디코딩 단계가 이어집니다. 이 단계에서 모델은 중지 기준이 충족될 때까지 캐시된 KV를 다시 사용하여 다음 토큰을 반복적으로 디코딩합니다.

사전 채우기 단계에서는 모든 Transformer 레이어가 프롬프트의 모든 토큰을 사용합니다. 현재 최고의 Transformer 기반 LLM은 깊고 넓기 때문에 프롬프트가 길면 TTFT가 느려질 수 있으며 주의 계산 비용은 프롬프트의 토큰 수에 따라 2차적으로 증가합니다. 예를 들어 Llama 2(버전 7B)는 32개의 Transformer 레이어를 쌓고 모델 치수는 4096입니다. 이 경우 TTFT에는 각 후속 디코딩 단계의 벽 시간이 21배 필요하며 이는 LongBench 벤치마크에서 총 생성 시간의 약 23%를 차지합니다.

따라서 LLM 추론을 효율적으로 수행하려면 TTFT를 최적화하는 것이 매우 중요한 단계입니다.

LLM 추론 최적화는 활발한 연구 분야이지만 많은 방법은 디코딩 단계의 추론 속도를 향상시키는 데 중점을 둡니다. 연구자들은 TTFT의 개선에 거의 관심을 기울이지 않았습니다. 일부 압축 기반 연구 결과는 LLM의 크기를 줄여 TTFT를 암시적으로 향상시킬 수 있습니다.

또 다른 연구 방향은 정적 Transformer 아키텍처에서 TTFT를 개선하는 것입니다. 이 연구 방향에 대해 자연스럽게 다음과 같은 질문이 제기됩니다. 첫 번째 토큰을 생성할 때 모든 프롬프트 토큰이 필수입니까?

그림 2는 LongBench 벤치마크의 LLM 분석 결과를 보여줍니다.



첫 번째 생성된 토큰의 경우 입력 토큰의 어텐션 점수가 매우 희박하다는 것을 알 수 있습니다. 이는 입력 프롬프트의 많은 토큰이 중복되고 제거되더라도 다음 예측에 영향을 미치지 않음을 나타냅니다. 토큰. 이 관찰은 팀이 제안한 LazyLLM의 기초입니다.

LazyLLM의 장점은 적용 범위가 넓고, 교육이 필요 없으며, 좋은 결과를 얻을 수 있다는 점입니다. 그림 3은 표준 LLM과 LazyLLM을 비교합니다.



레이지LLM

그림 4는 LazyLLM의 전체 프레임워크를 보여줍니다.



전체 컨텍스트에서 시작하여 LazyLLM은 점차적으로 토큰을 정리하여 최종 모델을 얻는 데 사용되는 계산 수를 점차적으로 줄입니다. LazyLLM을 사용하면 이전 단계에서 토큰 중 일부가 제거된 경우에도 모델이 다양한 생성 단계에서 다양한 토큰 하위 집합을 선택할 수 있습니다. 정적 가지치기(모든 토큰이 동시에 정리됨)와 비교하여 동적 가지치기는 각 생성 단계에서 다음 토큰 예측을 최적화하여 모델 성능을 유지하는 데 도움이 됩니다.

점진적인 토큰 가지치기

일부 이전 연구에서는 LLM 추론을 최적화하기 위해 토큰 프루닝을 성공적으로 사용했습니다. 그러나 이러한 방법은 정리가 시작되기 전에 프롬프트 토큰의 중요성을 분석하기 위해 예측된 처음 몇 개의 토큰에 대한 완전한 주의 맵을 축적해야 합니다. 따라서 사전 채우기 단계에서 여전히 모든 KV 캐시를 계산해야 하기 때문에 TTFT를 줄이는 데 적합하지 않습니다.

이에 비해 LazyLLM은 "매우 게으르다"고 첫 번째 추론 반복(사전 채우기 단계)부터 시작하여 다음 토큰을 예측하는 데 중요한 토큰만 계산합니다.

첫 번째 반복에서 주요 과제는 각 토큰의 중요성을 결정하는 것이었습니다. 토큰 숨겨진 상태가 Transformer 레이어를 통과하면서 진화한다는 것을 보여주는 이전 연구에서 영감을 받아 팀의 솔루션은 각 생성 단계에서 레이어별 토큰 가지치기를 사용하는 것입니다. 구체적으로, 각 레이어의 어텐션 맵을 사용하여 예측할 토큰에 대한 입력 토큰의 중요성을 결정합니다.

토큰의 신뢰도 점수를 계산한 후 또 다른 어려운 문제는 토큰을 잘라내기 위한 임계값을 결정하는 것입니다.

특히, 다양한 레이어와 다양한 작업의 경우 주의 점수가 변경됨에 따라 이 임계값도 변경될 수 있습니다. 팀의 솔루션은 상위 k 백분위수 선택 전략을 사용하는 것입니다. 특히 토큰의 신뢰도 점수가 입력 토큰의 k번째 백분위수보다 작은 경우 해당 토큰은 제거됩니다. 토큰이 정리되면 더 이상 모든 후속 레이어 계산에 참여하지 않습니다.

즉, 후속 레이어에서 사용되는 토큰은 이전 레이어에서 사용되는 토큰의 하위 집합입니다.

이후의 실험에서는 가지치기 레이어의 위치와 가지치기된 토큰의 수가 다르면 성능도 변경되는 것으로 나타났습니다. 구체적으로, 동일한 Transformer 레이어에 대해 가지치기를 통해 제거되는 토큰이 많아질수록 모델의 성능은 점차 저하됩니다.

또한 초기 레이어의 가지치기에 비해 이후 레이어에서 가지치기를 수행할 때 더 나은 성능이 달성되었으며 이는 이후 레이어가 토큰 가지치기에 덜 민감하다는 것을 나타냅니다. 속도와 정확성의 균형을 높이기 위해 팀은 그림 4와 같이 점진적인 가지치기를 사용하여 초기 레이어에 더 많은 토큰을 유지한 다음 이후 레이어로 흐르면서 토큰 수를 점차 줄였습니다.

보조 캐시(보조 캐시)

사전 채우기 단계에는 KV 캐시가 없으며 각 토큰은 숨겨진 상태로 표시됩니다. 따라서, 정리된 토큰의 숨겨진 상태를 제거함으로써 점진적인 토큰 정리를 달성할 수 있습니다. 그러나 점진적인 토큰 가지치기를 후속 디코딩 단계로 확장하는 것은 간단하지 않습니다. 그 이유는 각 디코딩 단계에서는 Pre-fill 단계에서 계산된 KV 버퍼를 사용하여 어텐션을 계산하기 때문입니다. LazyLLM은 채우기 전 단계에서 점진적인 토큰 정리를 수행하므로 특정 레이어에서 정리된 토큰의 KV는 다음 레이어의 KV 캐시에 나타나지 않습니다.

참고로, LazyLLM 프레임워크에서는 각 생성 단계에서 이전 단계에서 제거되었는지 여부에 관계없이 각 단계의 전체 입력 토큰 시퀀스에서 서로 다른 토큰 하위 집합을 선택할 수 있습니다. 예를 들어, 후속 디코딩 단계에서는 KV 캐시에 존재하지 않는 프루닝된 토큰을 재선택하여 어텐션 계산을 수행할 수 있다. 이 경우 모델은 해당 토큰에 대한 KV 캐시를 검색할 수 없습니다.

이에 대한 직관적인 해결책은 Transformer의 원본을 통해 토큰을 전달하는 것입니다. 그러나 이로 인해 동일한 토큰이 이중으로 계산되고 궁극적으로 전체 생성 속도가 느려집니다.

이 문제를 해결하기 위해 팀은 원래 KV 캐시 외에 또 다른 캐시인 Aux 캐시(보조 캐시)를 도입했습니다.

정리된 토큰(예: 그림 4의 T4 및 T7)이 후속 레이어의 KV 캐시에 표시되지 않는 경우 해당 숨겨진 상태는 후속 반복에서 검색을 위해 Aux 캐시에 저장됩니다.

그림 4에 표시된 것처럼 각 디코딩 단계에서 각 변환기 계층은 먼저 과거 토큰(존재하는 경우)의 KV 캐시를 검색합니다. KV 캐시에 없는 토큰의 경우 이전 레이어를 다시 거치지 않고 이전 레이어의 Aux 캐시에서 숨겨진 상태를 직접 검색합니다. 보조 캐시는 각 변환기 계층에서 각 토큰이 최대 한 번 계산되도록 보장하고 LazyLLM이 가장 느릴 때 표준 LLM보다 빠르도록 보장합니다.

실험

팀은 Llama 2 7B 및 XGen 7B라는 두 가지 대규모 언어 모델에서 이 새로운 "게으른" 접근 방식을 테스트했습니다. 비교를 위한 표준 LLM은 추가 교육 없이 공개적으로 출시된 사전 교육된 체크포인트 모델과 동일합니다.

실험적 벤치마크는 긴 콘텐츠 이해를 위한 멀티태스킹 벤치마크인 LongBench입니다. LongBench 벤치마크에는 16개의 데이터 세트가 포함되어 있으며 단일 문서 Q&A, 다중 문서 Q&A, 요약, 소수 학습, 합성 작업 및 코드 완성을 포함한 6개 작업이 포함됩니다.

평가 지표는 TTFT 가속 대 정확도 균형 측면에서 각 방법의 효율성과 효율성을 나타냅니다.

결과

표 1은 LazyLLM, 표준 LLM 및 기타 기본 방법에 대한 TTFT 속도 향상 및 정확도 결과를 제공합니다.



이 표에서 기준선은 표준 LLM 추론을 나타냅니다. 무작위 토큰 드롭은 토큰에 대해 무작위 가지치기를 수행하는 것을 의미합니다. 정적 토큰 프루닝은 사전 채우기 단계에서 이전 Transformer 레이어의 Attention 방식을 기반으로 입력 토큰에 대해 일회성 프루닝을 수행하는 것을 의미합니다. 프롬프트 압축은 LLM을 사용하여 입력 컨텍스트에서 중복성을 제거하는 프롬프트 압축 방법입니다.

표 1에서 볼 수 있듯이 LazyLLM은 TTFT 가속에서 전반적으로 우수하지만 정확도 감소는 기본적으로 무시할 수 있습니다. 프롬프트를 압축하기 위해 LLM을 사용하려면 많은 계산이 필요하다는 점을 지적해야 합니다. 따라서 프롬프트 압축을 사용하면 추론 속도가 빨라지지만 실제 TTFT는 표준 LLM보다 깁니다.

전체 빌드 속도에 미치는 영향

새로운 방법이 전체 생성 속도에 미치는 영향을 평가하기 위해 팀은 계산에 사용된 프롬프트 토큰의 비율과 생성 가속화를 분석했습니다(표 2 참조).



LazyLLM 계산에 사용된 토큰의 비율은 항상 100% 미만임을 알 수 있습니다. 이는 LazyLLM이 생성이 끝날 때 프롬프트의 모든 토큰을 사용하지 않았음을 나타냅니다. 그러나 이론적으로는 모델이 모든 토큰을 사용할 수 있습니다. 이는 다양한 작업에 대한 전체 생성 프로세스를 추가로 가속화할 수 있습니다.

다양한 레이어의 드롭률

또한 팀은 가지치기 레이어의 위치와 가지치기된 토큰의 개수가 미치는 영향을 분석했습니다. 결과는 그림 6에 나와 있습니다.



동일한 Transformer 계층에서 가지치기를 수행할 경우 남은 토큰이 적을수록 모델 성능이 저하되는 것을 볼 수 있습니다. 이는 우리의 직관적인 이해와도 일치합니다. 또한, 이전 Transformer 레이어에서 프루닝을 수행한 것과 비교하여, 이후 레이어에서 프루닝을 수행하면 더 나은 성능을 얻을 수 있으며, 이는 이후 레이어가 토큰 프루닝에 덜 민감함을 보여줍니다.

이러한 관찰을 바탕으로 프로그레시브 토큰 프루닝의 효율성이 입증되었다고 할 수 있습니다.

점진적인 KV 성장

마지막으로 팀은 토큰 가지치기 논리를 사용하여 모델의 내부를 이해하려고 노력했습니다. 구체적으로 프롬프트 토큰의 누적 사용 비율과 그에 따른 미사용 비율을 알고 싶어합니다. 이 "누적 토큰 사용량"은 각 단계의 KV 캐시 크기로 동일하게 정의될 수 있습니다. 그림 7은 LazyLLM의 각 단계에서 이러한 프롬프트 토큰의 누적 사용량을 보여줍니다.



이 결과는 모델이 이론적으로 프롬프트에서 모든 토큰을 사용할 수 있음에도 불구하고 많은 토큰이 모델에 의해 선택되지 않을 것이라는 가설을 뒷받침합니다.

모델이 작업 수행의 정확성을 여전히 유지할 수 있다는 점을 고려하면, 모델은 출력 품질에 영향을 주지 않는 토큰을 효과적으로 폐기할 수 있다는 결론을 내릴 수 있습니다.