LG AI Research, TokenGT 그래프 트랜스포머 메커니즘 해석 성공

LG AI Research가 ICML 2026 Mechanistic Interpretability Workshop에서 그래프 트랜스포머 모델인 TokenGT에 메커니즘 해석 기법을 적용한 연구 결과를 발표했다. 대부분의 해석 연구가 LLM에 집중된 반면, 이 연구는 화학 및 생물학에서 활용되는 그래프 트랜스포머의 미개척 영역을 다룬다. 연구팀은 노드와 엣지 토큰 시퀀스로 그래프를 처리하는 바닐라 트랜스포머의 내부 계산 경로를 역추적했다.

공유된 계산 메커니즘

차수 계산, 사이클 감지, 최단 경로 거리 모델 분석 결과 초기 단계에서 공유되는 메커니즘이 발견됐다. 첫 번째 트랜스포머 레이어에서 노드 토큰은 일치하는 엔드포인트 식별자를 가진 엣지 토큰에 강하게 어텐션한다. 이 ID 매칭 동작은 내장된 그래프 연산 없이 인접 엣지를 복원하게 한다. 선형 프로브 확인 결과, 첫 어텐션 직후 노드 차수 정보가 정확히 식별되며 잔차 스트림의 쓰기 벡터 크기가 실제 차수에 비례한다.

작업별 적응 방식

초기 지역 구조 계산은 공유되지만, 특정 작업을 위해 기능을 조합하는 방식은 모델마다 다르다.

• 링 멤버십: 모델은 명시적인 사이클 추적 알고리즘을 구현하지 않는다. 대신 노드를 링의 일부로 분류한 후, 링에 속하지 않은 경우 예측을 뒤집기 위해 증거를 점진적으로 수집한다. 어블레이션 연구 결과, 두 번째 레이어 제거는 비링 예측에 주로 손상을 줬다.
• 최단 경로 거리: 모델은 더 정교한 메커니즘을 사용한다. 세 번째 레이어의 세 번째 어텐션 헤드는 BFS 트리의 부모와 일치하는 인접 노드에 집중한다. 이 어텐션 대상을 해당 BFS 부모로 대체해도 성능이 대부분 유지되어, 모델이 입력 그래프를 루트 지정된 BFS 트리로 인식함을 시사한다.

AI4Science의 향후 방향

이 연구는 아키텍처 편향 없이 트랜스포머가 그래프 계산을 수행하는 방식에 대한 초기 메커니즘적 설명을 제공한다. LG AI Research는 실제 분자 그래프로 연구를 확장해 모델이 화학적 사실과 규칙을 학습하는 방식을 이해할 계획이다. 목표는 해석 분석을 통해 AI 모델과 도메인 과학자 간 간극을 좁히고 과학적 지식 추출을 용이하게 하는 것이다.