GradCuit: 테스트 시점에 LLM의 내부 추론을 최적화하다

도입

대규모 언어모델의 추론 성능을 높이는 일반적인 방법은 더 나은 프롬프트를 쓰거나, 여러 답을 샘플링하거나, 후보를 재정렬하거나, 추가 학습을 수행하는 것입니다. GradCuit은 다른 방향을 제시합니다. 모델 가중치는 그대로 두고, 테스트 시점에 현재 입력에 대응하는 연속 잠재 상태를 최적화합니다. 즉, 답을 반복 생성하는 데 그치지 않고, 해당 문제를 풀기 위한 모델 내부의 추론 경로를 조정하려는 접근입니다.

핵심 아이디어는 프롬프트의 hidden representation과 생성될 continuation 사이, 즉 선택된 Transformer 중간층에 최적화 가능한 잠재 상태를 넣는 것입니다. 인과적 self-attention과 이후 Transformer 블록을 통해, continuation의 각 token 로그확률은 앞선 잠재 상태로 이어지는 미분 가능한 경로를 갖습니다. 따라서 최종 출력에 대한 보상이나 피드백을 잠재 상태로 더 직접적으로 되돌릴 수 있습니다.

핵심 포인트

• 파라미터 업데이트 없는 테스트 시점 적응: GradCuit은 기반 LLM을 미세조정하지 않습니다. 입력별 잠재 상태만 추론 중 최적화하므로, 일종의 추론 단계 적응 메커니즘에 가깝습니다.
• 더 직접적인 credit assignment: 기존 잠재 추론 방법은 잠재 상태와 추론 과정을 디코딩된 token을 통해 연결하는 경우가 많아, 전체 sequence의 성과를 내부 상태에 어떻게 배분할지 불명확했습니다. GradCuit은 전체 continuation에서 잠재 상태로 gradient가 흐르는 구조를 만듭니다.
• 여러 설정에서의 성능 향상: 논문은 5개 instruction-tuned backbone, 3개 추론 벤치마크, 2개 답변 형식에서 평균 정확도 64.5%를 보고했습니다. 이는 chain-of-thought prompting보다 6.6%포인트 높고, 연구 내 최강 비교 방법보다 2.4포인트 높은 수치입니다.
• 학습률 변화에 강한 안정성: 7개 learning rate 설정에서 GradCuit은 LatentSeek를 지속적으로 앞섰고, 정확도 표준편차도 1.53에서 0.82로 낮췄다고 합니다. 또한 random-walk 변형도 LatentSeek와 경쟁적인 성능을 보였습니다.
• 해석 가능성 단서 제공: token 수준 gradient attribution 결과, 잠재 상태의 영향은 추론을 이어주는 connector token에 집중되는 경향을 보였습니다. layer 분석에서는 Transformer의 초기~중간층이 최적화 공간으로 효과적이라는 결과도 제시됐습니다.

의미와 영향

GradCuit의 의미는 단순히 정확도가 조금 올랐다는 데 있지 않습니다. 테스트 시점 확장을 더 많은 샘플 생성이나 후보 재정렬을 넘어, 내부 추론 상태 최적화라는 축으로 넓혔습니다. 복잡한 추론 작업에서는 가중치를 바꾸지 않고도 특정 입력에 맞춰 추론 방식을 조정할 수 있다는 가능성을 보여줍니다.

물론 실제 적용에는 고려할 점이 있습니다. 테스트 시점 최적화는 보통 추가 계산을 요구하므로, 정확도 향상과 지연 시간, 비용 사이의 균형이 중요합니다. 또한 보상 또는 피드백 신호의 품질이 낮거나 불안정하면, 잠재 상태 업데이트가 항상 도움이 된다고 보장하기 어렵습니다.

연구 관점에서 GradCuit은 성능, 견고성, 해석 가능성을 하나의 구조 안에서 다룬다는 점이 흥미롭습니다. 향후 추가 추론 비용을 줄이고 더 개방적인 실제 과제에서 검증이 이뤄진다면, 테스트 시점 잠재 최적화는 LLM 추론 강화를 위한 중요한 방향이 될 수 있습니다.

출처: Hugging Face Daily Papers