Quest Query-Aware Sparsity for Efficient Long-Context LLM Inference

논문 : https://arxiv.org/abs/2406.10774

개인의견 : 역시 songhan.. KV cache는 항상 볼때마다 압축을 하거나 줄이고 싶지만 그 특성상 query와 관계를 알아야 이게 중요한 KV였다는걸 알 수 있어 고육지책으로 들고다녀야하는 부담이라고 생각이 드는데 이걸 참 재밋게 풀어 경량화를 하는거 같아서 흥미롭게 읽었습니다.

아래글은 Scholar GPT와 대화를 나눈 내용입니다.

Q : 논문의 핵심 내용과 강점, 알고리즘 설명, 그리고 한계점

Answer

논문의 강점과 독창성

이 논문에서 제안하는 “Quest” 알고리즘은 긴 문맥을 가진 대형 언어 모델(LLM)의 추론 속도를 획기적으로 향상시키는 독창적인 방법을 제시합니다. 주목할 만한 강점은 다음과 같습니다:

  1. 질문-기반 희소성(Query-Aware Sparsity): Quest는 단순히 기존의 희소성을 사용하는 대신, 현재의 쿼리(token)를 기준으로 중요한 토큰을 동적으로 선택하여 계산 부하를 줄입니다. 이는 기존의 고정된 희소성 접근법보다 효율적이며, 각 쿼리에 맞춘 더 높은 정확성을 제공합니다.
  2. KV 캐시 선택 최적화: 기존 방법들이 과거의 주목(attention) 기록을 기반으로 중요 토큰을 식별하는 반면, Quest는 쿼리의 동적 특성에 맞춰 중요한 토큰을 실시간으로 파악합니다. 이로 인해 높은 정확도를 유지하면서도 KV 캐시에서 꼭 필요한 데이터만 활용하여 메모리와 연산 부담을 줄입니다.
  3. 효율적 메모리 사용: Quest는 긴 문맥에서도 전체 캐시를 유지하는 대신 일부 페이지를 선택해 메모리 이동을 줄입니다. 이 방식은 대규모 문서에서 최적의 성능을 발휘하도록 설계되었습니다.

핵심 알고리즘 설명: 예시 입력과 과정

Quest의 작동 과정을 예시를 들어 설명하겠습니다.

1단계: 중요 페이지 추정

예를 들어, 쿼리 벡터 ( Q )가 주어지고, KV 캐시가 페이지 단위로 나누어져 있으며, 각 페이지마다 최솟값(( mi ))과 최댓값(( Mi ))이 저장되어 있다고 가정합니다.

  • 현재 쿼리 ( Q )와 각 페이지의 ( mi ) 및 ( Mi )를 활용해 중요도를 추정합니다.
  • 예를 들어, ( Q = [1, 0.5, -0.2] )이고, 특정 페이지의 ( mi )와 ( Mi )가 각각 ( [-0.5, -0.1, -0.3] )와 ( [0.7, 1.2, 0.9] )일 때, Quest는 각 차원에서 최댓값을 ( \max(Q_i \cdot mi, Q_i \cdot Mi) )로 계산합니다. 이 과정에서 계산된 값을 합산하여 페이지의 중요도를 결정합니다.

2단계: 상위 K개의 중요 페이지 선택 후 주목 수행

모든 페이지에 대해 계산된 중요도 값 중 상위 K개만 선택합니다. 이후 선택된 페이지들에 대해 일반적인 self-attention을 수행합니다. 이렇게 하면 불필요한 데이터를 줄이고 쿼리와 관련된 정보만을 활용해 더 빠른 추론이 가능합니다.

한계점

이 논문의 한계점도 존재합니다:

  1. 초기 비용: Quest는 초기 메타데이터로 각 페이지의 최소 및 최대값을 저장해야 하므로 초기 단계에서 메모리 비용이 발생합니다.
  2. 하이퍼파라미터 K의 조정 필요성: 상위 K개의 중요 페이지를 선택하는 것이므로, 최적의 K 값 설정이 성능에 큰 영향을 미칠 수 있습니다. 다양한 데이터셋에 따라 최적의 K 값을 찾는 것이 필요합니다.
  3. 모델 구조 의존성: Quest는 페이지 단위의 캐시 관리와 같이 특정 구조에 최적화된 방법이기 때문에, 다른 구조의 모델에 적용할 때는 제한이 있을 수 있습니다.

종합적으로 Quest는 긴 문맥 추론의 효율성을 크게 향상시키는 독창적인 접근법으로, 효율성과 정확성의 균형을 맞추기 위해 쿼리 인지적 방식의 희소성을 성공적으로 활용하고 있습니다.

Q : 중요도를 계산하는 방법에 대해서 구체적으로 알려줘 그리고 KV cache 에서 eviction이 된 cache는 완전히 사라지게 되는건지 다음 decode step에서만 사용되지 않는지가 궁금해

Answer :

중요도 계산 방법

