에너지 기반 트랜스포머(EBT) 해설: 추론을 경사하강법으로 푸는 새 학습 패러다임

핵심 메시지

• 에너지 기반 트랜스포머(EBT)는 입력과 후보 답의 ‘궁합’을 하나의 숫자로 매기는 에너지 함수를 학습하고, 추론 시점에는 그 숫자를 낮추는 최적화를 돌려 답을 만든다.
• 덕분에 한 번의 순전파로 답이 끝나지 않고, 계산을 더 넣을수록 답을 더 다듬을 수 있는 구조가 된다.
• 학습할 때부터 이 추론 최적화 과정을 통째로 역전파에 포함시키며, 이 때문에 2차 미분이 필요하지만 헤시안-벡터 곱으로 모델 크기에 선형적으로 처리한다.
• 에너지 지형이 울퉁불퉁해지지 않도록 리플레이 버퍼, 최적화 단계마다의 잡음 주입, 스텝 크기와 단계 수의 무작위화 같은 정규화 기법이 함께 쓰인다.
• 실험 규모는 아직 작고 고정 비용도 크지만, 데이터·배치·깊이를 키울 때의 기울기가 기존 트랜스포머보다 가팔라 대규모에서 역전될 가능성이 제기된다.

쉽게 이해하기

논문이 던지는 질문은 크다. 사람의 느리고 명시적인 사고, 이른바 ‘시스템 2 사고’를 지도 데이터 없이 순수한 비지도 학습만으로 만들어 낼 수 있는가다. 리뷰어는 이 표현이 다소 과한 철학이라고 보면서도, 논문이 말하는 ‘생각’이 사실은 매우 구체적인 정의를 갖는다는 점을 짚는다. 훈련된 모델을 추론 시점에 한 번보다 여러 번 쓰는 것, 그리고 그렇게 계산을 더 넣었을 때 답이 얼마나 좋아지는지가 곧 이 논문에서 말하는 사고이고 그 척도다.

핵심 장치는 에너지 함수다. 입력 x와 후보 답 y를 함께 받아 하나의 실수를 내놓는데, 둘이 잘 어울리면 낮은 값을, 어울리지 않으면 높은 값을 준다. 다음 토큰 예측이라면 x는 지금까지의 문맥이고 y는 다음 토큰에 대한 확률 분포다. 이미지 잡음 제거라면 x는 잡음이 낀 이미지, y는 복원된 이미지가 된다. 중요한 것은 학습되는 모델이 이 에너지 함수 하나뿐이라는 점이다. 생성자와 판별자를 따로 두는 적대적 생성 신경망과 달리, 여기서는 판정자 역할을 하는 함수가 곧 모델 전부다.

그렇다면 숫자 하나만 알려주는 함수로 어떻게 답을 만드는가. 답은 최적화다. 문맥 x를 고정한 채 무작위 분포에서 출발해, 에너지의 기울기를 후보 답 쪽으로 역전파해 조금씩 내려가는 경사하강을 반복한다. 단계를 많이 밟을수록 더 정교한 답을, 적게 밟으면 거친 답을 얻는다. 확산 모델이 떠오르는 방식이지만, 여기서는 학습된 판정 기준 자체를 상대로 최적화한다는 점이 다르다.

학습은 이 추론 절차를 염두에 두고 이뤄진다. 대조 학습 방식도 가능하지만 확장성이 떨어지고, 대신 최적화를 한 단계 밟아 나온 출력에 교차 엔트로피 같은 손실을 걸고 그 최적화 과정 전체를 관통해 역전파한다. 기울기 계산을 다시 미분해야 하므로 2차 미분이 필요하지만, 헤시안-벡터 곱을 이용하면 모델 크기에 선형으로 늘어나 폭발적으로 비싸지지는 않는다. 여기에 에너지 지형을 매끄럽게 만들기 위한 장치들이 더해진다. 리플레이 버퍼를 두고, 최적화 각 단계마다 잡음을 섞고, 스텝 크기와 단계 수를 무작위로 바꾼다. 학습 때 밟는 경로를 일부러 넓게 번지게 해 두면 처음 보는 데이터에서도 길을 잃지 않는다는 발상이다.

에너지 기반 모델은 전통적으로 병렬화·안정성·확장성에서 약했고, 트랜스포머는 바로 그 세 가지에 강하다. 두 계열을 결합한 것이 EBT이며, 특히 디코더 전용 삼각 어텐션에서 여러 번의 추론 단계 사이에 정보가 새지 않게 만드는 일이 만만치 않은 공학적 과제였다고 소개된다. 실험은 두 축으로 나뉜다. 하나는 학습 확장성으로, 데이터양·배치 크기·모델 깊이를 키울 때 EBT가 처음에는 불리하게 출발하지만 기존 트랜스포머보다 빠르게 따라잡는다. 다른 하나는 ‘사고 확장성’으로, 순전파 횟수를 늘려도 변하지 않는 기존 트랜스포머와 달리 EBT는 횟수를 늘릴수록 성능이 올라가 결국 앞선다.

리뷰어는 기대와 유보를 함께 남긴다. 학습 연산량 기준으로 보면 EBT는 여전히 기존 트랜스포머보다 한참 오른쪽, 즉 훨씬 비싼 위치에 있다. 다만 기울기가 아주 조금 가파르기 때문에 추세가 이어진다면 언젠가 교차할 수 있다는 것이지, 지금 규모에서 증명된 이야기는 아니다. 또한 논문이 사고의 사슬이나 강화학습 기반 추론 모델과 자신을 대비시키는 구도에 대해서도, 그것들은 오히려 서로 배타적이지 않고 결합 가능한 축이라고 선을 긋는다.

주요 인사이트

• 이 논문에서 ‘생각’은 신비한 능력이 아니라 측정 가능한 양으로 정의된다. 순전파를 여러 번 돌렸을 때 답이 얼마나 좋아지는가가 그 척도이고, 그 정의를 받아들이면 확산 모델이나 순환 신경망도 같은 잣대 위에 놓인다.
• 손실 함수와 에너지 함수의 차이는 수식이 아니라 쓰이는 시점에 있다. 손실 함수는 학습할 때 쓰는 목표이고, 에너지 함수는 추론할 때 최적화하는 목표다. 이 구분이 EBT 전체 설계를 이해하는 열쇠다.
• 에너지 값 자체가 부가 정보를 준다. 문장 안에서 쉬운 토큰일수록 에너지가 낮게 나오는데, 이는 모델이 자기 예측의 확신도를 스스로 평가한다는 뜻이고, 토큰마다 계산량을 달리 배분하는 동적 추론의 실마리가 된다.
• 성능 향상의 대부분은 최적화 첫 단계에서 나온다. 이후 단계는 개선 폭이 급격히 줄어드는데, 이는 최적화에서 흔히 관찰되는 양상이기도 하다.
• 리뷰어가 가장 경계하는 지점은 외삽이다. 그래프의 기울기가 유리하다는 사실과 대규모에서 실제로 역전한다는 사실은 다르며, 현재 실험은 작은 규모에 머물러 있다.

자주 묻는 질문

관련 AI 소식