SWE-Pruner Pro: The Coder LLM Already Knows What to Prune

Pruning long context for coding agents has been a vital technology for efficient context management. While existing context pruning methods such as SWE-Pruner realize this by attaching a separate code classifier, we find the agent itself encodes internal representations indicating the relevance of code context when reading tool output. Based on this finding, we propose SWE-Pruner Pro, which prunes tool outputs directly inside the agent. Concretely, a small head turns the agent's own internal representations into a keep-or-prune label for each line, with a length-aware embedding keyed to each tool output's line count. Across two open-weight backbones and four multi-turn benchmarks, SWE-Pruner Pro saves up to 39% of prompt and completion tokens while preserving task quality, with bounded inference overhead. Notably, on MiMo-V2-Flash SWE-Pruner Pro additionally raises the SWE-Bench Verified resolve rate by +3.8% and the long-context Oolong accuracy by +2.2 points.

한 줄 요약

SWE-Pruner Pro는 코드 생성 에이전트 내부 표현을 활용해 툴 출력을 직접 줄이며, 토큰 소비를 최대 39% 절감하고 성능을 유지한다.

핵심 기여도

기존 별도의 코드 분류기 대신 에이전트 내부 표현을 활용해 줄이기 (SWE-Pruner Pro).
길이 인식 임베딩(length-aware embedding)과 샘플별 균형 focal loss를 도입.
MiMo-V2-Flash 모델에서 SWE-Bench Verified 해결률 +3.8%, Oolong 정확도 +2.2 포인트 향상.
2개 오픈 가중치 백본과 4개 멀티턴 벤치마크에서 평균 30~39% 토큰 절감.

핵심 아이디어

기존 SWE-Pruner는 별도의 분류기를 사용해 툴 출력을 줄였지만, 본 연구는 에이전트가 툴 출력을 읽을 때 이미 내부 표현에 줄일지 유지할지를 판단하는 정보가 포함되어 있다는 점을 발견했다. 이에 따라 SWE-Pruner Pro는 에이전트 내부의 마지막 레이어 히든 상태(last-layer hidden states)를 활용해 줄이기 결정을 내리는 가벼운 헤드를 도입한다. 이 헤드는 툴 출력의 줄 수에 따라 다르게 작동하도록 길이 인식 임베딩을 학습하고, 샘플별 균형 focal loss를 통해 유지 줄이 비중을 조정한다. 이는 별도 모델 호출 없이 서버 내에서 실행되어 추론 오버헤드를 최소화한다.

기술적 접근법

**SWE-Pruner Pro 헤드**: 에이전트의 last-layer hidden states를 기반으로 줄이기/유지 결정.
**Length-aware embedding**: 툴 출력의 줄 수에 따라 헤드의 예측이 달라지도록 설계.
**Per-sample balanced focal loss**: 샘플 내 유지 줄이 비중을 균형 있게 학습, 유지 줄이 드물수록 더 높은 가중치 부여.
**Frozen backbone 활용**: 기존 prefill 과정에서 생성된 토큰 표현 재사용, 별도 모델 호출 없이 실행.
**하이퍼파라미터**: 헤드 실행 시 전체 생성 시간 대비 15.0%의 추가 벽 시간만 소요.

주요 결과

**SWE-QA-Pro**: 39% 토큰 절감, MiMo-V2-Flash 정확도 +2.2 포인트.
**Oolong**: 30% 토큰 절감, MiMo-V2-Flash 정확도 +2.2 포인트.
**SWE-Bench Verified**: MiMo-V2-Flash 해결률 +3.8%.
**추론 오버헤드**: 전체 생성 시간 대비 15.0% 추가 벽 시간.

의의 및 한계

SWE-Pruner Pro는 별도의 분류기나 쿼리 없이 에이전트 내부 표현만으로 줄이기 결정을 내리는 새로운 접근법을 제시한다. 이는 에이전트가 이미 툴 출력의 중요도를 인식하고 있다는 통찰을 기반으로 하며, 토큰 소비를 줄이면서도 성능을 유지하는 데 성공했다. 그러나 이 방법은 특정 백본 모델과 벤치마크에서 평가되었으며, 다른 환경이나 더 복잡한 툴 출력에서는 동일한 효과를 보장하지 못할 수 있다. 또한, 헤드의 학습이 줄이기 결정의 정확도에 직접적으로 영향을 미치므로, 헤드 설계의 최적화가 중요하다.

실용적 활용

SWE-Pruner Pro는 코드 생성 에이전트가 긴 툴 출력을 효율적으로 관리할 수 있도록 해주며, 클라우드 기반 개발 환경, 자동화된 코드 리뷰 시스템, 대규모 리포지토리 관리 등에서 유용하게 활용될 수 있다. 특히, 토큰 비용을 줄이고 추론 시간을 최소화해야 하는 산업 현장에서 실용적 가치가 높다.