Quest 알고리즘에서 중요도 계산은 주어진 쿼리 ( Q ) 벡터와 각 KV 캐시 페이지의 최소값 ( m_i )와 최대값 ( M_i )을 이용하여 이루어집니다. 구체적인 계산 방식은 다음과 같습니다:

  1. 페이지의 메타데이터 확인:
    • 각 페이지에는 키(Key) 벡터의 각 차원별 최솟값 ( m_i )와 최댓값 ( M_i )이 저장됩니다. 이를 통해 해당 페이지의 키 값의 범위를 미리 알 수 있습니다.
  2. 중요도 추정 계산:
    • 주어진 쿼리 벡터 ( Q = [q_1, q_2, \dots, q_d] )에 대해, 각 차원 ( i )에서 중요도를 추정하는 상한값 ( U_i )를 계산합니다.
    • 각 차원의 중요도 상한값 ( U_i )는 다음과 같이 계산됩니다: [ U_i = \max(q_i \cdot m_i, q_i \cdot M_i) ]
    • 이렇게 계산된 ( U_i )는 현재 쿼리 ( Q )가 해당 페이지의 모든 키 값 중 가장 높은 중요도를 가지는 경우의 상한을 나타냅니다. 이 상한값을 사용하면 각 페이지의 최대 중요도를 효율적으로 추정할 수 있습니다.
  3. 페이지 중요도 점수 계산:
    • 모든 차원에 대해 계산된 ( U_i )의 합을 구하여 페이지의 중요도를 정의합니다: [ \text{중요도 점수} = \sum_{i=1}^{d} U_i ]
    • 이렇게 계산된 페이지별 중요도 점수로, Quest는 전체 페이지 중에서 상위 K개를 선택합니다. 이 상위 K개 페이지에 대해서만 self-attention을 수행하여 계산 자원을 절약합니다.

KV 캐시에서의 Eviction 처리 방식

Quest에서는 KV 캐시의 모든 데이터를 삭제(eviction)하는 대신, 특정 페이지의 중요도에 따라 필요 없는 페이지는 다음 decode 단계에서 제외되는 방식으로 관리합니다. 따라서 캐시에서 제거된 데이터는 완전히 삭제되지 않고, 현재 decode 단계에서만 사용되지 않으며, 필요 시 다시 사용 가능합니다.

이 방식의 장점은, 특정 토큰이 현재 단계에서는 중요하지 않지만 이후 단계에서 중요해질 가능성이 있는 경우를 대비할 수 있다는 점입니다. 예를 들어, 질문이 바뀌거나 새로운 정보가 추가되면 해당 정보가 필요해질 수 있으므로, 캐시 데이터를 완전히 삭제하지 않고 유지함으로써 필요한 순간에 바로 사용할 수 있게 합니다.

따라서, Quest의 Eviction은 일시적으로 사용을 제한할 뿐, KV 캐시에 대한 완전한 삭제를 의미하지는 않습니다.

Q : 음 그렇다면 중요도를 추산하는 방법의 의미는 Query와 Value 의 correlation이 큰값이 attention이 되어야할 중요한 값이므로 eviction을 하지 않고 남겨두겠다고 이해하면 될까?

Answer:

네, 정확하게 이해하셨습니다. Quest에서 중요도 추산 방법은 쿼리(Query)와 키(Key) 값 간의 상관관계를 활용해 중요도를 평가합니다. 이 상관관계가 클수록 해당 토큰이 현재 쿼리에 대해 중요한 정보를 제공하므로, Quest는 이런 토큰들을 주목(attention) 대상에 포함시키고 캐시에서 유지합니다.

구체적으로, 쿼리 벡터와 키 벡터의 값이 높은 상관관계를 가지는 페이지는 해당 쿼리에 중요한 영향을 미치는 정보로 간주되어, 이 정보를 보존하여 다음 decode 단계에서 사용되도록 합니다. 반대로, 상관관계가 낮은 페이지는 다음 단계에서는 사용되지 않도록 하여 연산 비용을 절약합니다.

이렇게 함으로써 Quest는 중요한 정보만을 선택적으로 유지하고 나머지는 일시적으로 생략하여, 긴 문맥에 대한 추론 과정에서 메모리 이동과 연산 부하를 줄이면서도 모델의 성능을 유지할 수 있습니다.

Q: 이 논문에서 하이퍼파라미터 K는 context length가 길어짐에 따라 바뀌게 되는지 아니면 고정인지가 궁금해 그리고 이것이 동적으로 바뀌어야할 필요성이 있는지도 검토해줘ㅏ

Answer :

논문에서의 하이퍼파라미터 ( K )는 고정된 값으로 사용되며, 이는 중요한 KV 캐시 페이지 수를 의미합니다. 이 값은 모델과 데이터셋의 특성에 맞춰 실험적으로 선택되며, 보통 사전 설정된 일정한 값으로 유지됩니다.

( K )를 고정하는 이유

  1. 일관된 메모리 관리: 고정된 ( K ) 값은 메모리 사용량을 예측 가능하게 해줍니다. 이는 장기적 문맥 추론을 처리할 때, 메모리 이동 비용을 안정적으로 유지하는 데 유리합니다.
  2. 모델 안정성 유지: ( K )를 고정하면 모델이 과도하게 많은 페이지를 선택하거나 너무 적은 페이지를 선택하는 것을 방지해, 일관된 성능을 보장할 수 있습니다.

