ATACompressor: Adaptive Task-Aware Compression for Efficient Long-Context Processing in LLMs

기술 요약: ATACompressor

문제 정의

대규모 언어 모델(LLM)은 긴 컨텍스트 입력을 처리할 때, 시퀀스 길이가 지나치게 길어짐에 따라 중요한 정보가 희석되거나 무시되는 "lost in the middle(중간에서 길을 잃는)" 현상이라는 심각한 문제에 직면합니다. 검색 증강 생성(RAG)은 외부 컨텍스트를 제공하지만, 원본 문서 토큰을 그대로 이어 붙이는 단순한 방식은 모델에 과도한 부담을 줄 수 있습니다. 기존의 컨텍스트 압축 방법들은 이를 완화하려고 시도하지만 다음과 같은 구체적인 한계를 가집니다:

• Hard-prompt 방식 (예: Selective-Context, LongLLMLingua)은 불필요한 토큰을 제거하지만, 일반적으로 압축률이 낮습니다.
• Soft-prompt 방식 (예: AutoCompressor, ICAE, 500Compressor)은 텍스트를 밀집된 토큰 시퀀스로 압축하여 높은 압축률을 달ages 하지만, 전체 컨텍스트를 무차별적으로 압축하기 때문에 작업 관련 정보의 손실이 발생할 수 있습니다. 이는 압축 과정에서 작업 특이적 인지(task-specific awareness)가 부족하기 때문입니다. 또한, 이러한 방법들은 대개 고정된 수의 압축 토큰을 사용하므로, 서로 다른 작업의 다양한 정보 밀도에 동적으로 적응하지 못합니다.

방법론: ATACompressor

저자들은 특정 작업 요구 사항에 따라 압축을 동적으로 조정하도록 설계된 프레임워크인 **Adaptive Task-Aware Compressor (ATACompressor)**를 제안합니다. 이 아키텍처는 Selective Encoder, Adaptive Allocation Controller (AAC), 그리고 타겟 LLM(디코더)의 세 가지 주요 구성 요소로 이루어져 있습니다.

1. Selective Encoder ($\varphi$)

전체 컨텍스트를 처리하는 기존의 soft-prompt 압축기와 달리, Selective Encoder는 쿼리($Q$)와 관련된 컨텍스트의 부분($\hat{C}_{Rel}$)만을 식별하고 압축하도록 설계되었습니다.

• 아키텍처: 학습 가능한 LoRA 파라미터를 가진 프리징된(frozen) 대규모 언어 모델(LLM)을 활용합니다.
• 입도(Granularity) 처리: "관련성"의 모호함(문장 또는 문서 수준까지 다양할 수 있음)을 해결하기 위해, 입력 컨텍스트는 훈련 중 정답 입도(ground-truth granularity)와 일치하도록 균일한 청크(예: 단락 또는 문장)로 분할됩니다. 인코더는 쿼리에 답하는 데 필요한 청크만을 유지하고 압축하는 법을 학습합니다.
• 프로세스: 인코더는 결합된 청크와 쿼리를 처리하여, 관련 정보를 나타내는 압축된 토큰 집합($c_1, ..., c_k$)을 출력합니다.

2. Adaptive Allocation Controller (AAC)

AAC는 기존 soft-prompt 방식의 고정된 압축률 문제를 해결합니다. 이는 관련 콘텐츠의 추정 길이에 따라 압축될 토큰의 수($k$)를 동적으로 결정합니다.

• 구성 요소: 경량화된 **Probe ($\zeta$)**와 **Policy Function ($\eta$)**으로 구성됩니다.
• 메커니즘: Selective Encoder가 입력을 처리한 후, Probe는 인코더의 은닉 상태(hidden states)를 분석하여 관련 텍스트의 길이($\hat{L}_{Rel}$)를 추정합니다. 그 후 Policy Function은 미리 정의된 정책 비율($r$)을 사용하여 압축 토큰의 수($k$)를 계산합니다:
  $$k = \min(\frac{\hat{L}_{Rel}}{r}, k_{max})$$
• 이점: 이를 통해 짧은 관련 구간에는 적은 토큰을 할당하고, 긴 구간에는 더 많은 토큰을 할당함으로써, 필수 정보를 보존하면서 자원 활용을 최적화할 수 있습니다.

3. 훈련 및 추론 워크플로우

