기술 요약: SWE-Pruner Pro

문제 정의

멀티 턴 환경에서 작동하는 코딩 에이전트는 상호작관 과정에서 방대한 양의 도구 출력값(예: cat, grep, ls, python 실행 로그)을 축적하게 된다. 이러한 출력값은 종종 상당한 중복성을 포함하며, 이는 에이전트의 턴당 토큰 예산을 대부분 소비하고 긴 컨텍스트 성능 저하(long-context degradation)를 유발한다. 기존의 컨텍스트 프루닝(context pruning) 솔루션은 두 가지 범주로 나뉜다:

1. 범용 압축: LLMLingua2나 Selective Context와 같은 방법들은 고정된 메트릭(예: perplexity, 구문 구조)을 기준으로 토큰을 점수화한다. 이러한 방식은 에이전트의 진화하는 작업 초점에 적응하지 못한다.
2. 태스크 특화 프루닝: SWE-Pruner와 같은 선행 연구는 에이전트의 의도(intent)를 조건으로 사용하지만, 외부 스코어링 모델과 매 턴마다 에이전트가 생성해야 하는 명시적인 "목표 힌트(goal-hint)" 쿼리를 필요로 한다. 이는 추가적인 오버헤드와 지연 시간을 발생시킨다.

두 접근 방식 모두 프루닝 신호를 외부에서 재구성해야 할 양으로 취급하며, 에이전트의 백본 모델이 이미 도구 출력을 처리했고 그 과정에서 내부 표현(internal representations) 안에 관련성 정보를 인코딩했을 가능성을 간과하고 있다.

방법론: SWE-Prner Pro

SWE-Pruner Pro는 외부 스코어링 모델이나 명시적인 "목표 힌트" 쿼리를 필요로 하지 않고, 에이전트 백본의 내부 은닉 상태(hidden states)로부터 직접 프루닝 신호를 읽어내는 방식을 제안한다. 이 방법은 에이전트의 추론 루프 내부에서 직접 작동한다.

핵심 메커니즘

1. 내부 표현 활용: 에이전트가 도구 응답($r_t$)을 읽을 때, 동결된(frozen) 백본은 표준 프리필(prefill)을 수행하여 각 토큰에 대한 마지막 레이어 은닉 상태($h_i$)를 생성한다. SWE-Pruner Pro는 이러한 기존 상태에 경량 프루닝 헤드를 부착한다.
2. 프루닝 헤드 아키텍처:
   • 길이 인식 임베딩(Length-Aware Embedding): 응답의 총 라인 수 $N$에 따라 조건화된 학습된 임베딩 $e(N)$이 각 은닉 상태에 브로드캐스트 방식으로 더해진다. 이를 통해 모델은 응답 길이에 따라 프루닝 전략을 조정할 수 있다 (예: 모든 라인이 중요한 짧은 출력의 경우 더 보수적으로 접근).
   • 비선형 분류기: 작은 피드 포워드 네트워크(LayerNorm + 두 개의 Linear-GELU-Dropout 블록)가 증강된 은닉 상태를 '유지 또는 삭제(keep-or-prune)' 로짓으로 매핑한다. 이러한 비선형성은 선형 프로빙 실험에서 관찰된 표현 공간의 중첩을 해결하기 위해 필요하다.
   • 결정 집계: 토큰 수준의 로짓은 각 라인 내에서 다수결 투표를 통해 집계되어, 전체 라인에 대한 이진 유지/삭제 결정을 생성한다.
3. 학습 전략:
   • 데이터: 헤드는 각 라인의 유지/삭제 라벨이 부여된 22,609개의 멀티 턴 궤적을 Claude Sonnet 4.6으로 어노테이션하여 학습되었다.
   • 손실 함수: 저자들은 **샘플별 균형 잡힌 포컬 로스(per-sample balanced focal loss)**를 사용한다. 전역적으로 균형을 맞추는 표준 포컬 로스와 달리, 이 방식은 각 샘플 내에서 "유지"와 "삭제" 토큰에 대해 손실을 별도로 계산하고 이를 동일한 가중치로 평균을 낸다. 이는 극단적인 유지율(예: 라인의 3%만 유지되는 경우)을 가진 샘elle이 평균 유지율(~30%)에 과적합되는 것을 방지하여, 모든 샘플이 목적 함수에 동일하게 기여하도록 보장한다.
4. 추론 통합: 헤드는 프리필 단계 동안 "인-서버(in-server)"에서 실행된다. 이는 도구 응답을 위해 생성된 은닉 상태를 재사용하며, 작은 헤드의 단일 순방향 패스(forward pass)만을 추가한다. 프루닝된 응답($\tilde{r}_t$)은 다음 턴의 컨텍스트에서 원본 응답을 대체하지만, 현재 턴의 생성은 여전히 전체 원본 응답을 참조(attend)한다.