( K )가 동적으로 바뀌어야 할 필요성

긴 문맥 길이와 복잡성이 증가하는 경우, 동적으로 ( K )를 조정하는 방안이 더 효율적일 수 있습니다. 이유는 다음과 같습니다:

  1. 문맥 길이에 따른 정보 밀도 변화: 긴 문맥일수록 특정 순간에 주목해야 할 정보가 더 많아질 수 있습니다. 따라서, 문맥 길이에 따라 ( K ) 값을 확장하면 필요한 정보를 놓치지 않고 유지할 가능성이 높아집니다.

  2. 모델 효율성 개선: 문맥이 짧은 경우에는 많은 KV 페이지를 유지할 필요가 없을 수 있습니다. 반면에 긴 문맥에서는 더 많은 정보를 포함하여 추론 성능을 향상시킬 수 있으므로, 상황에 따라 ( K )를 늘려 효율적으로 중요한 정보를 보존하는 것이 유리할 수 있습니다.

동적 ( K ) 적용의 도전 과제

동적으로 ( K )를 변경하면 계산 복잡성이 증가할 수 있으며, ( K )를 결정하는 기준도 추가로 설계해야 합니다. 예를 들어, 문맥 길이 외에도 현재 쿼리의 복잡도나 입력 데이터의 중요도에 따라 ( K )를 유연하게 조정할 수 있는 메커니즘이 필요합니다. 이를 통해 시스템 성능을 높일 수 있지만, 제어 메커니즘이 추가되면서 시스템의 복잡도도 증가할 수 있다는 점은 고려해야 합니다.

결론적으로, ( K )를 동적으로 설정하면 긴 문맥 추론에서 더 높은 효율성을 얻을 수 있을 가능성이 있지만, 구현의 복잡성이 증가하는 단점이 있습니다. 현재는 고정된 ( K ) 값이 사용되지만, 향후 연구에서는 동적 ( K ) 설정의 가능성을 탐구해볼 가치가 있습니다.

Q : page의 크기에 따라서 정확도 및 효율성이 변할꺼 같은데 논문에 관련 실험이 있으면 설명해주고 어떤 값을 사용하는게 좋은지 근거를 들어 알려주겠니?

Answer:

논문에서는 페이지 크기에 따른 정확도와 효율성의 관계를 실험적으로 평가하고 있습니다. 페이지 크기(Page Size)는 각 페이지에 포함되는 토큰 수로, 이는 모델이 한 번에 관리하는 KV 캐시의 데이터 블록 크기를 의미합니다. 페이지 크기는 효율성(속도)과 정확도(성능)에 직접적인 영향을 미치며, 다음과 같은 이유로 중요합니다.

실험 결과 요약: 페이지 크기와 효율성 및 정확도의 관계

  1. 페이지 크기가 작을수록 정확도 유지에 유리:
    • 페이지 크기를 줄이면 더 세밀하게 중요한 토큰을 선택할 수 있어, 정확도가 향상될 수 있습니다. 세부적인 데이터 구분이 가능해지기 때문에 중요한 정보를 더 많이 보존할 수 있게 됩니다.
  2. 페이지 크기가 클수록 효율성(속도) 개선:
    • 반면에 페이지 크기가 크면 한 번에 더 많은 데이터를 처리하므로 메모리 이동과 계산이 줄어들어 효율성이 높아집니다. 큰 블록 단위로 데이터를 가져오면 모델이 필요한 메모리 접근 횟수가 줄어들어 속도가 빨라지기 때문입니다.

논문에서의 권장 페이지 크기

논문에서는 페이지 크기 16을 사용하여 최적의 균형을 맞추고 있습니다. 이 값은 실험적으로 다양한 페이지 크기를 테스트한 결과로 결정된 것으로, 높은 정확도를 유지하면서도 효율적인 추론 속도를 제공하는 값으로 나타났습니다. 페이지 크기가 16일 때 모델이 긴 문맥을 처리하는데 있어 효율적이며, 필요할 때 중요한 정보를 충분히 포함하여 정확도를 유지하는 데 도움을 줍니다.

페이지 크기 선택의 기준 및 고려 사항

  • 작은 페이지 크기 (예: 8 이하): 높은 정확도를 제공하지만, 메모리 이동 횟수가 많아져 효율성이 떨어질 수 있습니다.
  • 큰 페이지 크기 (예: 32 이상): 높은 효율성을 제공하지만, 정확도가 낮아질 위험이 있습니다. 중요한 정보가 상대적으로 적은 페이지에 포함될 가능성이 증가하여 정보 손실이 발생할 수 있습니다.

따라서, 페이지 크기를 16으로 설정하는 것이 정확도와 효율성 사이의 균형을 맞추는 데 적합한 선택으로 보이며, 이는 논문의 실험적 근거에 기반한 권장 설정입니다.