[CVPR 2026] Thinking Diffusion: Penalize and Guide Visual-Grounded Reasoning in Diffusion Multimodal Language Models

1. Introduction

Qwen-VL이나 LLaVA와 같은 Vision-Language Model(VLM)은 이미지를 이해하고, 이미지에 대한 질문에 답하거나 복잡한 시각적 추론을 수행하는 등 다양한 vision-language task에서 좋은 성능을 보여주고 있습니다. 여기에 Chain-of-Thought(CoT)를 적용하면 단순히 정답만 바로 생성하는 것이 아니라, 중간 reasoning 과정을 단계적으로 생성한 뒤 최종 답을 내도록 만들 수도 있습니다.

한편 최근에는 기존 autoregressive LLM과는 조금 다른 방식으로 문장을 생성하는 Diffusion Large Language Model, 즉 dLLM이 등장하고 있습니다. 기존 AR LLM은 왼쪽에서 오른쪽으로 토큰을 하나씩 생성합니다. 예를 들어 첫 번째 단어를 생성하고, 그 결과를 이용해 두 번째 단어를 생성하고, 다시 그 결과를 이용해 세 번째 단어를 생성하는 방식입니다. 반면 dLLM은 처음부터 여러 개의 [MASK] 토큰을 만들어놓고, 이 토큰들을 여러 diffusion step에 걸쳐 동시에 복원해나갑니다. 즉 한 번의 step에서 여러 토큰을 parallel하게 처리할 수 있기 때문에 기존 AR 방식보다 빠른 inference가 가능하다는 장점이 있습니다.

이러한 장점을 multimodal 환경까지 확장한 것이 Diffusion Multimodal Large Language Model, 즉 dMLLM입니다. dMLLM은 이미지와 텍스트를 함께 입력받으면서도 기존 AR 방식이 아니라 diffusion 방식으로 답변을 생성합니다. 최근에는 이러한 모델들도 등장하면서 multimodal reasoning을 수행하는 새로운 방향으로 주목받고 있습니다.

그런데 여기서 한 가지 문제가 있습니다.

과연 dMLLM은 실제로 어떤 방식으로 reasoning을 하고 있을까요?

기존 AR 기반 VLM의 reasoning 과정에 대해서는 이미 상당히 많은 연구가 이루어졌습니다. 모델이 reasoning 과정에서 이미지 정보를 얼마나 사용하는지, 어느 layer에서 visual information을 활용하는지, 또는 어떤 visual evidence를 참고해 답을 생성하는지 등에 대한 분석들이 존재합니다. 하지만 dMLLM의 경우에는 이야기가 조금 다릅니다. dMLLM은 여러 토큰을 동시에 복원하기 때문에 기존 AR VLM처럼 앞에서 생성한 reasoning token을 기반으로 다음 reasoning token을 순차적으로 생성하지 않습니다. 따라서 AR VLM을 위해 만들어진 reasoning enhancement 방법을 그대로 적용하기 어렵습니다. 더욱이 dMLLM이 실제 reasoning 과정에서 이미지 정보를 어떻게 사용하는지에 대한 정량적인 분석 역시 거의 이루어지지 않았습니다.

저자는 이 부분에 주목합니다. 즉, dMLLM이 CoT reasoning을 수행할 때 실제로 어떤 과정을 거쳐 답을 만들어내는지 먼저 분석해보자는 것입니다.

그리고 분석 결과 두 가지 중요한 문제를 발견합니다.

첫 번째는 Early Answer Generation입니다.

dMLLM은 diffusion step이 충분히 진행되기도 전에 최종 answer token을 너무 빠르게 결정하는 경향을 보입니다. 특히 diffusion step 수가 적은 경우 이러한 현상이 더 강하게 나타납니다.

예를 들어 Figure 1을 보면 모델이 reasoning을 충분히 진행하기도 전에 이미 최종 답을 만들어놓는 모습을 확인할 수 있습니다. 이후의 diffusion step에서는 새로운 reasoning을 통해 답을 찾는다기보다는, 이미 결정된 답을 정당화하기 위한 중간 reasoning을 채워 넣는 것처럼 동작합니다.

일반적으로 우리가 기대하는 CoT는

Reasoning → Reasoning → Reasoning → Answer

의 형태입니다.

그런데 실제 dMLLM에서는 경우에 따라

Answer 결정 → Reasoning 생성 → Answer 정당화

에 가까운 현상이 나타나는 것이죠.

두 번째 문제는 초기 diffusion step에서 이미지 정보에 대한 dependency가 낮다는 것입니다.

