역증류 정규화로 마스크 확산 모델 보상 튜닝을 소수 스텝으로 단축한다

무엇인가

마스크 확산 모델(MDM)은 토큰을 병렬로 언마스킹해 시퀀스를 생성하므로 언어·생물학적 서열·이산 이미지 표현에 쓰인다. 문제는 두 가지다. 반복 샘플링이 비싸고, 시퀀스 우도가 모든 디노이징 궤적을 주변화해야 해서 계산 불가능하다. 보상만 최적화하면 보상이 측정하지 않는 성질이 파괴되므로 사전학습 참조 모델로의 KL 정규화가 표준 안전장치인데, MDM에서는 그 KL을 직접 쓸 수 없다. 기존 연구는 우도 대리값과 증거 하한을 쓰거나 역방향 전이 사이의 KL을 벌점으로 주는데, 후자는 롤아웃을 따라 참조 모델을 평가해야 하고 정규화 항이 참조의 샘플링 경로에 묶인다는 한계가 있다.

어떻게 동작하나

IDRF의 핵심은 역증류(inverse distillation)를 정규화 항으로 재해석하는 것이다. 참조 디노이저를 고정한 채 학생 분포 위에서 참조의 디노이징 손실과 그 분포에 가장 잘 맞춘 보조 디노이저의 손실을 비교해, 그 차이를 벌점으로 쓴다. 논문의 명제 1은 보조 디노이저가 최적일 때 이 모집단 역증류 손실이 참조 분포로의 시퀀스 수준 KL 발산을 상한한다는 것을 증명한다. 시퀀스 길이와 학생 샘플러의 형태에 관계없이 성립하며, 참조가 자기 엔드포인트 분포의 베이즈 디노이저일 때 등식이 된다. 실전 목적함수는 학생이 자기 롤아웃에서 만든 예측을 소프트 타깃으로 삼아 참조와 보조 디노이저 각각의 NELBO를 평가하는 min-max 형태이고, 롤아웃 상태와 손상 문맥은 고정한 채 학생의 소프트 타깃 확률로만 기울기를 흘리는 세미그래디언트로 갱신한다. 참조 모델의 우도도, 참조 롤아웃도, 사전학습 데이터도, 보상의 기울기도 필요하지 않다.

무엇과 다른가

보상 최적화는 소수 스텝 생성기를 유한 지평 마르코프 결정 과정으로 본다. 1=t_K에서 t_0=0까지의 고정 격자에서 각 스텝마다 마스킹된 위치를 확률 b_k로 공개하고 값을 학생 디노이저에서 뽑으며, 이미 공개된 위치는 복사한다. 상태는 (스텝, 부분 마스킹 시퀀스), 행동은 다음 스텝의 시퀀스, 보상은 마지막 스텝에서만 받는 R(x_0)이다. 전이가 결정론적이고 b_k가 파라미터에 의존하지 않으므로 정책 경사는 각 스텝에서 새로 공개된 토큰에만 걸린다. 실제 학습은 M개의 롤아웃에 대해 클리핑된 leave-one-out 대리 목적함수를 최적화하며, 어드밴티지는 해당 샘플의 보상에서 나머지 롤아웃 보상 평균을 뺀 값이다. 정규화 항은 보상용으로 뽑은 롤아웃을 그대로 재사용하고, 궤적을 따라 스텝 j를 무작위로 골라 학생이 실제로 방문하는 상태에서 평가한다. 비용은 학생 평가 1회에 참조·보조 모델 순전파 각 1회가 추가되는 수준이다.

어떻게 쓰나

DNA 실험에서는 DRAKES를 따라 200bp 조절 서열을 생성하고 HepG2 세포의 인핸서 활성을 최적화했다. IDRF는 λ_ID=0.2, K=8 스텝으로 샘플링해 ancestral 디코딩에서 ATAC-Acc 95.3%로 표에서 최고를 기록했고(SDPO 94.8%), max-confidence 변형은 활성 7.31로 두 번째로 높았다. 비교 베이스라인은 모두 128 스텝을 쓰므로 IDRF가 16배 적다. 활성만 보면 SEPO+GF가 7.66으로 더 높지만 ATAC-Acc가 47.4%에 그친다. 정규화 항을 제거하면 같은 8 스텝에서 활성은 7.68로 오르지만 ATAC-Acc가 37.8%로 떨어지고 3-mer 상관이 0.68에서 0.38로, 엔트로피가 1.04로 무너지며 CA/CG 편중 서열이 된다. 128 스텝 샘플러에서는 붕괴가 완전해져 ATAC-Acc 0.1%가 되는데, 이 붕괴 모델들이 오히려 더 높은 추정 우도(-204, -163)를 받았다. λ_ID를 쓸어보면 활성은 단조 감소하고 ATAC-Acc는 37.8%에서 95.3%까지, 3-mer 상관은 0.38에서 0.68까지 올라 하나의 가중치로 동작점을 고를 수 있다.

