Quantifying Logical Consistency in Transformers via Query-Key Alignment

대규모 언어 모델은 인간의 대화를 흉내 내고, 시를 쓰고, 복잡한 텍스트를 요약하는 데 놀라울 정도로 능숙해졌습니다. 이 모델들은 방대한 양의 학습 데이터를 바탕으로 문장에서 다음에 올 단어를 예측함으로써 이를 달성합니다. 그러나 다단계 논리 퍼즐을 풀라고 요청받으면, 이 시스템들은 종-종 비틀거립니다. 그들은 매우 설득력 있게 들리는 추론 과정을 생성할 수도 있지만, 결과적으로 완전히 틀린 결론에 도달하거나 프롬프트에 포함된 무관한 세부 사항에 의해 주의가 흐트러질 수 있습니다. 연구자들은 모델에게 "단계별로 생각하라"고 요청함으로써 모델의 성능을 끌어올리는 방법들을 찾아냈지만, 근본적인 문제는 여전히 남아 있습니다. 모델 자체에는 자신의 추론이 실제로 일관성이 있는지 확인할 수 있는 내장된 방법이 없다는 점입니다. 모델은 논리적인 경로처럼 보이는 것을 만들어낼 수는 있지만, 그 단계들이 실제로 제대로 연결되는지 진정으로 이해하고 있는지는 알 수 없습니다.

이러한 불확실성은 인공지능의 사각지대를 만듭니다. 우리는 모델이 내놓는 최종 답변은 볼 수 있지만, 모델이 실제로 논리의 규칙을 따르고 있는지 아니면 단순히 추측하고 있는지를 나타내는 내부 신호를 쉽게 볼 수는 없습니다. 이를 해결하기 위해 스콜코보 과학 기술 연구소(Skolkovo Institute of Science and Technology)와 다른 기관의 연구진은 모델이 추론하는 과정 중에 모델의 "마음"을 들여다볼 수 있는 새로운 방법을 개발했습니다. 모델이 답변을 마칠 때까지 기다렸다가 그것이 맞는지 확인하는 대신, 그들은 모델이 생각하는 동안 만드는 구체적인 내부 연결을 관찰합니다. 그들은 모델의 아키텍처 내 특정 부분들이 신뢰할 수 있는 체크포인트 역할을 하여, 최종 답변이 입 밖으로 나오기도 훨씬 전에 논리적 단계가 유효한지를 신호로 알려준다는 사실을 발견했습니다.

연구진은 이러한 모델들이 정보를 어떻게 처리하는지에 초점을 맞추었습니다. 대규모 언어 모델 내부에서는 정보가 수학적 연산의 층(layers)을 통해 흐릅니다. 이 과정의 핵심에는 '어텐션 헤드(attention heads)'가 있는데, 이는 매 순간 입력의 어느 부분에 집중할지를 결정하는 작은 스위치와 같은 역할을 합니다. 연구팀은 이 헤드들 내에서 두 가지 특정 신호 사이의 정렬(alignment)을 측정할 수 있다는 것을 발견했습니다. 하나는 테스트 중인 명제를 나타내는 신호이고, 다른 하나는 잠재적인 답을 나타내는 신호입니다. 이 두 신호가 얼마나 잘 일치하는지를 기반으로 간단한 점수를 계산함으로써, 연구진은 모델이 전제와 결론 사이의 논리적 연결을 인식했는지 여부를 판단할 수 있었습니다.

이 아이디어를 테스트하기 위해 연구진은 다양한 논리적 추론 과제를 대상으로 실험을 진행했습니다. 그들은 롬푸스(rompus)와 좀푸스(jompus)에 관한 허구적인 규칙을 바탕으로 폴리(Polly)라는 캐릭터가 불투명한지 결정하는 것과 같이, 규칙이 명확하고 답이 확정적으로 참 또는 거짓인 합성 데이터셋을 사용했습니다. 또한 실세계의 언어와 다단계 연역이 포함된 더 복잡한 데이터셋에서도 이 방법을 테스트했습니다. 이 테스트에서 그들은 모델의 표준적인 방식, 즉 단순히 확률이 가장 높은 옵션을 선택하는 방식과 연구진의 새로운 내부 점수 산출 방식을 비교했습니다.

