StreamingLLM - 어텐션 싱크로 효율적인 스트리밍 언어 모델 구현

 (github.com/mit-han-lab)
1P by GN⁺ | ★ favorite | 댓글 1개
  • StreamingLLM은 효율성과 성능을 희생하지 않고 LLM을 무한 길이 입력 스트리밍 애플리케이션에 배포하기 위한 프레임워크임
  • 긴 상호작용이 필요한 다중 라운드 대화에서 이전 토큰의 Key/Value 상태 캐싱이 큰 메모리를 사용하고, 일반적인 LLM이 학습 시퀀스 길이보다 긴 텍스트에 일반화하지 못하는 문제를 다룸
  • 최근 KV만 캐싱하는 window attention은 텍스트 길이가 캐시 크기를 넘으면 실패하며, StreamingLLM은 초기 토큰의 KV를 유지하는 attention sink로 window attention 성능을 크게 회복함
  • 유한 길이 attention window로 학습된 LLM이 fine-tuning 없이 무한 시퀀스 길이에 일반화하도록 하며, Llama-2, MPT, Falcon, Pythia에서 최대 4백만 토큰 이상으로 안정적이고 효율적인 언어 모델링을 수행함
  • 스트리밍 설정에서 sliding window recomputation baseline보다 최대 22.2배 speedup을 보임
  • 컨텍스트 윈도우를 확장하지 않으며, 최근 토큰과 attention sink만 유지하고 중간 토큰은 버림
    • Llama-2가 4096 토큰 컨텍스트 윈도우로 사전학습된 경우 StreamingLLM의 Llama-2 최대 캐시 크기도 4096임
    • 긴 책을 입력하면 모델은 최신 토큰만 인식하므로 결론 부분만 요약할 수 있음
  • 적합한 사용 사례는 지속 동작이 필요하고 과거 데이터 의존이나 큰 메모리 사용을 피해야 하는 스트리밍 애플리케이션이며, 예시로 다중 라운드 대화와 LLM 기반 일일 어시스턴트가 제시됨
  • 최근 컨텍스트 확장 방법과 직교적이며 통합 가능하고, StreamingLLM 맥락의 context extension은 더 큰 캐시 크기로 더 많은 최신 토큰을 저장하는 가능성을 뜻함
  • 실행 예시는 examples/run_streaming_llama.py --enable_streaming이며, 환경 설정에는 Python 3.8, torch, transformers==4.33.0, accelerate, datasets, evaluate, wandb, scikit-learn, scipy, sentencepiece가 사용됨
  • 핵심 코드는 Llama-2, MPT, Falcon, Pythia를 포함해 공개 완료 상태이며, perplexity 평가 코드와 Streaming Llama Chatbot 데모도 공개 완료, StreamEval 데이터셋과 평가 코드는 미공개 상태임

함께 보면 좋은 글 β

