SWE-Pruner Pro: 코더 LLM이 이미 무엇을 제거해야 하는지 알고 있다

코딩 에이전트는 다중 턴 동안 툴 출력으로 인해 컨텍스트가 빠르게 비대해져 비용과 추론 품질 저하가 발생한다. SWE-Pruner Pro는 별도 스코어러나 추가 질의 없이 이미 에이전트가 형성한 표현에서 라인별 중요도를 직접 읽어 컨텍스트를 압축하므로 불필요한 외부 호출과 그에 따른 오버헤드를 제거한다. 이 접근은 토큰 사용량과 응답 시간의 트레이드오프를 줄여 대화형 코드 작업에서 처리량과 비용 효율을 동시에 개선할 가능성을 제시한다.

코딩 에이전트는 툴 호출을 통해 수백 줄의 텍스트를 받아 컨텍스트 창에 누적하며 이후 턴에서도 이 정보를 조건으로 사용한다. Transformer 기반 백본은 툴 응답을 읽을 때 attention 가중치와 포워드 연산을 통해 어떤 토큰과 라인이 이후 행동에 중요한지를 내재적으로 반영한 은닉 표현을 생성한다. 따라서 중요한 라인 여부는 이미 모델 내부 표현에 인코딩되어 있으며, 이를 별도 모델이 재생성할 필요가 없다는 것이 출발점이다. SWE-Pruner Pro는 이 출발점에 근거해 백본의 마지막-레이어 토큰 임베딩을 읽어 토큰별 keep 로그it를 계산하고 라인 단위로 집계하는 방식을 채택했다. 추가 모델 호출이나 명시적 goal-hint를 요구하는 기존 task-specific pruner와 달리, 프리필 과정에서 생성된 은닉 상태를 재사용하므로 외부 오버헤드가 제거된다. 길이 의존성은 임베딩으로 처리하고 토큰 단위의 불확실성은 비선형 헤드가 흡수하도록 설계해 중간 점수 영역의 판별 성능을 개선한다. 이 접근은 토큰 절감과 품질 보존 사이의 트레이드오프를 완화한다는 점에서 핵심적이다. 실험에서는 평균 keep-rate 약 30%가 관찰되었고, 헤드는 긴 응답에서는 보다 공격적인 제거를, 짧은 응답에서는 보수적 결정을 하도록 길이 인식 임베딩을 통해 동작이 조정되었다. 이로 인해 다음 턴에서의 재포워드 비용 감소가 초기 프루닝 오버헤드를 상쇄해 장기적으로 처리량과 비용 효율을 개선했다.

전체 파이프라인은 각 턴에서 에이전트가 툴 호출 c_t을 수행하고 환경이 raw response r_t를 반환하면 시작된다. 에이전트는 [H_{t-1}, c_t, r_t]를 prefill하여 KV 캐시에 채우고, 이 과정에서 새로 전달된 r_t의 토큰들만이 포워드되어 마지막-레이어 은닉 상태 h_i를 생성한다. SWE-Pruner Pro는 이 시점에서 추가 백본 포워드 없이 h_i를 읽어 토큰별 keep 로그it를 예측하고 라인 단위 다수결로 최종 keep/prune 결정을 만든다. 프루닝 헤드는 두 요소로 구성된다. 첫째는 라인 수 N을 입력으로 하는 길이 인식 임베딩 e(N)이며, 이는 모든 토큰의 마지막-레이어 은닉에 덧셈되어 길이 조건을 부여한다. 둘째는 토큰별 분류기 f_(theta)로 LayerNorm, 두 블록의 Linear→GELU→Dropout, 마지막 Linear를 포함하며 각 증강된 은닉 tilde{h}_i에 대해 로그it z_i를 산출하고 시그모이드로 확률 p_i를 얻는다. 라인 결정은 각 라인 내 토큰별 이진화 결과의 다수결로 얻으며, 학습 시에는 라인 레이블을 토큰 레이블로 확장해 토큰 수준 손실로 학습한다. 학습 전략은 공개된 다중 턴 궤적에서 라인별 keep/prune 레이블을 얻어 22,609 샘플로 헤드를 학습한 것이다. 손실 함수는 샘플별 균형 포컬 손실로, 토큰별 포컬 스케일링을 적용한 뒤 각 샘플에서 keep 토큰과 prune 토큰의 손실을 따로 평균해 두 클래스를 균형 있게 반영한다. 백본은 완전히 고정(frozen)되어 헤드만 학습하므로 백본 동작 자체는 유지되며 캐시된 특징으로부터 오프라인으로 헤드를 학습할 수 있다.