AR 기반 VLM의 경우 reasoning을 시작하는 초기 단계부터 visual information을 적극적으로 참고하는 경향을 보입니다. 반면 dMLLM을 분석해보면 초기 timestep에서는 이미지에 대한 의존도가 상당히 낮고, diffusion 과정이 어느 정도 진행된 이후에야 visual information을 본격적으로 활용하기 시작합니다.

결국 두 현상을 함께 보면 상당히 흥미로운 문제가 생깁니다. 이미지는 아직 제대로 보지도 않았는데 답부터 먼저 결정하는 것입니다. 저자는 이러한 관찰을 기반으로 두 가지 training-free 방법론을 제안합니다.

첫 번째는 Position & Step Penalty, PSP입니다.

앞서 살펴봤듯이 dMLLM은 초기 timestep에서 final answer token을 너무 빠르게 생성하는 문제가 있습니다. PSP는 이를 막기 위해 diffusion의 초기 단계에서 answer가 위치할 것으로 예상되는 token position에 penalty를 부여합니다.

초기 timestep에서는 answer token 생성 확률을 낮춰 reasoning을 먼저 진행하게 만들고, diffusion step이 진행될수록 이러한 penalty를 완화합니다. 이를 통해 reasoning이 timestep에 따라 조금 더 점진적으로 진행되도록 유도합니다.

두 번째는 Visual Reasoning Guidance, VRG입니다.

VRG는 diffusion image generation에서 많이 사용되는 classifier-free guidance에서 아이디어를 가져옵니다. 모델이 이미지를 입력받았을 때의 prediction과 이미지를 제거했을 때의 prediction을 비교하고, 이 둘의 차이를 이용해 visual information이 반영되는 방향을 더욱 강하게 만들어주는 방식입니다. 즉 visual prompt에 의해 발생하는 conditional signal을 증폭시켜 모델이 reasoning 과정에서 이미지 정보를 더욱 적극적으로 활용하도록 만드는 것입니다.

중요한 점은 두 방법 모두 training-free라는 것입니다. 기존 dMLLM을 다시 학습할 필요 없이 inference 과정에 바로 적용할 수 있습니다. 실험 결과 PSP와 VRG를 적용했을 때 기존 reasoning 방식보다 최대 7.5% 높은 accuracy를 달성했습니다.

일반적으로 dLLM에서 reasoning 성능을 높이는 가장 단순한 방법은 diffusion step을 많이 사용하는 것입니다. step을 많이 사용할수록 모델이 토큰을 반복적으로 수정할 기회를 더 많이 얻기 때문입니다. 하지만 당연히 inference time도 함께 증가합니다. 반면 저자가 제안한 방법은 diffusion step을 4배 많이 사용하는 baseline보다도 높은 성능을 보이면서 3배 이상 빠른 inference를 달성합니다.

즉 이 논문의 핵심은 단순히 diffusion step을 늘려 reasoning을 개선하는 것이 아니라,

“dMLLM이 reasoning을 잘못 수행하는 원인을 먼저 분석하고, answer generation timing과 visual information 활용 방식을 직접 조절하자.”라고 정리할 수 있습니다.

2. Background & Related Work

2.1. Diffusion Language Models

본격적인 방법론을 보기 전에 Diffusion Language Model이 어떻게 동작하는지 먼저 간단히 살펴보겠습니다. Diffusion model은 원래 image generation 분야에서 큰 성공을 거둔 방법입니다. 이미지에 점점 noise를 추가한 뒤, 반대로 noise를 제거하면서 원본 이미지를 복원하도록 모델을 학습합니다.

이 아이디어를 text generation에 적용하려는 연구도 꾸준히 진행되어 왔습니다.

초기에는 text token을 continuous embedding으로 변환한 뒤 이 embedding space에서 diffusion을 수행하는 방법들이 제안되었습니다. 하지만 text는 기본적으로 단어 혹은 token이라는 discrete한 값을 갖기 때문에 image처럼 continuous space에서 diffusion을 수행하는 것과는 차이가 있습니다.

이를 해결하기 위해 등장한 것이 Discrete Diffusion Model입니다. Discrete diffusion에서는 vocabulary 자체를 discrete state로 보고 token을 직접 corruption하고 복원합니다. 최근에는 이러한 방식을 대규모 language model까지 확장한 LLaDA나 Dream과 같은 dLLM이 등장하면서 AR LLM과 비교 가능한 수준의 성능까지 보여주고 있습니다.

