EndoCoT: 디퓨전 모델의 내생적 사고 사슬 추론 확장

기존 디퓨전 모델은 텍스트를 한 번만 해석하여 복잡한 논리 문제를 풀 때 한계가 있었다. 이 논문은 모델 내부에서 단계별로 생각을 다듬는 '사고 사슬' 기법을 적용해 미로 찾기나 스도쿠 같은 고난도 시각 추론 문제를 해결하며, 생성 AI의 논리적 한계를 돌파하는 새로운 방향을 제시한다.

기존 디퓨전 모델은 텍스트 인코더(MLLM)가 생성 시작 시점에 단 한 번만 임베딩을 생성하는 '정적 가이드' 방식에 의존한다. 이는 미로나 스도쿠처럼 단계별 논리가 필요한 작업에서 Attention 메커니즘이 복잡한 제약 조건을 한 번에 처리하지 못해 물리적 벽을 뚫거나 숫자를 중복 배치하는 등의 오류를 범하게 만든다.

EndoCoT는 이를 해결하기 위해 MLLM의 내부 잠재 상태(Latent State)를 반복적으로 업데이트하는 방식을 도입한다. 마치 사람이 문제를 풀 때 중간 과정을 메모하듯, 모델이 이전 단계의 '생각'을 다음 단계의 입력으로 다시 넣어 점진적으로 논리를 구체화한다. 이 과정에서 Embedding 레이어를 거치지 않고 히든 스테이트를 직접 전달하여 정보 손실을 최소화한다.

이렇게 정제된 '생각'은 Diffusion Transformer(DiT)의 디노이징 과정에 실시간으로 주입되어, 매 단계마다 변화하는 상황에 맞는 동적인 가이드를 제공한다. 그 결과, 단순한 패턴 매칭을 넘어선 실제적인 문제 해결 능력을 갖추게 되며, 추론 시간을 더 많이 할당할수록 성능이 향상되는 Scaling Law를 보여준다.

EndoCoT는 MLLM과 DiT를 결합한 구조로, 반복적인 잠재 사고 정제 과정을 통해 추론을 수행한다. 전체 과정은 Flow Matching 프레임워크를 기반으로 하며, 선형적인 확률 경로를 따라 노이즈에서 데이터로 변환되는 과정을 학습한다.

Iterative Thought Guidance는 $h_\tau = e^\top_{L+1} f_\phi([P; h_{\tau-1}])$ 수식을 통해 구현된다. [이전 단계의 사고 상태 $h_{\tau-1}$와 프롬프트 임베딩 $P$를 결합하여] → [MLLM 레이어 $f_\phi$를 통과시킨 뒤] → [$(L+1)$번째 위치의 히든 스테이트를 추출하여] → [현재 단계의 새로운 사고 상태 $h_\tau$를 생성한다]. 이 값은 DiT의 디노이징 가이드로 사용된다.

Terminal Thought Grounding은 최종 사고 상태 $h_T$를 텍스트 정답 임베딩 $h_{ref}$와 정렬한다. [최종 사고 상태와 정답 텍스트 임베딩 사이의 L2 거리를 계산하여] → [Semantic Loss $\mathcal{L}_{align}$를 산출하고] → [이를 전체 손실 함수에 추가함으로써] → [모델의 추론 경로가 텍스트 의미론적 맥락에서 벗어나지 않도록 강제한다].

학습은 2단계 점진적 전략을 따른다. 1단계에서는 모든 중간 추론 단계에 대해 Flow Matching Loss와 Alignment Loss를 적용하여 추론 능력을 개발한다. 2단계에서는 중간 단계의 그래디언트 전파를 차단한 채 최종 출력의 시각적 품질만을 최적화하여 성능을 극대화한다.

Maze, TSP, Sudoku, VSP 등 4가지 시각 추론 벤치마크에서 평균 92.1%의 정확도를 기록하며 기존 SOTA 모델인 DiffThinker(83.8%)를 8.3%p 차이로 앞질렀다. 특히 난이도가 높은 Maze-32와 Sudoku-35에서는 각각 90%, 95%의 정확도를 달성하여 기존 모델 대비 25~40%의 성능 향상을 보였다.

Ablation Study 결과, Semantic Loss를 제거했을 때 Maze-32 정확도가 90%에서 14%로 급락하는 것이 확인되었다. 이는 텍스트 기반의 가이드가 장기적인 추론 경로 유지에 필수적임을 입증한다. 또한 명시적인 텍스트 토큰을 생성하는 방식보다 잠재 공간에서의 연속적인 토큰 업데이트 방식이 오류 누적 방지에 더 효과적임이 확인됐다.

추론 시간 확장(Inference-Time Scaling) 분석에서 사고 단계 $\tau$를 늘릴수록 정확도가 꾸준히 상승하는 현상이 관찰되었다. Maze-32 기준 $\tau=2$일 때 11%였던 정확도가 $\tau=50$일 때 90%까지 상승하며, 계산 자원을 더 투입할수록 더 복잡한 문제를 해결할 수 있는 확장성을 증명했다.

EndoCoT 아키텍처는 Qwen-Image-Edit-2511을 베이스 모델로 하며, MLLM(텍스트 인코더)과 DiT(디퓨전 백본) 모두에 LoRA(rank=32)를 적용하여 공동 최적화한다. 이는 고수준의 논리 계획과 저수준의 시각적 구현 사이의 시너지를 극대화하기 위한 설계다.

핵심 차별점은 '내생적(Endogenous)' 추론 구조다. 외부에서 텍스트 추론 결과를 생성하여 다시 주입하는 대신, MLLM의 내부 히든 스테이트를 직접 DiT의 컨디셔닝 신호로 사용함으로써 이산적인 토큰 생성 과정에서 발생하는 정보 손실과 오류 누적 문제를 해결했다.

Flow Matching 프레임워크를 채택하여 확산 과정의 효율성을 높였다. 각 추론 단계 $\tau$는 독립적인 디노이징 궤적을 가지며, 이전 단계의 사고 상태 $h_{\tau-1}$가 다음 단계의 초기 컨텍스트로 작용하는 재귀적 구조를 취한다.

레이어별 민감도 분석(Layer-Wise Sensitivity Analysis)을 통해 논리 추론의 부하가 주로 MLLM의 마지막 레이어와 DiT의 초기 레이어 접합부에 집중됨을 확인했다. 이를 바탕으로 해당 구간의 상호작용을 강화하는 설계를 반영하여 추론 효율을 높였다.

최적의 사고 단계(Reasoning Steps) 수를 수동으로 설정해야 하며, 학습을 위해 고품질의 중간 단계 감독 데이터셋이 필수적이라는 제약이 존재한다.