주요 벤치마크에서 SWE-Pruner Pro는 토큰 절감과 품질 보존을 동시에 달성했다. Qwen3-Coder-Next의 SWE-QA-Pro 셀에서 SWE-Pruner Pro는 39.4% 토큰 절감을 기록했고 SWE-QA 계열 전체에서 최대 39% 토큰 절감 수치를 보였다. Oolong 장문 집계 작업에서는 MiMo-V2-Flash에서 30%에 이르는 토큰 절감과 함께 정확도 +2.2 포인트 향상이 관찰되어 절감이 품질 저하로 이어지지 않음을 확인했다. 코드 수정 벤치마크인 SWE-Bench Verified에서는 백본에 따라 결과의 비대칭성이 나타났다. MiMo-V2-Flash에서는 모든 프루너가 resolve rate를 개선했고 SWE-Pruner Pro는 +3.8%의 추가 개선을 보이며 다른 방법 대비 입력 토큰 증가폭이 상대적으로 작았다. 반대로 Qwen3-Coder-Next에서는 대부분의 프루너가 resolve 수를 감소시켰고, SWE-Pruner Pro는 6건만 손실하면서 입력 토큰을 최대 -13.5% 줄여 트레이드오프 측면에서 가장 유리한 프로파일을 보였다. 설계 구성 요소의 기여도를 평가한 소규모 소거(ablation)에서는 샘플별 균형 포컬 손실과 길이 인식 임베딩이 성능 향상을 주도했다. 손실을 BCE로 대체하면 judge 점수와 F1이 하락했고, 길이 임베딩은 F1을 크게 바꾸지 않으면서도 GPT-5.4-mini judge 점수에서 6.86에서 7.08로 향상시켜 길이-의존적 오류 재분배가 실제 평가에서 유의미함을 보였다.

전체 아키텍처는 백본을 고정하고 마지막-레이어 은닉 벡터 H mathbb{R}^{B\times T\times D}를 입력으로 받는 구조다. 길이 인식 임베딩 S mathbb{R}^{B\times D}를 선형적으로 버킷 색인화해 T에 브로드캐스트 더한 뒤 합산하여 H'를 얻고, H'는 FFN Pruning Head로 들어가 토큰 수준 로그it z를 출력한다. 헤드 구조는 LayerNorm, 두 개의 Linear-GELU-Dropout 블록(숨김 폭 d, backbone hidden size와 동일), 마지막 Linear 출력으로 구성되어 있다. 핵심 메커니즘의 수학적 기반은 길이 임베딩의 덧셈과 토큰별 분류기의 비선형성에 있다. 길이 임베딩은 수식 tilde{h}i = h_i + e(N)로 표현되며, 이는 동일한 h_i가 응답 길이 N에 따라 다른 판별 경계를 갖게 만드는 역할을 한다. 토큰별 확률은 z_i=f(theta)(tilde{h}i), p_i=(z_i)로 계산되고, 라인 결정은 각 라인 내 바이너리화된 p_i 임계값(τ=0.5)에 대한 다수결로 이루어진다. 손실 함수는 토큰별 포컬 형태 L_i^{tok}=(1-p{t,i})^{}\cdot BCE(p_i,y_i), =2의 형태로 정의되며 샘플별로 keep과 prune 토큰의 평균을 따로 구해 두 값을 1/2 가중치로 합친다. 수식 설명 절차에 따라 기존 수식과 변수는 다음과 같다: 1) 원문 수식 L_i^{tok}=(1-p_{t,i})^{}\cdot BCE(p_i,y_i), =2. 2) 변수 의미: p_i는 모델의 토큰별 예측 확률, y_i는 토큰 레이블(0 또는 1), p_{t,i}는 정답에 따라 p_i 또는 1-p_i로 변환한 값, BCE는 이진 교차엔트로피를 뜻한다. 3) 계산 흐름: 각 토큰에 대해 p_i를 계산하고 p_{t,i}를 정하면 포컬 스케일 (1-p_{t,i})^{}를 곱한 BCE를 구한다. 4) 수치 예시: y_i=1이고 p_i=0.9이면 p_{t,i}=0.9이고 (1-0.9)^2=0.01이므로 토큰 손실이 0.01 imes BCE(0.9,1)작게 된다. 5) 해석: 확신 있는 예측은 작은 스케일링 값을 받아 손실 기여가 줄고, 불확실한 예측은 더 큰 스케일링을 받아 학습 신호가 강화된다. 그 결과 샘플별 소수 클래스가 희석되지 않고 학습에 충분한 기여를 하게 된다. 구현 측면에서 백본은 완전히 고정되어 있으며 22,609개 샘플로 헤드만 학습했다. 라인 레이블은 LLM 어노테이터에서 얻은 것을 토큰 레벨로 확장해 사용했고 불확실한 행은 긍정 레이블로 유지하는 라벨링 프로토콜을 따랐다. 추론 측면에서는 in-server 통합으로 은닉 상태 전송 비용을 제거했고 재포워드는 pruned response에 대해서만 발생하므로 긴 응답에서 전체 비용 이득이 클 때 순이익이 발생한다.

평가 데이터와 실험은 공개된(open-weight) 백본을 대상으로 수행되었고, 백본이 내부 은닉 상태를 외부에 노출할 수 있어야 SWE-Pruner Pro를 그대로 적용할 수 있다. 언어 범위는 주로 Python 중심의 코딩 작업과 Oolong의 자연어 장문 집계로 제한되어 있어 다른 프로그래밍 언어나 도메인 전반에 대한 일반성은 추가 검증이 필요하다. 공격적 프루닝은 벤치마크에서 포착되지 않은 품질 저하를 유발할 수 있으므로 안전성이나 핵심 로직 손실이 민감한 환경에서는 배포 전 별도 검증이 요구된다.