dLLM의 중요한 특징 중 하나는 diffusion step 수를 조절할 수 있다는 것입니다. step을 많이 사용하면 모델이 여러 번 token을 수정할 기회를 얻기 때문에 일반적으로 generation quality가 좋아질 수 있지만 그만큼 inference가 느려집니다. 반대로 step을 줄이면 빠르게 생성할 수 있지만 성능이 떨어질 수 있습니다. 즉 diffusion step을 통해 speed-quality trade-off를 직접 조절할 수 있는 것입니다.

조금 더 구체적으로 dLLM의 generation 과정을 살펴보겠습니다. 길이가 (L)인 token sequence가 있다고 가정하겠습니다.

원래 문장을

라고 하겠습니다.

Forward diffusion 과정에서는 이 token들을 점점 [MASK] token으로 바꿔나갑니다.

diffusion time (t)가 증가할수록 mask 비율이 증가하고, 최종적으로 (t=1)이 되면 모든 token이

[X_1 = [M, M, … , M]] 과 같이 완전히 masking된 상태가 됩니다.

모델은 이 과정의 반대 방향을 학습합니다. 즉 완전히 masking된 문장에서 시작해서 원래 문장을 복원하도록 학습하는 것입니다.

Inference에서도 마찬가지입니다.

먼저 생성하고 싶은 sequence length를 정한 뒤 모든 위치를 [MASK]로 초기화합니다.

예를 들어 답변 길이를 8개 token으로 설정했다면 처음에는

[MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK]

와 같은 상태에서 시작합니다.

이후 모델은 현재 입력과 prompt, 그리고 이미지와 같은 multimodal condition을 참고하여 모든 [MASK] 위치의 token을 동시에 예측합니다. 여기서 AR 모델과 가장 큰 차이가 나타납니다. AR 모델이라면

The → cat → is → sitting ...

처럼 token 하나씩 순차적으로 생성하지만, dLLM은 한 timestep에서 여러 위치의 token을 동시에 예측합니다. 하지만 한 번 예측했다고 모든 token을 그대로 확정하는 것은 아닙니다. 예측한 token 중 일부는 유지하고, 일부는 다시 [MASK]로 돌려놓습니다. 그리고 다음 timestep에서 다시 예측합니다. 이 과정을 반복하면서 점점 mask의 개수가 줄어들고 최종적으로 완성된 sequence를 얻는 것입니다.

2.2. Multimodal Chain-of-Thought

다음으로 Multimodal Chain-of-Thought에 대해 살펴보겠습니다.

Chain-of-Thought, 즉 CoT는 LLM의 reasoning 성능을 크게 향상시킨 대표적인 방법입니다. 모델에게 바로 정답만 생성하도록 하는 대신 중간 reasoning 과정을 먼저 생성하도록 함으로써 복잡한 문제를 단계적으로 풀도록 만드는 것입니다. 이러한 CoT는 자연스럽게 multimodal model에도 적용되어 왔습니다. 그런데 multimodal CoT에서는 일반적인 text-only CoT에는 없는 문제가 하나 추가됩니다. 바로 서로 다른 modality의 정보를 어떻게 reasoning 과정에 함께 사용할 것인가입니다. 가장 간단한 방법은 이미지를 먼저 text로 바꾸는 것입니다. 예를 들어 이미지 caption을 생성한 뒤, Image caption + Question을 LLM에 입력하여 reasoning하도록 만들 수 있습니다. 이 방식은 기존 LLM의 강력한 reasoning 능력을 그대로 활용할 수 있다는 장점이 있습니다.

하지만 명확한 한계도 있습니다.

이미지에 존재하는 모든 정보를 caption 하나에 담는 것은 어렵습니다. 작은 object, 정확한 위치 관계, 세부적인 visual cue 등이 image-to-text 과정에서 손실될 수 있습니다. 그리고 한번 caption에서 빠진 정보는 이후 LLM이 아무리 reasoning을 잘하더라도 다시 복구할 수 없습니다. 이러한 한계를 해결하기 위해 이미지를 VLM에 직접 입력하면서 CoT reasoning을 수행하는 연구들이 등장했습니다. 하지만 이 방법들 역시 reasoning 자체는 대부분 text token을 기반으로 이루어지기 때문에 이미지와 텍스트 사이의 fine-grained relation을 충분히 활용하기 어렵다는 문제가 존재합니다. 최근 연구들은 이를 해결하기 위해 reasoning 과정에 더 많은 visual information이나 structural information을 집어넣는 방향으로 발전하고 있습니다.

