[Paper] Tool Calling은 선형적으로 읽을 수 있고 언어 모델에서 제어 가능합니다

Source: arXiv - 2605.07990v1

개요

이 논문은 최신 지시‑튜닝된 언어 모델이 외부 도구를 호출하는 선택이 모델의 은닉 상태 내부에서 선형적으로 읽히고 조정될 수 있다는 놀라울 정도로 단순한 특성을 밝혀냅니다. 내부 활성화를 탐색하고 미세 조정함으로써, 저자들은 모델의 도구 선택을 거의 완벽한 정확도로 뒤바꿀 수 있으며, 도구 호출 오류가 발생할 직전인 경우도 예측할 수 있습니다. 이 통찰은 AI 에이전트를 보다 신뢰할 수 있고 제어 가능하게 만드는 새로운 길을 열어줍니다.

주요 기여

• 도구 정체성의 선형 판독 – 은닉 활성화에 대한 단일 선형 투영을 사용하면, 튜닝되지 않은 기본 모델에서도 69‑82 % 정확도로 모델이 사용하려는 도구를 복원할 수 있습니다.
• 도구 선택의 선형 조정 – 두 도구의 활성화 중심 사이의 평균 차이 벡터를 모델의 은닉 상태에 추가하면, 77‑100 %의 경우(파라미터가 4 B 이상인 모델에서는 ≥ 93 %) 목표 도구를 선택하도록 강제할 수 있습니다.
• 오류 예측 신호 – top‑1과 top‑2 도구 로짓 간의 차이가 실패를 예측합니다: 차이가 작을수록 잘못된 도구 호출이 발생할 확률이 14‑21배 높아집니다.
• 효과의 위치 지정 – 활성화 패칭을 통해 중간 및 후반 레이어의 몇몇 어텐션 헤드와 출력 레이어 행의 단일 방향이 도구 토큰을 구동한다는 것을 확인했으며, 이는 현상이 단순히 주제 전환에 국한되지 않음을 보여줍니다.
• 모델 및 도구 간 일관성 – 선형 판독은 270 M에서 27 B 파라미터에 이르는 12개의 인스트럭션‑튜닝 모델(Gemma 3, Qwen 3/2.5, Llama 3.1)과 14개의 항공 분야 도구 전반에 걸쳐 작동합니다.
• 사전학습 vs. 인스트럭션‑튜닝에 대한 통찰 – 기본 모델은 이미 생성 전에 올바른 도구를 인코딩하고 있으며(높은 코사인 유사도), 인스트럭션 튜닝은 해당 잠재 표현을 실제 출력 토큰과 정렬시킵니다.

방법론

1. Tool‑calling benchmark – 여러 JSON‑schema 도구(예: “search flight”, “book seat”) 중 하나로 각 질의에 답변할 수 있는 고정 메뉴, 단일 턴 프롬프트 세트를 구축했습니다.
2. Probing – 모델이 프롬프트를 처리하는 동안 모든 레이어의 숨겨진 상태를 수집했습니다. 이 활성화를 사용해 의도된 도구를 예측하는 간단한 선형 분류기(로지스틱 회귀)를 학습했습니다.
3. Steering via activation injection – 각 도구에 대한 평균 활성화 벡터(“도구 중심”)를 계산했습니다. 도구 A에서 도구 B로 전환하려면 선택한 레이어의 숨겨진 상태에 차이 Δ = mean_B – mean_A 를 더하고 모델이 생성 과정을 계속하도록 했습니다.
4. Activation patching – 개별 어텐션 헤드 출력을 “올바른” 실행에서 나온 값으로 교체하고 도구 선택에 미치는 영향을 측정했습니다.
5. Error‑gap analysis – 상위 두 도구 토큰 간의 로짓 차이를 기록했습니다; 차이가 작은 질의는 고위험으로 표시했습니다.
6. Cross‑model validation – 동일한 프로브와 스티어링 벡터를 크기와 아키텍처가 다른 12개의 모델에 적용해 견고성을 테스트했습니다.

모든 단계는 표준 트랜스포머 내부(숨겨진 상태, 어텐션 헤드, 출력 로짓)를 활용하며, 그래디언트 업데이트 없이 순전파와 간단한 벡터 연산만 필요합니다.

결과 및 발견

• 단일 방향 제어: 목표 도구의 첫 번째 토큰에 해당하는 출력‑레이어 행과 정렬된 단위 벡터를 주입하면 이미 93 % 이상의 스티어링을 달성할 수 있어, 효과의 대부분이 한 방향에 집중되어 있음을 확인했다.
• 어텐션 헤드 핫스팟: 5‑8개의 중‑후반 레이어 헤드만 패치하면 스티어링 효과를 재현할 수 있어, “도구 선택 회로”가 비교적 컴팩트하게 존재함을 시사한다.
• 오류 예측: top‑1/top‑2 로짓 차이가 최하 사분위에 해당하는 쿼리는 최상 사분위에 비해 잘못된 도구 호출이 14‑21배 더 자주 발생한다.
• 베이스 vs. 튜닝: 튜닝되지 않은 베이스 모델은 올바른 도구를 (높은 코사인 유사도) 인코딩하지만 실제로는 거의 출력하지 않는다 (생성 정확도 2‑10 %). 인스트럭션 튜닝은 잠재 표현을 출력 토큰과 정렬시켜 실제 도구 사용을 크게 향상시킨다.