전제와 한계

이미지에서는 ImageNet으로 사전학습된 클래스 조건부 MaskGIT를 CLIP 이미지-텍스트 유사도로 파인튜닝했고, 10개 클래스에 속성 프롬프트를 붙여 K=4 스텝, 1,000장으로 평가했다. IDRF는 4 스텝에서 classifier-free guidance(w=3)를 쓴 128 스텝 참조 모델을 ImageReward, PickScore, HPSv2.1 모든 홀드아웃 지표에서 앞섰고 ImageReward를 -0.73에서 0.88로 올렸다. 정규화 항을 빼면 CLIP 26.12, ImageReward -0.49로 떨어지고 PickScore와 HPSv2.1에서 참조보다 낮아진다. ReDGE는 추론 시 이미지별 최적화로 CLIP은 가장 높지만 홀드아웃 지표에서는 IDRF보다 낮다. 텍스트에서는 EOS 토큰을 포함해 Amazon Fashion 리뷰로 학습한 MDLM을 동결된 BERT 분류기 아래에서 log P(1 star) 최대화로 파인튜닝했다. 정규화 없는 RLOO는 보상은 가장 높지만 1,000개 중 906개가 EOS를 내지 않고 127 토큰 한계까지 가며 정규화 엔트로피가 실제 리뷰 0.95에 비해 0.72로 떨어진다. GenPPL만 보면 RLOO가 참조보다 낮은 혼란도(76.9 대 91.7)를 받아 이 실패를 잡아내지 못한다. IDRF는 16·32 스텝 모두 모든 샘플이 EOS를 내고 정규화 엔트로피 0.90~0.92를 유지하며, 32 스텝에서 보상은 RLOO보다 약간 낮지만(-0.044 대 -0.008) 실제 1-star 리뷰까지의 GM 거리가 2.6배 작고(10.9 대 28.8) 128 스텝 참조보다 낮은 GenPPL을 낸다.

실무 관점에서 이 논문이 주는 시사점은 명확하다. 보상 모델이 미분 불가능하거나 참조 모델 롤아웃 비용이 부담스러운 상황에서, 소수 스텝 생성기를 그대로 유지한 채 보상과 정규화를 같은 롤아웃에서 동시에 계산할 수 있다. 다만 정규화 항의 가중치 λ_ID가 보상과 품질 사이의 실질적 조작 변수이므로, DNA 사례처럼 보상만 보면 좋아 보이는 체크포인트가 분포 붕괴를 숨기고 있을 수 있다는 점을 전제로 홀드아웃 지표를 반드시 함께 봐야 한다. 또한 정규화를 완성된 학생 샘플에 순방향 마스킹으로 적용하면 DNA 설계에서 ATAC-Acc가 95.3%에서 52.5%로 떨어졌다는 절제 실험은, 정규화를 어느 상태에서 평가하느냐가 성능을 크게 가른다는 실무적 경고다.

저자들이 밝힌 한계는 다음과 같다. IDRF는 학생과 보조 디노이저 사이의 적대적 min-max 게임을 포함하므로 다른 적대적 학습처럼 최적화 불안정성과 두 갱신 사이 균형에 대한 민감도가 생길 수 있고, 보조 디노이저가 학습 비용을 추가한다. 이론적으로도 실전 대리 목적함수는 학생 롤아웃을 따라가는 원스텝 완성 분포들의 혼합에 대한 KL 보장만 제공하며, 최종 샘플링 분포 자체를 통제하는 것은 미해결 문제로 남는다. 향후 과제로 더 큰 모델로의 확장과 다른 도메인 적용을 제시한다.