예를 들어 CCoT는 이미지에서 scene graph를 생성합니다. 이미지 안에 어떤 object가 있고, object 간 relation이 무엇이며, 각각의 spatial relationship이 어떻게 형성되어 있는지를 구조화한 뒤 이를 prompt에 넣어 reasoning을 수행합니다. DDCoT는 문제를 reasoning이 필요한 부분과 visual recognition이 필요한 부분으로 나누고, 각각을 단계적으로 해결하도록 모델을 유도합니다. 또 ICoT는 더 직접적인 방식을 사용합니다. 텍스트로 표현하기 어려운 세밀한 visual cue를 reasoning 과정에 image patch 형태로 직접 삽입합니다. 따라서 reasoning sequence가

Text reasoning → Image patch → Text reasoning → Image patch

와 같이 text와 visual information이 interleave되는 형태를 갖게 됩니다. 이러한 방법들은 AR 기반 VLM에서는 좋은 성능을 보여주었습니다. 하지만 중요한 문제가 있습니다.

이 방법론들은 기본적으로 token이 순차적으로 생성된다는 AR generation structure를 전제하고 있습니다. 앞에서 생성된 reasoning token을 보고 다음 reasoning token을 생성하면서 reasoning이 단계적으로 진행되는 것입니다.

반면 dMLLM은 여러 위치의 token을 parallel하게 복원합니다. 따라서 기존 AR VLM에서 효과적이었던 step-by-step multimodal CoT 방법을 dMLLM에 그대로 적용한다고 해서 동일한 효과를 기대하기 어렵습니다.

2.3. CoT in Diffusion Language Models

마지막으로 diffusion language model에서 CoT reasoning이 어떻게 이루어지는지 살펴보겠습니다.

앞서 설명했듯이 기존 AR language model의 CoT는 왼쪽에서 오른쪽으로 token을 하나씩 생성합니다. 예를 들어

문제 이해 → 첫 번째 추론 → 두 번째 추론 → 결과 계산 → 최종 답

과 같은 reasoning chain을 순서대로 만들어갑니다.

이러한 구조의 장점은 이전에 생성한 reasoning을 다음 token generation에 그대로 활용할 수 있기 때문에 모델 스스로 만든 rationale을 기반으로 reasoning을 계속 이어나갈 수 있습니다.

하지만 단점도 존재합니다. 우선 reasoning이 길어지면 token을 하나씩 생성해야 하기 때문에 inference time과 computational cost가 거의 선형적으로 증가합니다. 또 한 가지 중요한 문제는 이미 생성된 token을 수정하기 어렵다는 것입니다.

반면 dLLM은 구조적으로 조금 다른 특성을 갖습니다. dLLM에서는 매 timestep마다 일부 token을 unmask하여 예측한 뒤, 아직 확신하지 못한 token을 다시 mask하고 다음 timestep에서 다시 생성할 수 있습니다. 즉 잘못 생성된 token도 이후 diffusion step에서 다시 수정할 수 있는 것입니다. 이 때문에 diffusion model의 reasoning에서는 어떤 token을 다시 mask할 것인가, 즉 remasking strategy가 굉장히 중요합니다.

가장 직관적인 방법은 confidence를 사용하는 것입니다. 모델이 높은 confidence로 예측한 token은 유지하고 confidence가 낮은 token은 다시 [MASK]로 만든 뒤 다음 step에서 재생성하는 방식입니다. 이외에도 entropy가 높은 token을 remask하거나, 첫 번째와 두 번째로 높은 probability 사이의 margin을 이용하는 등의 방법이 연구되어 왔습니다. 하지만 어떤 remasking strategy가 reasoning에 가장 적합한지에 대한 연구는 아직 초기 단계입니다.

현재 diffusion CoT는 전반적으로 AR CoT만큼 강력한 reasoning 성능을 보여주지는 못하고 있습니다. 그럼에도 diffusion CoT에는 AR 방식이 갖기 어려운 두 가지 중요한 장점이 있습니다.

첫 번째는 parallel generation을 통한 빠른 reasoning입니다.

두 번째는 speed와 quality 사이의 trade-off를 유연하게 조절할 수 있다는 점입니다. 빠른 응답이 중요하다면 diffusion step을 줄일 수 있고, 더 높은 reasoning 성능이 필요하다면 step을 늘릴 수 있습니다.

따라서 이 논문에서는 단순히 dMLLM에 CoT를 적용하는 것에서 끝나는 것이 아니라, dMLLM 특유의 이러한 speed-quality trade-off를 유지하면서도 reasoning을 더 효과적으로 수행할 수 있는 방법을 찾는 것을 목표로 합니다.

