최적화 주도형 AI, 범용 모델을 향한 새로운 길로 부상하다

일리노이 대학교 어바나-샴페인 캠퍼스와 버지니아 대학교 연구진이 추론 능력을 강화한 더 탄력적인 AI 시스템의 길을 열 수 있는 새로운 모델 아키텍처를 개발했다.

에너지 기반 트랜스포머(EBT)로 명명된 이 아키텍처는 복잡한 문제를 해결하기 위해 추론 시간 확장성을 자연스럽게 활용합니다. 기업에게는 이는 특수한 미세 조정 모델 없이도 새로운 시나리오에 적응할 수 있는 비용 효율적인 AI 애플리케이션을 의미합니다.

시스템 2 사고의 과제

심리학에서 인간의 인지 능력은 일반적으로 두 가지 모드로 구분됩니다: 빠르고 본능적인 시스템 1과 더 느리고 신중하며 분석적인 시스템 2입니다. 오늘날의 대규모 언어 모델(LLM)은 시스템 1 작업에 탁월하지만, AI 분야는 이제 더 복잡한 추론 문제를 처리하기 위해 시스템 2 사고 개발에 집중하고 있습니다.

추론 모델들은 어려운 질문에 대한 성능을 향상시키기 위해 다양한 추론 시간 확장 방법을 사용합니다. 일반적인 기법은 강화 학습(RL)으로, DeepSeek-R1이나 OpenAI의 "o-series" 같은 모델에서 사용되며, AI가 올바른 해결책에 도달할 때까지 추론 토큰을 생성하여 보상을 얻습니다. 또 다른 전략인 '베스트 오브 엔(best-of-n)'은 여러 가능한 답안을 생성하고 검증 시스템을 통해 최상의 답안을 선택하는 방식입니다.

그러나 이러한 접근법에는 주목할 만한 한계가 존재합니다. 수학이나 코딩처럼 검증하기 쉬운 좁은 범위의 문제에 주로 국한되며, 창의적 글쓰기 같은 다른 작업에서는 성능이 저하될 수 있습니다. 또한 최근 연구에 따르면 RL 기반 방법은 모델에게 새로운 추론 능력을 가르치는 것이 아니라, 이미 알고 있는 성공적인 추론 패턴을 재사용하도록 장려할 뿐일 수 있습니다. 이는 훈련 범위를 넘어 진정한 탐색이 필요한 문제 해결 능력을 제한합니다.

에너지 기반 모델(EBM)

이 새로운 아키텍처는 에너지 기반 모델(EBM)로 알려진 모델 클래스를 기반으로 하는 다른 경로를 취합니다. 핵심 개념은 간단합니다: 모델은 답을 직접 생성하기보다는 검증자 역할을 하는 "에너지 함수"를 학습합니다. 이 함수는 입력(프롬프트 등)과 후보 예측을 받아 값 또는 "에너지"를 할당합니다. 낮은 에너지 점수는 높은 호환성을 의미하며, 예측이 입력에 잘 부합함을 나타냅니다. 반대로 높은 에너지 점수는 부적합함을 시사합니다.

연구진은 논문에서 이 개념을 AI 추론에 적용하며, 개발자들이 "학습된 검증자에 대한 최적화 절차로서 사고를 고려해야 한다"고 제안합니다. 이 검증자는 입력과 후보 예측 간의 호환성(정규화되지 않은 확률)을 평가합니다. 이 과정은 무작위 예측으로 시작하여 에너지 점수를 최소화하고 가능한 해법을 탐색함으로써 점차 정교화되며, 궁극적으로 높은 호환성을 지닌 답에 수렴합니다. 이 방법은 해법을 검증하는 것이 처음부터 생성하는 것보다 훨씬 단순하다는 개념에 기반합니다.

이 '검증자 중심' 설계는 AI 추론의 세 가지 주요 과제를 해결한다. 첫째, 동적 컴퓨팅 할당을 가능케 하여 모델이 어려운 문제에는 더 오래, 쉬운 문제에는 덜 '생각'하도록 한다. 둘째, EBM은 단일 명확한 답이 존재하지 않는 현실 세계 문제의 불확실성을 자연스럽게 수용한다. 셋째, 자체 검증자 역할을 수행하여 외부 모델이 필요 없다.

별도의 생성기와 검증기를 사용하는 다른 시스템과 달리, EBM은 이 둘을 단일 통합 모델로 통합합니다. 이 구성의 주요 이점은 향상된 일반화 능력입니다. 새로운 분포 외(OOD) 데이터에서 해법을 검증하는 것이 일반적으로 올바른 답을 생성하는 것보다 쉽기 때문에, EBM은 익숙하지 않은 상황을 더 잘 처리할 수 있습니다.

잠재력에도 불구하고 EBM은 역사적으로 확장성 문제에 직면해 왔습니다. 이를 해결하기 위해 연구진은 이 프레임워크를 위해 설계된 특수 트랜스포머 모델인 EBT(Enhanced Binary Transformer)를 도입했습니다. EBT는 먼저 컨텍스트와 예측 간의 호환성을 확인한 후, 가장 낮은 에너지(가장 호환성 높은) 출력을 식별할 때까지 예측을 정제하도록 훈련됩니다. 이 절차는 모든 예측에 대해 사고 과정을 효과적으로 모방합니다. 연구팀은 두 가지 EBT 변형을 개발했습니다: GPT 아키텍처에서 영감을 받은 디코더 전용 모델과 BERT와 유사한 양방향 모델입니다.