GN⁺   [-]
Hacker News 의견들

  • 이 작업을 전체 밀집 어텐션과 같은 것으로 오해하는 것 같음
    여기서 말하는 건 멀리 떨어진 내용을 볼 수 있게 한다는 게 아니라, 슬라이딩 윈도우를 매번 다시 계산하면서 L² 비용을 T번 치르는 대신 캐시를 재사용해 퍼플렉서티를 유지하는 효율 개선에 가까움
    테스트도 Q A Q A Q A Q A...처럼 이어 붙여 측정한 것이지, Q Q Q Q A A A A...처럼 한참 뒤의 답을 찾게 한 게 아님
    퍼플렉서티 측정은 “읽을 만한 텍스트”, 즉 지역적으로 그럴듯한 문장을 만든다는 뜻이지, 어텐션이 닿지 않는 큰 삼각형 공백에서 뭔가를 “추출”한다는 증거는 아님
    책을 주고 각 문단의 첫 단어를 쓰게 하거나, 각 장을 한 문장으로 요약하게 하면 실패할 것 같음

    • 저자들이 README에 이 부분을 직접 다루는 FAQ를 추가했음: https://github.com/mit-han-lab/streaming-llm#faq
      직접 테스트해보니 컨텍스트 길이 확장을 해주는 것 같지는 않았고, 실행은 빠른 편이었음
      A100 메모리를 약 35GB 사용했고 실행 내내 사용량이 고정돼 있었음
      Project Gutenberg에서 책을 가져와 문단 단위로 나눠 하나씩 넣고 각 문단마다 “okay”라고 답하게 한 뒤 마지막에 질문했더니 답을 완전히 환각했음
      참고로 10분 정도 만지는 동안 기본 모델인 lmsys/vicuna-13b-v1.3이 영어로 답하게 만드는 것도 잘 안 됐음
      https://gist.github.com/bluecoconut/9cae9e91fe3b1616ed650a96...
    • 맞는 말이지만, “무한 길이 입력”이라는 표현은 독자 입장에서 오해를 부르기 쉬움
      그래도 흥미로운 작업이고, 핵심은 Figure 2의 발견으로 보임
      첫 두 레이어는 최근 토큰에 더 주목하는 지역 패턴을 보이지만, 그 아래를 지나면 모델이 모든 레이어와 헤드에서 초기 토큰에 강하게 주목함
      저자들은 이를 “attention sinks”라고 부르며, 의미적으로 중요하지 않아도 Softmax 때문에 어텐션 점수의 합이 1이 되어야 하므로 남는 어텐션이 어딘가로 가야 한다고 봄
      자기회귀 언어 모델에서는 초기 토큰이 거의 모든 이후 토큰에서 보이기 때문에 이런 싱크 역할을 하도록 학습되기 쉽다는 설명임
      StreamingLLM은 LLM의 어텐션 윈도우를 잘라 쓸 때 생기는 이 이상한 동작을 보정하는 “해킹”에 가깝고, Softmax 사용의 균열이 드러난 사례라서 컨텍스트 길이에 유연한 LLM을 원한다면 다른 함수가 더 나을 수도 있겠다는 생각이 듦

  • 처음 훑어보면 너무 좋아서 사실인지 의심될 정도지만, 작업 품질은 괜찮아 보이고 기법도 놀랄 만큼 단순함
    아이디어는 각 레이어에서 첫 토큰과 슬라이딩 컨텍스트 윈도우에만 어텐션을 적용하고, 그 사이 토큰은 무시하는 것임
    이 말은 각 레이어가 관련 정보를 시퀀스 뒤쪽으로 조금씩 밀어 올려서, 최상위 레이어 끝부분의 슬라이딩 어텐션 윈도우가 그것을 볼 수 있게 한다는 뜻으로 보임
    다만 모든 슬라이딩 윈도우가 커버하는 범위가 전체 시퀀스를 잇기에 충분하지 않으면 중요한 정보를 모두 앞으로 전달하지 못할 수 있음
    예를 들어 모든 윈도우 길이가 같을 때 모델 깊이 × 윈도우 길이 < 시퀀스 길이라면 한계가 생김

    • 시퀀스 끝을 상수인 “중립 값”으로 패딩할 수도 있지 않을까 싶음

  • 이건 Softmax가 합을 1로 만들어야 한다는 관찰 덕분에 가능해진 것으로 보임
    빠르게 보기엔 모델이 이전 토큰에 주목할 필요가 없을 때 첫 토큰을 자리표시자처럼 쓰는 경향이 있음
    이 문제를 처음 본 건 Evan Miller의 HN 글이었는데, 어텐션 헤드가 반드시 이전 토큰들에 모든 어텐션을 배분하게 하는 건 잘못이고, Softmax 분모에 1을 더해 “주목하지 않기”를 허용해야 한다는 내용이었음
    재학습 없이 이 관찰을 활용한 점은 좋고, Evan의 제안을 따랐다면 모델이 어떻게 달라졌을지도 궁금함
    [2] https://news.ycombinator.com/item?id=36851494

    • 실제로는 그 제안을 비슷하게 시험해본 듯함
      모든 값이 0인 전용 싱크 토큰을 두고 모델을 학습시켰는데, 그래도 다른 초기 토큰들이 싱크로 쓰이게 되므로 전용 싱크 토큰이 있는 편이 낫다는 결론처럼 보임
    • HN에서 처음 본 건 그 글이었지만, 그 글에서도 지적됐듯이 Softmax + 1이 처음 제안된 건 아님
      내가 알기로는 실제 성능을 더 좋게 만든 적은 없었음
      학습 후 어텐션 윈도우를 조작할 때는 Softmax + 1이 더 잘 맞을 수도 있지만, 대규모로 시험한 사람이 있는지는 모르겠음

  • 어텐션 캐시 메모리를 추가하는 건 이 문제에 대해 매우 흥미로운 해법임
    며칠 전 Vision Transformer에서 관련 관찰을 한 논문도 나왔음
    Transformer 모델은 전역 정보를 저장할 토큰을 고르는 것처럼 보이고, 일종의 “생각할 토큰”이 필요해 보임
    이 목적의 특정 토큰을 제공하면 성능이 조금 좋아지고, 설명용 시각화도 꽤 흥미롭게 나옴
    [0] https://arxiv.org/pdf/2309.16588.pdf

    • 이미 학습된 모델에 추가 유닛을 넣고 계속 학습하거나 미세조정하는 지점으로 흥미로워 보임
      미세조정에서는 원래 모델 파라미터를 고정하고, 새 “튜닝” 캐시 유닛으로 들어오고 나가는 파라미터만 조정하게 할 수 있음
      그러면 서로 다른 튜닝 유닛 세트를 교체하거나 함께 쓸 수 있음
      비속어 회피 유닛 + 특정 용어 유닛 + 간결하게 쓰는 유닛처럼, 일종의 슈퍼 프롬프트를 섞어 쓰는 방식임
      새 파라미터 수가 충분히 적다면 메모리는 더 들지만 고차 최적화로 빠르고 효과적인 튜닝도 가능할 수 있음
      학습 중 시퀀스 길이와 유닛 수를 함께 늘리는 방법도 생각해볼 수 있음
      짧은 시퀀스에는 몇 개 유닛만 쓰고, 학습 시퀀스 길이를 늘리면서 유닛을 추가하고 계속 학습하는 식임
      임의의 스케줄 대신 성능이나 기울기 분석으로 캐시 확장을 제어할 수도 있을 듯함

  • 저자들이 FAQ를 올렸고, 혼란을 어느 정도 정리해줄 수 있음: https://github.com/mit-han-lab/streaming-llm/blob/main/READM...

    • 업데이트가 좋고, 특히 3번 질문이 핵심을 많이 정리해줌
      “책 같은 긴 텍스트를 StreamingLLM에 넣어 요약할 수 있나?”에 대해, 긴 텍스트를 넣을 수는 있지만 모델은 최신 토큰만 인식하므로 책을 넣으면 마지막 문단들만 요약할 수 있고 별로 유용하지 않을 수 있다고 답함
      즉 LLM의 컨텍스트 윈도우를 확장하거나 장기 기억을 강화하는 것이 아니며, StreamingLLM의 강점은 캐시 새로고침 없이 최근 토큰에서 유창한 텍스트를 생성하는 데 있음

  • 내가 틀릴 수도 있지만, 사람들이 생각하는 것처럼 학습된 길이를 넘어선 내용을 LLM이 참조하게 해주는 기술은 아닌 것 같음
    아마도 긴 텍스트에서 모델 성능을 유지하는 문제, 정확히는 여전히 컨텍스트 윈도우 안에 있는 내용에 대한 성능에 가까워 보임
    모델이 텍스트의 초기 토큰 어텐션에 일종의 짐을 넣도록 학습하고, 그게 윈도우 밖으로 사라지면 깨진다는 설명인데, 왜 그런지는 확실히 모르겠음
    명령 입력이 아니라면 중간 텍스트도 시작 텍스트만큼 좋지 않나 싶음

  • 이런 슬라이딩 윈도우 기법들은 예상하지 못한 명령이 끝에만 나오는 경우를 어떻게 처리하는지 궁금함
    예를 들어 책을 모델에 넣고 마지막 문장에 “이전 입력에서 문자 m의 개수를 반환하라”는 명령이 있다고 해보면, 사람은 한숨을 쉬고 다시 읽으면서 세겠지만 LLM은 입력을 되돌아가 다시 읽을 능력이 없음
    이 예시에서는 문자 세기 자체의 LLM 한계는 무시하더라도, 진짜로 해결하려면 LLM이 임의로 루프를 돌고 점프할 수 있어야 할 것 같음
    물론 그러면 완전히 새로운 문제들이 생기고, 아예 새 아키텍처가 필요할 수도 있음

    • 비슷한 맥락에서, LLM이 읽을 수 있고 접근 가능한 모든 연구 논문을 소화해 색인에 맞는 형식으로 “노트”를 남긴 뒤, 사람이 제한된 말뭉치를 공부한 것처럼 질문에 답할 수 있으면 좋겠음
      질문을 관련 키워드로 바꾸고, 찾아보고, 내용을 다시 훑어 관련 정보를 찾는 방식임
      필요한 전처리가 있다면 “가서 충분히 조사한 뒤 답하는” LLM은 매우 강력할 수 있음
      지난 약 1만 년 동안 우리는 개인 두뇌의 용량과 시간을 넘어서도록 지식 관리 기술을 개선해왔으니, 언어 모델도 단순한 Bing 검색이 아니라 실제 연구 방식과 사전 소화를 활용하게 해야 함
      단기 기억이 어떤 코드 조각이 무엇을 했는지 기억할 필요는 없고, 읽을 때 태그를 붙인 뒤 확장 가능한 공유 태그 색인에 의존하면 됨
      다만 생각할수록 일반적인 LLM 사전학습과 비슷하고, 지식 색인이 거대한 LLM 가중치 덩어리인 것처럼 느껴짐
    • 한 가지 방법은 함수 호출과 비슷하게, LLM이 컨텍스트 파싱 방식을 바꾸는 출력을 낼 수 있게 하는 것임
      LLM 자체의 동작을 바꾸기보다 그 위에 올리는 계층에 가까움
    • 일반적인, 즉 비슬라이딩 컨텍스트 윈도우에서도 LLM이 입력을 다시 읽으려고 되돌아갈 필요가 있나 싶음
      내가 오해한 걸 수도 있지만, 이 경우에는 은닉 상태가 조회 문제를 해결하는 것 아닌가?
      답하기 전에 어차피 전체 입력을 흡수해야 하므로, 명령이 앞에 있든 끝에 있든 어텐션 외에는 영향이 크지 않아 보임
    • 사용자에게 명령을 처음에 넣으라고 하는 게 그렇게 어려운 일인가 싶음
      Claude 100K는 사용자에게 명령을 끝에 넣으라고 요청함
      아니면 빠른 모델로 끝에 명령이 있는지 확인하고 앞쪽으로 가져오면 됨
    • 이 예시는 좀 이상한 경계 사례처럼 보임
      현재 모델이 짧은 입력에서도 이걸 할 수 있는지 잘 모르겠음

  • 약간 농담 섞어 말하면, LLM은 RNN을 다시 발명하려고 정말 열심히 노력 중이고, 도구만 주면 결국 그렇게 할 것 같음

    • RNN이 올바른 해법이지만 실행 비용이 감당하기 어려울 정도로 큼
      다르게 보면 Transformer 모델은 자원 제약이 있을 때 RNN 네트워크의 어느 부분을 “보존할 가치가 있는지” 예측하려는 것임
      현재 Transformer는 단순한 휴리스틱을 쓰고 있고, 이 결과는 그 휴리스틱을 더 좋게 만듦
      많은 NP-완전 문제처럼 완벽하게 맞지는 않아도 유용한 근사가 있을 수 있으며, Transformer는 신경망에서도 그게 가능하다는 걸 보여줌
    • 그런 프로젝트 중 하나가 RWKV임
      오픈소스 리더보드에서 한동안 중간 정도 위치에 있었으니 꽤 정당한 접근이고, 단지 유행하지 않을 뿐임
      [1]: https://huggingface.co/blog/rwkv
    • 많은 사람이 그렇게 믿는 것 같음
      Transformer가 RNN보다 가진 주된 장점은 학습 병렬화임
      RNN은 학습 중 기울기 소실이 있고, 전체 활용률을 높이기도 어려워 큰 배치가 필요하기 때문에 까다로움
      RWKV 같은 모델의 존재는 Transformer처럼 학습하고 RNN처럼 추론하는 미래가 있을 수 있음을 보여줌
    • 지난 30년 동안 더 작은, 요즘 용어로는 “극도로 작은” 신경망에서 배웠던 많은 것들이 이런 대형 모델에서 다시 검토되고 있음

  • 관련해서 MIT의 Han 교수가 공개 TinyML 강의를 진행 중임
    https://news.ycombinator.com/item?id=37620507
    https://efficientml.ai

답변달기