3. Analysis of Reasoning in dMLLMs

3.1. Experimental Setup for Case Study

앞선 Introduction에서 dMLLM이 CoT reasoning을 수행할 때 답을 너무 빨리 생성하고, 동시에 초기 단계에서는 이미지 정보를 충분히 활용하지 않는 것 같다는 두 가지 문제를 말씀드렸습니다. 그럼 실제로 이런 현상이 나타나는지 정량적으로 확인해보겠습니다.

저자는 dMLLM의 reasoning process를 분석하기 위해 M3CoT Validation Set을 사용합니다. M3CoT는 science, mathematics, commonsense 등 다양한 분야의 문제를 포함하고 있으며, 단순히 이미지를 보고 바로 답할 수 있는 문제가 아니라 multimodal input을 기반으로 여러 단계의 reasoning이 필요한 benchmark입니다.

분석에서는 크게 두 가지 설정을 변화시킵니다.

첫 번째는 diffusion step (T)이고, 두 번째는 모델이 생성할 수 있는 response length (L)입니다. 쉽게 말하면 diffusion step은 모델에게 답을 몇 번에 걸쳐 수정할 기회를 줄 것인지에 해당하고, response length는 reasoning과 answer를 포함해 최대 몇 개의 token을 생성할 수 있도록 할 것인지에 해당합니다.

저자는 이러한 두 설정을 변화시키면서 dMLLM이 실제로 언제 answer를 결정하는지, 그리고 각 timestep에서 visual information을 얼마나 참고하는지를 분석합니다.

3.2. Observation 1: Early Answer Generation

먼저 첫 번째로 살펴볼 것은 dMLLM이 언제 최종 answer를 생성하는가입니다. 일반적으로 CoT reasoning을 기대한다면 모델이 먼저 충분한 reasoning을 수행하고 마지막 부분에서 최종 답을 생성해야 합니다.

예를 들어

Reasoning → Reasoning → Reasoning → Final Answer

와 같은 순서가 자연스럽습니다.

그런데 dMLLM에서는 diffusion 방식으로 sequence 전체를 동시에 복원하기 때문에 반드시 왼쪽부터 순차적으로 reasoning이 만들어지는 것은 아닙니다. 그래서 저자는 각 sample마다 최종 answer token이 처음 등장하는 timestep을 기록하고, 전체 sample에 대해 그 분포를 Figure 3에 시각화합니다.

먼저 Figure 3(a)는 response length가 (L=64)이고 diffusion step이 상대적으로 작은 설정에서의 결과입니다. 결과를 보면 상당히 많은 sample에서 최종 answer가 diffusion process의 매우 이른 시점에 등장합니다. 특히 전체 sample의 30% 이상에서 answer가 매우 초기 step에 이미 생성되는 모습을 확인할 수 있습니다. 이는 모델이 충분한 reasoning을 통해 answer를 찾아가는 것이 아니라, reasoning이 완성되기 전에 answer부터 먼저 정해버릴 가능성이 높다는 것입니다. 이 경우 뒤에서 생성되는 rationale 역시 실제로 answer를 도출하기 위한 reasoning이라기보다는, 이미 결정된 answer를 설명하거나 정당화하는 방향으로 만들어질 가능성이 있습니다.

반면 Figure 3(b)를 보면 response length를 (L=256)으로 늘리고 diffusion step 역시 (T=128)처럼 충분히 크게 설정했을 때는 다른 양상이 나타납니다. 대부분의 sample에서 answer가 더 뒤쪽 timestep에서 생성됩니다. 즉 모델에게 충분한 generation space와 diffusion step을 제공하면 intermediate reasoning이 어느 정도 형성된 이후 최종 answer가 등장하는 보다 안정적인 reasoning pattern이 나타나는 것입니다.

여기서 한 가지 직관적인 결론을 얻을 수 있습니다.

dMLLM도 충분한 diffusion step과 response length를 주면 reasoning을 잘 할 수 있습니다.

문제는 그렇게 하면 diffusion model을 사용하는 핵심 장점 중 하나인 빠른 inference가 약해진다는 것입니다. reasoning 성능을 안정화하기 위해 무조건 (T)와 (L)을 크게 설정해야 한다면, 결국 계산량과 inference time이 크게 증가하게 됩니다. 따라서 저자는 단순히 diffusion step을 늘리는 대신, 적은 diffusion step에서도 answer를 너무 일찍 확정하지 않도록 generation process 자체를 조절할 필요가 있다고 봅니다.