주요 기여

1. 신호 발견(Signal Discovery): 본 논문은 에이전트 백본의 내부 표현이 이미 도구 출력의 라인별 중요도를 인코딩하고 있음을 입증한다. 단순한 선형 프로브가 AUC 0.83을 달해, 외부 재구성이 없이도 "유지 또는 삭제" 신호가 존재함을 증명했다.
2. 아키텍처: 길이 인식 임베딩과 샘플별 균형 잡힌 포컬 로스를 활용하여, 동결된 백본의 상태로부터 직접 정밀한 프루닝 결정을 내리는 경량 헤드인 SWE-Pruner Pro를 소개한다.
3. 효율성: 추가적인 모델 호출이나 명시적인 목표 힌트 쿼리를 피함으로써, 제한된 추론 오버헤드를 추가하면서도 상당한 토큰 절감을 달ien한다.

실험 결과

방법론은 두 가지 오픈 웨이트 백본(MiMo-V2-Flash 및 Qwen3-Coder-Next)을 사용하여 네 가지 멀티 턴 벤치마크(SWE-Bench Verified, SWE-QA, SWE-QA-Pro, Oolong)에서 평가되었다.

• 토큰 절감: SWE-Pruner Pro는 평가된 7가지 방법 중 모든 설정에서 엔드 투 엔드 토큰 사용량을 줄이는 유일한 방법이다.
  • SWE-QA-Pro (Qwen3-Coder-Next)에서 최대 39% 감소.
  • 롱 컨텍스트 벤치마크인 Oolong (MiMo-V-Flash)에서 30% 감소.
  • 반면, 다른 방법들(예: LLMLingua2)은 롱 컨텍스트 태스크에서 오버헤드로 인해 토큰 사용량이 최대 190%까지 증가하기도 했다.
• 태스크 품질:
  • SWE-QA/Pro: 프루닝되지 않은 베이스라인과 매우 좁은 범위 내에서 태스크 품질을 유지한다 (예: SWE-QA에서 +0.02 점수 변화).
  • Oolong: MiMo-V-Flash에서 정확도를 +2.2 포인트 향상시킨다.
  • SWE-Bench Verified: MiMo-V-Flash에서, 프루닝되지 않은 베이스라인 대비 입력 토큰을 7.4% 줄이면서 해결률(resolve rate)을 +3.8% 높였다. Qwen3-Coder-Next에서는 가장 큰 입력 토큰 감소(-13.5%)와 함께 가장 적은 해결률 저하(-1.2%)를 달성했다.
• 지연 시간(Latency): 인-서버 헤드는 전체 생성 시간에 약 **15.0%**의 총 벽시계 시간(wall time)을 추가하지만, 이는 후속 턴의 컨텍스트 크기 감소로 상쇄된다.

의의 및 주장

본 논문은 컨텍스트 프루닝에 필요한 "관련성 판단"이 추가적인 쿼리나 모델을 통해 이끌어내야 하는 외부 신호가 아니라, 백본이 관찰을 수동적으로 읽는 동안 형성된 표현 속에 이미 존재한다는 점을 주장한다.

• 패러다임 전환: 프루닝을 에이전트 주변에 구축하는 것(외부 모델이나 명시적 쿼리 요구)이 아니라, 에이전트가 이미 형성한 신호를 단순히 읽는 것이다.
• 실질적 영향: 이 접근 방식은 선행 연구에서 관찰된 품질-압축 트레이드오프 문제를 해결하여, 태스크 성능을 저하시키지 않으면서도 공격적인 프루닝(최대 39% 토큰 절감)을 가능하게 하며, 때로는 롱 컨텍스트 성능 저하를 완화함으로써 성능을 오히려 향상시킨다.
• 일반화 가능성: 파이썬 중심의 코딩 태스크에 대해 평가되었으나, 파이프라인은 언어에 구애받지 않으며 자연어 롱 컨텍스트 태스크(Oolong)에서도 토큰 절감과 품질을 보존하는 것으로 나타나, 이 패러다임이 다양한 도메인으로 전이될 수 있음을 시사한다.

저자들은 코딩 에이전트를 위한 가장 효율적인 컨텍스트 관리 경로은 백본 자체의 내부 상태를 활용하는 것이며, 이를 통해 외부 스코어링 메커니즘의 지연 시간과 복잡성을 피할 수 있다고 결론짓는다.