EBT의 아키텍처는 다양한 추론 시간 확장 방법에 적응하고 호환될 수 있도록 합니다. 일리노이 대학교 어바나-샴페인 캠퍼스 컴퓨터과학 박사 과정생이자 본 논문의 주저자인 알렉시 글래드스톤은 VentureBeat와의 인터뷰에서 "EBT는 더 긴 CoT를 생성하고, 자체 검증하며, 베스트 오브 엔(best-of-N)을 수행하거나 여러 EBT에서 샘플링할 수 있습니다"라고 설명했습니다. "가장 큰 장점은 이러한 모든 기능이 사전 훈련 과정에서 학습된다는 점입니다."

실전 적용 사례

연구진은 기존 아키텍처와 EBT를 비교 테스트했습니다: 텍스트 생성(이산 모달리티)을 위한 대중적인 트랜스포머++ 레시피와 영상 예측 및 이미지 노이즈 제거(연속 모달리티) 같은 작업을 위한 확산 트랜스포머(DiT)가 그 대상이었습니다. 두 가지 주요 기준에 따라 모델을 평가했습니다: "학습 확장성"(훈련 효율성)과 "추론 확장성"(추론 시 더 많은 연산으로 성능이 얼마나 향상되는지 측정).

프리트레이닝 과정에서 EBT는 데이터, 배치 크기, 매개변수, 컴퓨팅 자원을 아우르는 모든 측면에서 Transformer++ 대비 최대 35% 높은 확장률을 달성하며 탁월한 효율성을 입증했습니다. 이는 EBT가 더 빠르고 경제적으로 훈련될 수 있음을 의미합니다.

추론 단계에서도 EBT는 추론 작업에서 기존 모델을 능가했습니다. "더 오래 생각하는"(더 많은 최적화 단계 사용) 방식과 "자기 검증"(여러 후보를 생성하고 에너지 값이 가장 낮은 것을 선택)을 통해 EBT는 Transformer++보다 언어 모델링 성능을 29% 더 향상시켰습니다. 연구진은 "이는 기존 피드포워드 트랜스포머가 예측마다 추가 연산을 동적으로 할당할 수 없어, 더 오래 생각하는 방식으로 토큰별 성능을 개선할 수 없다는 우리의 주장과 일치한다"고 설명합니다.

이미지 노이즈 제거 작업에서 EBT는 DiT보다 우수한 결과를 보여주면서도 전방 전달(forward pass) 횟수를 99% 줄였습니다.

중요한 점은 EBT가 다른 아키텍처보다 효과적으로 일반화한다는 사실이 연구를 통해 밝혀졌다는 것입니다. 동일한 또는 더 약한 사전 훈련 성능을 가졌음에도 EBT는 다운스트림 작업에서 기존 모델을 능가했습니다. 시스템 2 사고 방식의 성능 향상은 훈련 데이터와 달리 분포에서 더 멀리 떨어진 데이터에서 가장 두드러지게 나타났으며, 이는 EBT가 새롭고 어려운 과제에 직면했을 때 특히 강건함을 시사합니다.

연구진은 "EBT 사고의 이점은 모든 데이터에 걸쳐 균일하지 않지만, 분포 변화의 규모에 따라 긍정적으로 확대되며, 이는 훈련 분포를 넘어선 견고한 일반화를 위한 핵심 메커니즘으로서 사고의 중요성을 부각시킨다"고 지적했습니다.

EBT의 장점은 두 가지 핵심 이유로 중요합니다. 첫째, 오늘날의 대규모 기초 모델 규모에서 EBT가 LLM에 사용되는 기존 트랜스포머 아키텍처를 크게 능가할 수 있음을 시사합니다. 저자들은 "1,000배 더 많은 데이터로 훈련되고 1,000배 더 큰 모델을 사용하는 현대 기초 모델 규모에서는 EBT의 사전 훈련 성능이 트랜스포머++ 방식보다 현저히 우수할 것으로 예상한다"고 관찰합니다.

둘째, EBT는 훨씬 뛰어난 데이터 효율성을 보여줍니다. 이는 고품질 훈련 데이터가 AI 확장의 주요 병목 현상으로 부상하는 시대에 결정적인 장점입니다. 논문은 "데이터가 추가 확장의 주요 제약 요인 중 하나가 됨에 따라, 이는 EBT를 특히 매력적으로 만든다"고 설명합니다.

추론 메커니즘이 다르더라도 EBT 아키텍처는 트랜스포머와 높은 호환성을 지녀 현재의 대규모 언어 모델(LLM)을 대체할 수 있습니다.

글래드스톤은 "EBT는 GPU나 TPU에서 피드포워드 모델을 활용한 추측적 디코딩을 포함해 현재 하드웨어/추론 프레임워크와 매우 호환된다"고 말했습니다. 그는 LPU 같은 전용 가속기나 FlashAttention-3 같은 최적화 알고리즘에서도 실행 가능하며, vLLM 같은 일반 추론 프레임워크를 통해 배포될 수 있다고 확신한다고 덧붙였습니다.

개발자와 기업에게 EBT의 강력한 추론 및 일반화 능력은 차세대 AI 애플리케이션 구축을 위한 강력하고 신뢰할 수 있는 기반이 될 수 있습니다. 글래드스톤은 "더 긴 사고는 거의 모든 기업 애플리케이션에 광범위하게 도움이 될 수 있지만, 가장 흥미로운 부분은 더 중요한 의사 결정, 안전성 또는 제한된 데이터를 가진 애플리케이션이 필요할 때일 것"이라고 말했습니다.