그리고 이를 해결하기 위해 뒤에서 설명할 새로운 remasking strategy인 Position & Step Penalty(PSP)를 제안합니다.

3.3. Observation 2: Weak Visual Grounding

두 번째로 저자는 dMLLM이 reasoning 과정에서 이미지 정보를 실제로 얼마나 활용하는지를 분석합니다.

이를 위해 기존 연구에서 사용된 Visual Prompt Dependency Measure, PDM을 사용합니다.

PDM은 간단히 말하면,

“이미지를 넣었을 때와 이미지를 빼고 넣었을 때 모델의 token prediction이 얼마나 달라지는가?”

를 측정하는 값이라고 생각하시면 됩니다. 수식으로는 다음과 같이 표현됩니다.

여기서 (v)는 visual input, (c)는 text prompt, 그리고 (X_t)는 현재 diffusion timestep에서의 partially masked sequence를 의미합니다. 수식을 자세히 볼 필요 없이 핵심만 보면, 이미지를 입력했을 때와 입력하지 않았을 때 prediction distribution의 차이가 크면 PDM이 높아집니다.

그러니까 PDM이 높다는 것은 해당 token을 생성할 때 모델이 이미지에 많이 의존하고 있다는 뜻이고, PDM이 낮다면 이미지를 제거해도 prediction이 크게 변하지 않는다는 뜻입니다. 저자는 먼저 AR 기반 VLM인 LLaVA-1.5를 살펴봅니다.

Figure 4를 보면 LLaVA-1.5는 generation 초기 단계에서 높은 PDM을 보이고, reasoning이 진행될수록 PDM이 점차 감소합니다. 즉 AR VLM은 처음에 이미지를 보고 상황을 파악한 뒤, 이후에는 앞에서 생성한 reasoning을 기반으로 답을 이어가는 형태를 보이는 것입니다.

그런데 dMLLM은 완전히 다른 모습을 보입니다. 초기 diffusion timestep에서 PDM이 매우 낮습니다. 즉 reasoning의 시작 단계에서는 이미지를 넣든 빼든 token prediction이 크게 달라지지 않는다는 것입니다. 쉽게 말하면 초반에는 이미지를 거의 보지 않고 reasoning을 시작한다고 해석할 수 있습니다. 그리고 diffusion process가 진행될수록 PDM이 점점 증가합니다. 즉 dMLLM은 오히려 뒤쪽 timestep으로 갈수록 visual information을 더 적극적으로 사용하기 시작합니다. 여기까지만 보면 단순히 AR VLM과 dMLLM의 visual information 사용 pattern이 다르다고 생각할 수도 있습니다. 그런데 앞에서 확인한 Early Answer Generation과 함께 보면 문제가 훨씬 명확해집니다.

Observation 1에서는 dMLLM이 diffusion 초기 단계에서 이미 final answer를 생성하는 경우가 많았습니다. 그런데 Observation 2에서는 바로 그 초기 단계에서 visual dependency가 가장 낮게 나타납니다. 이미지 정보는 아직 충분히 활용하지 않았는데 answer는 이미 만들어지고 있는 것입니다. 이 두 결과를 종합하면 dMLLM의 reasoning process를 다음처럼 해석할 수 있습니다.

Weak Visual Grounding → Early Answer Generation → 이후 Visual Information 반영 → 이미 정해진 Answer를 중심으로 Reasoning 형성

즉 모델이 이미지에 충분히 grounding된 상태에서 reasoning을 통해 answer를 찾는 것이 아니라, image information을 제대로 활용하기 전에 answer를 먼저 결정하고 이후 reasoning trajectory를 그 answer에 맞춰 구성할 가능성이 있다는 것입니다. 저자는 이것을 dMLLM reasoning의 두 번째 핵심 문제로 정의합니다.

따라서 첫 번째 문제인 Early Answer Generation에는 PSP, 두 번째 문제인 Weak Visual Grounding에는 VRG를 적용하여 각각 직접적으로 해결합니다.

4. Method

4.1. Position & Step Penalty

먼저 Position & Step Penalty, PSP부터 살펴보겠습니다. dMLLM의 generation 과정을 다시 간단히 생각해 보면 모델은 길이가 (L)인 response sequence 전체를 [MASK] token으로 초기화합니다.

[X_1=[M,M,…,M]]

그리고 총 (K)개의 diffusion timestep [{t_1,t_2,…,t_K}]

을 거치면서 점차 token을 복원해나갑니다. 각 timestep에서는 모델이 현재 masked sequence를 기반으로 전체 token을 예측합니다.