전체적으로, 이번 연구는 도구 선택이 모델 내부에 명시적으로 표현되어 있으며 최소한의 오버헤드로 읽고, 편집하고, 모니터링할 수 있음을 보여준다.

실용적 함의

• Debug‑friendly agents – 개발자는 실행 전에 의도된 도구를 확인하기 위해 숨겨진 상태를 읽는 가벼운 “watchdog”을 추가할 수 있으며, 이를 통해 불일치를 조기에 포착하고 비용이 많이 드는 실수를 방지할 수 있습니다 (예: 잘못된 수신자에게 이메일을 보내는 경우).
• Runtime steering – 적절한 Δ 벡터를 주입함으로써 시스템은 모델을 다시 프롬프트하거나 파인튜닝하지 않고도 요청을 더 안전하거나 적절한 도구로 동적으로 재라우팅할 수 있습니다. 컴플라이언스(프라이버시 보호 도구 강제 적용)나 다양한 도구 구현을 A/B 테스트할 때 유용합니다.
• Safety layers – logit‑gap 메트릭은 도구 호출에 대한 저비용 실시간 신뢰도 점수를 제공하여 조건부 폴백을 가능하게 합니다 (예: gap이 낮을 때 사용자에게 확인을 요청).
• Model‑agnostic tooling – 이 현상이 다양한 아키텍처와 규모에 걸쳐 나타나므로, 라이브러리는 (read_tool, steer_tool, tool_confidence)와 같은 일반 API를 제공하여 최신 LLM 백엔드와 연동할 수 있습니다.
• Efficient fine‑tuning – 도구 사용을 위한 전체 모델 RLHF 대신, 개발자는 식별된 소수의 어텐션 헤드만 파인튜닝하거나 도구‑centroid 벡터를 출력 토큰에 직접 매핑하는 선형 어댑터를 추가하여 계산량과 데이터를 절감할 수 있습니다.
• Pre‑training diagnostics – 기본 모델이 이미 도구 정체성을 인코딩하고 있다는 사실은 프리트레이닝 데이터 품질(도구와 유사한 패턴의 존재)이 다운스트림 도구 신뢰성에 영향을 미친다는 것을 시사하며, 향후 LLM을 위한 데이터셋 큐레이션에 방향을 제시합니다.

제한 사항 및 향후 연구

• 단일 턴, 고정 메뉴 설정 – 실험은 정적인 도구 목록을 가진 원샷 프롬프트에 초점을 맞추었습니다. 다중 턴 대화와 동적 도구 탐지는 여전히 취약하며 더 깊은 연구가 필요합니다.
• 도구 범위 – 항공 분야 JSON 도구만 평가했으며, 결과가 더 복잡하거나 계층적인 도구 모음(예: 코드 실행, 데이터베이스 쿼리)으로 어떻게 일반화될지는 불분명합니다.
• 스티어링 부작용 – 도구 이름은 깔끔하게 전환되지만, 하위 JSON 인자는 프롬프트에 스키마가 포함되어 있기 때문에 새로운 스키마에 자동으로 맞춰집니다. 보다 모호한 경우에는 잘못된 인자가 생성될 수 있습니다.
• 해석 가능성 깊이 – 식별된 어텐션 헤드는 유망한 출발점이지만, “도구 회로”가 instruction tuning과 어떻게 상호작용하는지에 대한 완전한 메커니즘 모델은 아직 없습니다.
• 적대적 프롬프트에 대한 견고성 – 선형 스티어링 기법은 모델을 악의적인 도구 사용으로 강제할 위험이 있으며, 방어 및 탐지 메커니즘은 아직 연구가 필요한 분야입니다.

향후 연구에서는 프로빙/스티어링 프레임워크를 다중 턴 에이전트로 확장하고, 자동화된 도구 선택 교정 루프를 탐색하며, 사전 학습 중에 유사한 선형 제어 신호를 삽입하여 본질적으로 더 안전한 도구 호출 모델을 만드는 방법을 조사할 수 있습니다.

저자

• Zekun Wu
• Ze Wang
• Seonglae Cho
• Yufei Yang
• Adriano Koshiyama
• Sahan Bulathwela
• Maria Perez-Ortiz

논문 정보

• arXiv ID: 2605.07990v1
• 분류: cs.CL, cs.AI, cs.LG, cs.SE
• 출판일: 2026년 5월 8일
• PDF: PDF 다운로드