결과는 내부 점수 산출 방식이 모델의 자체적인 최종 추측보다 훨씬 더 신뢰할 수 있음을 보여주었습니다. 모델이 방해되는 무관한 정보가 제시된 시나리오에서, 표준 모델은 종종 혼란을 느껴 틀린 답을 내놓았습니다. 그러나 연구진이 추적하던 특정 내부 신호들은 안정적으로 유지되었으며, 모델의 최종 출력이 틀렸을 때조차도 유효한 논리적 경로를 정확하게 식별해 냈습니다. 이는 모델이 실제로 그 깊은 층 안에 올바른 추론 능력을 갖추고 있지만, 그 능력이 후속 처리 단계나 방해되는 맥락의 소음에 의해 가려질 수 있음을 시사합니다.

연구팀은 15억 개의 파라미터를 가진 매우 작은 모델부터 700억 개의 파라미터를 가진 거대 모델에 이르기까지 다양한 모델을 대상으로 이 접근 방식을 테스트했습니다. 모든 모델에 걸쳐, 특정 어텐션 헤드들이 일관되게 '검증 앵커(verification anchors)' 역할을 한다는 것을 발견했습니다. 이 헤드들은 추론에 필요한 단계가 얼마나 많든, 혹은 얼마나 많은 방해 요소가 있든 상관없이 유효한 추론과 무효한 추론을 높은 정확도로 구분해 낼 수 있었습니다. 어떤 경우에는 내부 점수가 90% 이상의 정확도를 보인 반면, 모델의 최종 답변은 정답률이 60% 미만인 경우도 있었습니다. 이는 모델이 내부적으로는 정답을 '알고' 있으면서도 최종 출력으로는 이를 제대로 표현하지 못하는 경우가 많음을 나타냅니다.

가장 중요한 발견 중 하나는 이 방법이 모델을 변경하거나 재학습시킬 필요가 없다는 점입니다. 연구진은 단순히 모델을 한 번 실행하여 내부 신호를 살펴보고, 모델 자신의 예측보다 더 정확한 점수를 도출할 수 있었습니다. 이는 모델의 일부를 비활성화하여 어떤 일이 일어나는지 관찰해야 했던 기존의 무거운 계산 비용이 드는 방법들과 달리, 가볍고 효율적인 방식으로 논리적 일관성을 평가할 수 있는 길을 열어줍니다. 이 연구는 이러한 특정 내부 구성 요소를 식별함으로써, 우리가 인공지능이 복잡한 관계를 어떻게 처리하는지에 대한 더 명확한 창을 가질 수 있으며, 이를 통해 단순히 언어에 유창한 것을 넘어 추론에서도 신뢰할 수 있는 시스템에 더 가까워질 수 있음을 시사합니다.

연구진은 이 방법이 특정 내부 신호를 신뢰할 수 있도록 하는 적절한 데이터를 충분히 확보했을 때 가장 잘 작동한다고 언급했습니다. 또한 이 접근 방식이 이 특정 부분들이 논리를 담당하는 유일한 부분임을 의미하는 것도 아니며, 인공지능의 모든 추론 문제를 해결하는 것도 아니라는 점을 인정했습니다. 그러나 이 작업은 논리가 어디에서 유지되고 어디에서 무너지는지를 볼 수 있는 구체적이고 측정 가능한 방법을 제공하며, 더 투명하고 신뢰할 수 있는 언어 모델을 구축하기 위한 새로운 도구를 제공합니다. 내부 신호의 정렬에 집중함으로써, 이 연구는 건전한 추론 능력이 모델 내부에 존재하며, 단지 적절하게 읽히기를 기다리고 있을 뿐이라는 사실을 밝혀냈습니다.