그리고 예측된 token의 confidence를 기준으로 일부 token은 그대로 유지하고, confidence가 낮은 token은 다시 mask하여 다음 timestep에서 재생성합니다. 즉 어떤 token을 유지하고 어떤 token을 다시 mask할 것인가를 결정하는 confidence가 굉장히 중요합니다. PSP는 바로 이 confidence score를 조절합니다. 기존에는 (j)번째 token의 confidence를 단순히 (C_ji)라고 사용했다면, PSP에서는 다음과 같이 penalty를 적용합니다.

여기서 각 항이 어떤 역할을 하는지 하나씩 보겠습니다. 먼저 τi는 현재 diffusion process가 얼마나 진행되었는지를 나타냅니다. 초기 timestep에서는 τi가 작고, 마지막 timestep으로 갈수록 1에 가까워집니다.

다음으로 rel(j)는 token이 response sequence의 어느 위치에 있는지를 나타냅니다. sequence 앞쪽 token은 0에 가깝고, 뒤쪽 token은 1에 가까운 값을 갖습니다.

왜 위치 정보를 사용할까요?

일반적인 CoT response를 생각하면 reasoning은 앞부분에 위치하고, final answer는 뒤쪽에 위치하는 경우가 많습니다. 따라서 PSP는 초기 diffusion step에서 response 뒤쪽에 있는 token의 confidence를 의도적으로 낮춥니다. 즉 answer가 위치할 가능성이 높은 뒷부분 token이 너무 빨리 확정되지 않도록 다시 mask될 가능성을 높여주는 것입니다. 예를 들어 diffusion 초반에 sequence의 마지막 부분에서 모델이 높은 confidence로 answer를 생성했다고 하더라도 PSP가 해당 confidence를 낮춰버립니다. 그러면 그 token은 그대로 유지되지 않고 다시 [MASK]가 될 가능성이 높아집니다.

여기서 중요한 것은 이 penalty가 모든 timestep에서 동일하게 적용되지 않는다는 점입니다. 초기 timestep에서는 ((1-τi))가 크기 때문에 강한 penalty가 적용됩니다. 하지만 diffusion이 진행될수록 τi가 1에 가까워지고 penalty는 자연스럽게 감소합니다. 그래서 이름도 Position & Step Penalty입니다.

4.2. Visual Reasoning Guidance

다음은 두 번째 문제인 Weak Visual Grounding을 해결하기 위한 Visual Reasoning Guidance, VRG입니다.

앞선 분석에서 dMLLM은 초기 diffusion timestep에서 visual input에 대한 dependency가 매우 낮다는 것을 확인했습니다. 특히 문제가 되는 것은 이 시기가 Early Answer Generation이 발생하는 시점과 겹친다는 것입니다.

따라서 저자는 reasoning 과정에서 visual information의 영향 자체를 더 강하게 만들어주는 방법을 제안합니다. 여기서 사용되는 아이디어가 Classifier-Free Guidance, CFG입니다. CFG는 diffusion image generation에서 매우 널리 사용되는 방법입니다. 기본적인 아이디어는 상당히 간단합니다.

예를 들어 “a red car”라는 text condition을 사용하여 이미지를 생성한다고 해보겠습니다. 모델에게 condition을 넣은 경우와 condition을 넣지 않은 경우의 prediction을 각각 계산합니다.

ϵ_cond 와 ϵ_uncond 를 얻은 뒤, 둘 사이의 차이를 증폭시킵니다.

여기서 ϵ_cond – ϵ_uncond는 condition이 모델의 prediction을 어느 방향으로 변화시켰는지를 의미합니다. 이 차이를 guidance scale (s)를 이용해 더 크게 만들어주면 모델이 condition을 보다 강하게 따르도록 만들 수 있습니다.

VRG는 이 아이디어를 visual reasoning에 그대로 적용합니다. 먼저 이미지가 있는 상태와 이미지가 없는 상태에서 각각 logits를 계산합니다. 이미지를 사용한 conditional logits를 logits_c 라고 하고, 이미지를 제외한 unconditional logits를 logits_u 라고 하겠습니다.

그리고 다음과 같이 visual guidance를 적용합니다.

VRG는 이 변화량을 증폭시킵니다. 결과적으로 visual evidence가 token prediction에 미치는 영향이 커지고, reasoning 과정에서 모델이 이미지에 더 강하게 의존하도록 유도할 수 있습니다.

마지막으로 저자는 두 방법을 결합합니다. 먼저 VRG를 적용한 logits에 softmax를 취해 각 token의 confidence를 계산합니다.

