StateBridge: Training-free Hidden-state Alignment for Latent Communication in LLM Multi-Agent Systems

기술 요약: StateBridge

문제 정의

대규모 언어 모델(LLM) 기반의 멀티 에이전트(Multi-Agent, MA) 시스템은 현재 에이전트 간 통신을 위해 자연어(이산적 토큰)에 의존하고 있다. 이러한 방식은 "이산적 병목 현상(discrete bottleneck)"을 유발한다. 즉, 송신자는 자신의 연속적인 내부 은닉 상태(hidden states)를 이산적 토큰으로 압축해야 하며, 이 과정에서 토큰의 정체성만으로는 포착할 수 없는 정보를 버리게 된다. 수신자가 이 토큰들을 임베딩 레이어를 통해 다시 벡터로 매핑할 때, 수신되는 것은 송신자의 원래 연속적 은닉 상태가 아니라 오직 토큰의 정체성뿐이다. 이는 정보 손실과 협업 능력 저하를 초래한다.

원시 은닉 상태(raw hidden states)를 직접 전송하는 것이 잠재적인 대안이 될 수 있으나, 이는 기하학적 불일치(geometric mismatch) 문제에 직면한다. 사전 학습된 LLM은 입력 레이어에서 토큰 임베딩을 처리하도록 학습되어 있는 반면, 디코더의 은닉 상태는 표현 공간(representation space)의 다른 영역을 점유한다. 은닉 상태를 직접 전달하면 수신자가 자신의 입력 임베딩 공간이 가진 특정한 기하학적 정렬을 이해하지 못하므로 의미론적 불일치(semantic mismatch)가 발생한다. 기존 솔루션들은 내부 상태를 레이어별로 주입하거나(이식성을 제한함), 특정 모델에 맞춘 프로젝터 학습을 요구한다(일반성을 제한함).

방법론: StateBridge

저자들은 StateBridge를 제안한다. 이는 폐쇄형 직교 변환(closed-form orthogonal transformation)을 통해 송신자의 최종 레이어 은닉 상태를 수신자의 입력 임베딩 공간에 정렬시키는 학습이 필요 없는(training-free) 통신 인터페이스이다. 이 방법은 세 단계로 작동한다:

1. 메시지 추출 (Message Extraction):
   송신 에이전트가 메시지를 생성한다. 전체 시퀀스를 전송하는 대신, 시스템은 생성된 토큰의 마지막 $K$개 은합 상태($S$)를 추출한다. 또한, 공유된 어휘 임베딩 행렬로부터 대응하는 참조 토큰 임베딩($R$)을 가져온다. 이 참조 임베딩은 정렬의 대상(target) 역할을 한다.

2. 정렬 인터페이스 (Alignment Interface):
   StateBridge의 핵심은 $S$를 수신자와 호환 가능한 정렬된 프리픽스(aligned prefix) $\bar{S}$로 변환하는 3단계 변환 과정이다:

   • 직교 프로크루스테스 정렬 (Orthogonal Procrustes Alignment): 이 방법은 $S$와 $R$을 $\mathbb{R}^d$ 공간의 점 구름(point clouds)으로 취급한다. 먼저 두 집합의 전역 오프셋을 제거하고 주 방향을 균등화하기 위해 중심화(centering) 및 백색화(whitening)를 수행한다. 그 후, 송신자의 정제된 상태를 백색화된 참조 임베딩으로 매핑하면서 쌍별 거리(pairwise distances)와 각도를 보존하는 최적의 회전 행렬 $Q^*$를 찾기 위해 직교 프로크루스테스 문제를 해결한다. 이를 통해 송신자 표현의 기하학적 구조(의미적 유사성)가 유지된다.
   • 노름 교정 (Norm Calibration): 최종 레이어의 은닉 상태는 일반적으로 입력 토큰 임베딩보다 노름(norm)이 훨씬 크다. StateBridge는 정렬된 각 벡터를 어휘 임베딩의 평균 $\ell_2$ 노름에 맞춰 재조정한다. 이는 프리픽스가 단순히 크기(magnitude) 때문에 수신자의 어텐션 점수를 독점하는 것을 방지한다.
   • 어휘 앵커링 (Vocabulary Anchoring): 정렬된 벡터가 사전 학습 중에 접했던 임베딩 공간 내의 영역에 위치하도록 하기 위해, 각 교정된 벡터를 코사인 유사도에 기반하여 가장 가까운 어휘 임베딩 쪽으로 미세하게 이동시킨다. 이 단계는 연속성을 유지하면서도 표현을 수신자가 기대하는 입력 분포에 고정(anchor)시킨다.