• 사전 훈련(Pretraining): Selective Encoder와 Probe는 공동으로 최적화됩니다. 인코더는 교차 엔트로피 손실(cross-entropy loss)을 사용하여 관련 컨텍스트(정답)를 재구성하도록 훈련되며, Probe는 Huber loss를 사용하여 관련 텍스트 길이를 예측하도록 훈련됩니다.
• 미세 조정(Finetuning): Selective Encoder는 타겟 LLM이 압축된 키-값(KV) 표현을 기반으로 답변을 생성하는 다운스트림 작업(예: QA)에 대해 추가로 훈련됩니다.
• 추론(Inference): 시스템은 병렬 방식으로 작동합니다. Selective Encoder가 입력을 처리함에 따라, AAC는 첫 번째 압축 토근이 생성되기 전에 필요한 토큰 수($k$)를 예측하여 자기회귀(autoregressive) 흐름에 방해가 없도록 합니다.

주요 기여

1. 작업 인지형 압축(Task-Aware Compression): 작업 관련 콘텐츠만을 필터링하고 압축하는 Selective Encoder를 도입하여, 글로벌 압축 방식에서 흔히 발생하는 핵심 정보 손실을 방지합니다.
2. 적응형 자원 할당(Adaptive Resource Allocation): 관련 콘텐츠의 추정 길이에 따라 압축률을 동적으로 조정하는 AAC를 개발하여, 고정된 토큰 제약을 넘어섰습니다.
3. 입도 인식 훈련(Granularity-Aware Training): 훈련 중에 컨텍스트를 균일한 청크로 사전 분할하는 전략을 통해 다양한 주석 입도(문장 vs 문서 수준)를 처리함으로써 인코더의 적응력을 높였습니다.

실험 결과

저자들은 LLaMA-2-7B와 Qwen-2-7B를 백본 모델로 사용하여 세 가지 QA 벤치마크인 HotpotQA(문서 수준), MSMARCO(단락 수준), SQUAD(문장 수준)에서 ATACompressor를 평가했습니다.

• 성능: ATACompressor는 Exact Match (EM) 및 F1 점수 모두에서 기존의 최첨단 베이스라인(Selective-Context, LongLLMLingua, AutoCompressor, ICAE, 500Compressor, QGC 포함)을 일관되게 능가했습니다. 예를 들어, Qwen-2-7B를 사용한 HotpotQA에서, ATACompressor는 차순위 방법인 QGC의 F1 점수인 72.36에 비해 80.23을 달성했습니다.
• 효율성: 이 방법은 경쟁력 있는 **처리량(Throughput, TP)**을 유지하면서 높은 압축률(Compression Ratios, CR)(최대 ~27배)을 달성했습니다. 복잡한 외부 검색기에 의존하는 QGC와 같은 방법들에 비해 우수한 효율성을 보여주었습니다.
• 절제 연구(Ablation Studies): AAC를 제거(토큰 수 고정)하거나 Selective Encoder를 제거(전체 컨텍스트 압축)했을 때 성능이 크게 저하되었으며, 이는 두 구성 요소의 필요성을 입증합니다.
• 강건성(Robustness): 압축된 토큰 수가 감소하더라도 성능 안정성을 유지하며 낮은 토큰 환경에서도 베이스라인보다 뛰어난 성능을 보였습니다. 또한 LongBench 작업(Summary 및 Few-shot)에서도 강력한 일반화 능력을 보여주었습니다.

의의 및 주장

이 논문은 ATACompressor가 LLM의 효율적인 롱 컨텍스트 처리를 위한 확장 가능한 솔루션을 제공한다고 주장합니다. 압축 과정에 작업 인지 기능을 통합하고 자원을 동적으로 할당함으로써, 정보 보존과 압축 효율성 사이의 트레이드오프를 효과적으로 균형 있게 조절합니다. 저자들은 이 접근 방식이 "lost in the middle" 문제를 완화할 뿐만 아니라 추론 비용과 지연 시간을 줄여, 기본 LLM의 아키텍처 수정 없이도 실시간 및 대규모 애플리케이션에 적합하게 만든다고 강조합니다. 이 연구는 RAG 및 롱 컨텍스트 LLM 애플리케이션을 더욱 실용적이고 효과적으로 만드는 데 있어 진일보한 단계로 제시됩니다.