ECHO: 1단계 블록 확산을 이용한 효율적인 흉부 엑스레이 판독문 생성

기존의 자동 판독 시스템은 토큰을 하나씩 순차적으로 생성하여 속도가 느렸으나, 이 논문은 확산 모델을 활용해 한 번에 여러 토큰을 생성함으로써 추론 속도를 8배 높였다. 특히 의료 현장에서 중요한 임상적 정확도를 유지하면서도 실시간에 가까운 보고서 생성이 가능하다는 점에서 큰 의미가 있다.

기존의 Transformer 기반 모델은 문장을 만들 때 단어를 하나씩 순서대로 예측하는 Autoregressive 방식을 사용한다. 이 방식은 앞 단어가 나와야 다음 단어를 계산할 수 있어 생성 속도가 느리다는 근본적인 한계가 있다. 이를 해결하기 위해 모든 단어를 동시에 예측하는 Diffusion 모델이 대안으로 떠올랐지만, 단어들이 서로 어떤 관계인지 고려하지 않고 독립적으로 예측하려다 보니 문장의 일관성이 깨지는 '평균장 편향(Mean-field Bias)' 문제가 발생한다.

ECHO는 이 문제를 해결하기 위해 '직접 조건부 증류(DCD)'라는 개념을 도입한다. 이는 숙련된 교사 모델이 여러 단계에 걸쳐 신중하게 단어 간의 관계를 파악하며 문장을 완성해가는 과정을 관찰하고, 그 핵심적인 '단어 간 연결 고리' 정보를 학생 모델에게 직접 전달하는 방식이다. 학생 모델은 이 정보를 통해 단 한 번의 연산만으로도 교사 모델이 여러 번 고민해서 만든 것과 같은 수준의 일관된 문장을 만들어낼 수 있게 된다.

결과적으로 ECHO는 딥러닝의 기초인 확률 분포 예측에서 토큰 간의 결합 확률(Joint Probability)을 직접 학습함으로써, 속도와 품질이라는 두 마리 토끼를 잡았다. 이는 의료 영상 분석처럼 방대한 양의 데이터를 빠르게 처리해야 하면서도 단 하나의 단어 오차도 허용되지 않는 정밀한 분야에서 혁신적인 변화를 가능하게 한다.

ECHO의 학습 파이프라인은 세 단계로 구성된다. 1단계에서는 Lingshu-7B 모델을 기반으로 흉부 엑스레이 데이터셋에 대해 지속적 사전 학습(CPT)을 수행하여 도메인 지식을 확보한다. 2단계에서는 Response-Asymmetric Diffusion(RAD) 기법을 적용하여 자기회귀 모델을 블록 확산 모델로 변환한다. RAD는 긴 시각적 토큰은 고정하고 텍스트 응답 블록에 대해서만 확산 과정을 적용하여 연산 효율을 극대화한다.

핵심 메커니즘인 Direct Conditional Distillation(DCD)은 두 단계로 작동한다. [교사 모델의 디노이징 궤적 입력 → 각 단계에서 확신도가 높은 토큰들을 추출 및 결합 → 비계수화된(Unfactorized) 타겟 분포 생성 → 학생 모델의 1단계 예측값과 KL Divergence를 통해 정렬] 과정을 거친다. 구체적으로 KL Divergence는 P_tch(교사 분포)와 Q_phi(학생 분포) 사이의 거리인 Σ P_tch * log(P_tch / Q_phi)를 계산하며, 이 값이 0에 가까워지도록 학생 모델을 최적화한다.

추론 단계에서는 Fused Block KV Cache 기술을 사용한다. [현재 블록의 디노이징 연산 수행 → 이전 블록의 KV 캐시 업데이트를 현재 연산에 통합 → 별도의 캐시 업데이트 단계 제거 → 추론 횟수 절반 감소] 순서로 동작하여 추가적인 연산량 증가 없이 지연 시간을 획기적으로 줄인다.

MIMIC-CXR, CheXpert-Plus 등 주요 벤치마크에서 ECHO는 기존 SOTA 모델들을 압도하는 성과를 보였다. 특히 임상적 정확도를 측정하는 SemScore에서 기존 자기회귀 모델인 Lingshu-7B(27.54%) 대비 ECHO는 53.4%를 기록하며 약 2배에 가까운 성능 향상을 보였다. 언어적 품질 지표인 ROUGE-L에서도 56.14를 기록하여 타 확산 기반 모델(CD4LM: 49.21)보다 우수한 결과를 나타냈다.

효율성 측면에서 ECHO는 8배의 추론 가속화를 달성했다. 기존 모델들이 초당 약 36~53개의 토큰을 생성할 때, ECHO는 초당 274.21개의 토큰을 생성(TPS)하는 놀라운 속도를 보여주었다. 또한 RAD 전략을 통해 학습 시 필요한 FLOPs를 72.3% 절감하면서도 성능 저하가 거의 없음을 확인했다.

Ablation Study 결과, DCD의 구성 요소인 단계별 토큰 재가중치(SW)와 토큰에 대한 교차 엔트로피 손실(CE)이 생성 안정성에 결정적인 역할을 함이 증명됐다. 특히 손실을 추가했을 때 Perplexity(PPL)가 23.72에서 18.83으로 낮아지며 문장의 종결 능력이 크게 개선되었다.

ECHO의 아키텍처는 Lingshu-7B를 백본으로 하며, 블록 단위로 토큰을 생성하는 Block Diffusion 구조를 채택한다. 핵심 차별점은 기존 확산 모델들이 토큰별로 독립적인 확률 분포를 가정하는 Mean-field Approximation을 사용하는 것과 달리, DCD를 통해 토큰 간의 조건부 의존성을 직접 학습한다는 점이다. 이는 수학적으로 Joint Posterior p(x0 | xt, t)를 직접 근사하는 것과 유사한 효과를 낸다.

학습 시 RAD(Response-Asymmetric Diffusion)는 시각적 특징(Vision Tokens)과 지시문(Instruction Tokens)을 조건부 컨텍스트로 고정하고, 오직 응답(Response) 부분에만 노이즈를 주입한다. 이때 블록 어텐션 마스크를 사용하여 각 노이즈 블록이 모든 시각 정보와 이전의 깨끗한 텍스트 블록을 참조할 수 있게 설계했다. 이는 연산량을 줄이면서도 강력한 조건부 생성을 가능하게 한다.

구현 세부사항으로는 2,250,000 픽셀 제한 내에서 고해상도 이미지를 처리하며, AdamW 옵티마이저와 1e-5의 학습률을 사용했다. 또한 'Reporting by exception' 관행으로 인한 데이터 편향을 해결하기 위해, 모든 해부학적 부위에 대해 긍정/부정 여부를 명시적으로 기록하는 데이터 정규화(Data Normalization) 과정을 거쳐 할루시네이션을 억제했다.

본 논문은 흉부 엑스레이(CXR) 도메인에 특화되어 있어, 다른 부위의 의료 영상이나 일반적인 멀티모달 작업으로의 확장 가능성에 대해서는 추가적인 검증이 필요하다. 또한 1단계 생성의 안정성을 위해 교사 모델의 품질에 크게 의존한다는 점이 한계로 작용할 수 있다.