3. 프리픽스 주입 (Prefix Injection):
   결과물인 정렬된 프리픽스 $\bar{S}$를 수신자의 프롬프트 임베딩 앞에 붙인다(prepend). 수신자는 별도의 구조적 변경이나 토큰화 없이 표준 언어 모델링 레이어를 사용하여 이 연속적인 프리픽스를 처리한다.

주요 기여

1. 표현 불일치 해결: 본 논문은 폐쇄형 정렬만으로도 최종 레이어 은닉 상태와 입력 임베딩 사이의 기하학적 불일치를 해결할 수 있음을 입증하였으며, 이를 통해 학습이나 구조 변경 없이 효과적인 잠재 통신(latent communication)을 가능하게 한다.
2. StateBridge 프레임워크: 저자들은 프로크루스테스 정렬, 노름 교정, 어휘 앵커링을 결합한 구체적인 파라미터 프리(parameter-free) 인터페이스를 제안한다. 이 접근 방식은 송신자 상태의 쌍별 기하학적 관계를 보존하는 동시에 수신자의 입력 공간과 호환성을 보장한다.
3. 경험적 우월성: Qwen3 및 OLMo3 제품군의 4개 모델과 8개 벤치마크(수학, QA, 코드) 전반에 걸쳐, StateBridge는 텍스트 기반 통신 및 KV-캐시 전송 베이스라인보다 일관되게 우수한 성능을 보인다. 이는 26개의 모델-태스크 쌍 중 22개에서 최고 또는 공동 최고 점수를 달 기록했다.

결과

• 성능: StateBridge는 모든 모델 설정에서 가장 좋은 베이스라인 대비 평균 2.4 ~ 2.9 퍼센트 포인트의 향상을 달성했다.
• 일반화 능력: 이 방법은 서로 다른 모델 아키텍처 전반에서 견고한 성능을 보였다. 특히 OLMo3-7B-Think 모델의 경우, KV-캐시 전송(LatentMAS)은 평균 55.1%로 텍스트(73.9%)보다 낮은 성능을 보였으나, StateBridge는 76.7%를 달성했다. 이는 StateBridge의 입력 레이어 전용 동작이 모든 트랜스포머 레이어에 걸쳐 상태를 주입하는 방식보다 아키텍처 차이에 더 강건하다는 것을 시사한다.
• 태스크 민감도: 성과는 심층 추론이 필요한 까다로운 벤치마크(예: MedQA, GPQA, AIME, 코드 생성)에서 가장 두드러졌는데, 이는 텍스트 통신이 유용한 연속 정보를 버릴 가능성이 가장 높은 영역이다.
• 절제 연구 (Ablation Studies):
  • 프로크루스테스 정렬을 리지 회귀(ridge regression, 비직교 선형 맵)로 대체했을 때 평균 성능이 82.4%에서 74.9%로 크게 하락하여, 쌍별 기하학적 구조를 보존하는 것이 매우 중요함을 확인했다.
  • 노름 교정이나 어휘 앵커링을 제거하는 것 또한 성능 저하를 일으켰으며, 이는 기하학적 정렬만으로는 규모(scale) 및 분포적 호환성을 확보하기에 불충분함을 나타낸다.
• 사례 연구: 재구성(reconstruction) 태스크를 통해, 정렬된 프리픽스가 가시적인 접미사(suffix) 토큰에는 존재하지 않는 의미적 정보(예: "koilonychia"와 같은 특정 의학 용어)를 포함하고 있음을 입증하였으며, 이는 프리픽스가 이산적 텍스트 출력 이상의 정보를 운반한다는 것을 증명한다.

의의 및 주장

본 논문은 잠재 통신의 주요 과제가 단순히 더 풍부한 상태를 전송하는 것이 아니라, 그 상태를 다음 에이전트가 읽을 수 있게(readable) 만드는 것이라고 주장한다. StateBridge는 전송된 상태가 정보적이면서도 수신자의 입력 공간과 기하학적으로 호환되도록 보장함으로써 이 문제를 해결한다.

저자들은 자신들의 방식이 학습이 필요 없고(training-free), 구조적 변경이 필요 없다는 점을 강조하며, 이를 통해 상용 LLM에도 적용 가능하다는 점을 내세운다. 또한 통신 인터페이스 설계가 동종(homogeneous) MA 시스템에서 매우 중요하지만 종종 간과되는 요소라고 주장한다. 현재의 연구는 공유된 모델 가중치와 고정된 프리픽스 길이를 가정하고 있지만, 결과는 정렬 기반 인터페이스가 특히 이산적 토큰화가 상당한 정보 손실을 일으키는 작업에서 다중 에이전트 협업을 크게 향상시킬 수 있음을 시사한다.