즉 전체 과정을 정리하면 다음과 같습니다. 먼저 VRG를 통해 visual information이 강하게 반영된 token probability를 만들고, 그다음 PSP를 통해 초기 timestep에서 answer-position token이 너무 빨리 확정되지 않도록 confidence를 조절합니다. 결국 두 방법을 함께 적용하면, 우리가 기대하는 reasoning process에 조금 더 가깝게 dMLLM의 generation을 유도할 수 있는 것입니다.

5. Experiments

5.1. Experimental Setup

저자는 제안한 PSP와 VRG의 효과를 검증하기 위해 M3CoT, ScienceQA, MMBench, V* Bench 총 4개의 benchmark를 사용합니다. 모델은 diffusion 기반 multimodal model인 LaViDa-llada-reason과 MMaDA-8B-MixCoT를 사용합니다. Baseline으로는 dMLLM에서 일반적으로 사용하는 세 가지 remasking strategy인 Low-confidence, Entropy, Margin을 비교하고, 기존 AR VLM을 위해 제안된 CoT 방법인 CCoT와 DDCoT도 함께 비교합니다.

5.2. Main Results

먼저 Table 1의 결과를 보면 기존 AR VLM에서 효과적이었던 CCoT와 DDCoT가 dMLLM에서는 큰 성능 향상을 보여주지 못합니다. DDCoT는 reasoning 과정을 여러 단계로 나누고, CCoT는 visual information을 구조화하여 reasoning을 강화하는 방법이지만, diffusion 방식에서는 단순 remasking strategy보다도 좋은 결과를 내지 못하는 경우가 많습니다.

반면 저자가 제안한 PSP + VRG는 모든 benchmark와 대부분의 (L/T) 설정에서 가장 높은 성능을 기록합니다. 특히 단순한 visual perception뿐 아니라 M3CoT나 V* Bench처럼 reasoning과 세밀한 visual grounding이 필요한 benchmark에서도 꾸준히 성능이 향상됩니다.

가장 중요한 것은 효율성입니다. 일반적으로 dMLLM의 성능을 높이기 위해서는 diffusion step이나 generation length를 크게 설정하면 됩니다. 하지만 그렇게 하면 inference가 느려집니다. 그런데 PSP와 VRG를 적용하면 훨씬 짧은 generation length와 적은 diffusion step으로도 더 큰 설정의 baseline보다 높은 성능을 얻을 수 있습니다.

예를 들어 MMBench에서 LaViDa의 Low-conf baseline은 (L=256)일 때 74.3%를 기록하지만, PSP와 VRG를 적용하면 (L=64)만 사용하고도 74.9%를 달성합니다. MMaDA 역시 같은 조건에서 PSP와 VRG를 사용했을 때 긴 generation length를 사용하는 baseline보다 최대 7.5% 높은 정확도를 기록합니다.

5.3. Ablation & Discussion

다음으로 PSP와 VRG가 각각 실제로 효과가 있는지 확인합니다.

Table 2를 보면 PSP만 적용하거나 VRG만 적용해도 Low-conf baseline보다 성능이 향상됩니다. 그리고 두 방법을 함께 사용했을 때 네 개 benchmark에서 가장 좋은 성능을 기록합니다. 즉 PSP와 VRG는 서로 다른 문제를 해결하면서 상호보완적으로 작동한다고 볼 수 있습니다.

Table 3에서는 Low-conf뿐 아니라 Entropy와 Margin remasking에도 PSP와 VRG를 적용합니다. 결과적으로 어떤 remasking strategy를 사용하더라도 성능이 꾸준히 향상됩니다. 즉 PSP와 VRG는 특정 remasking 방식에 종속되지 않고 plug-and-play 형태로 적용할 수 있는 방법이라는 것입니다.

Efficiency

Table 4에서는 inference time을 비교합니다. 먼저 DDCoT는 reasoning과 recognition을 따로 수행하기 때문에 가장 큰 inference cost를 요구합니다. 반면 PSP는 confidence score에 penalty만 적용하기 때문에 기존 Low-conf와 거의 동일한 inference time을 유지합니다. VRG는 이미지가 있는 경우와 없는 경우의 logits를 각각 계산해야 하기 때문에 PSP보다 약간의 추가 연산이 필요합니다. 그럼에도 PSP와 VRG를 적용한 작은 (L/T) 설정이 훨씬 큰 (L/T) baseline보다 높은 성능을 보이기 때문에 전체 speed-quality trade-off 관점에서는 충분한 효율성을 보여줍니다.