TL;DR: 학습 없이, 보조 모델 없이, 내부 상태 접근 없이. 미리 계산한 토큰 점수를 매 디노이징 step의 logit에 더해, Diffusion 언어 모델(DLM)의 출력 스타일을 직접 조종한다.

📄 Paper: https://arxiv.org/pdf/2605.29626 💻 Code: https://github.com/hsannn/swai


1. 기존 steering은 왜 DLM 앞에서 멈추는가

언어 모델을 실제로 배포할 때는 내용이 맞는 것만으로 부족하다. 독자 수준에 맞춘 글의 난이도, 사용자를 대하는 상황의 공손한 어투, 안전이 중요한 곳에서의 유해 표현 회피처럼 "무엇을 말하는가"가 아니라 *"어떻게 말하는가"*를 제어해야 한다. 이런 controllable generation은 이제 부가 기능이 아니라 핵심 역량이다.

그중 inference-time steering은 매력적이다. 베이스 모델을 재학습하지 않고 사후에 붙일 수 있고, 여러 속성에 재사용 가능하며, 파인튜닝보다 비용이 훨씬 낮다. 문제는 기존 기법 대부분이 autoregressive(AR) 모델을 전제한다는 점이다. 좌→우 next-token 디코딩을 가정하거나 보조 모델에 의존하는데, 이런 가정은 부분 마스킹된 시퀀스를 반복 디노이징하며 문장을 드러내는 Diffusion 언어 모델(DLM)에는 자연스럽게 옮겨오지 못한다.

image.png

DLM-SWAI는 위 세 조건을 모두 만족하는, 디자인 공간의 빈 자리를 채운다.


2. 핵심 아이디어: DLM-SWAI

DLM-SWAI(*Statistical Writing style Aligned Inference for Diffusion Language Models)*는 두 단계로 구성된다. (1) 오프라인에서 통계적 토큰 점수표를 구축하고, (2) 디코딩 중 매 디노이징 step의 logit에 그 점수를 더해 출력을 목표 특성 쪽으로 이동시킨다. 모델 파라미터와 내부 표현은 일절 건드리지 않는다.

image.png

Phase 1 — Statistical Score Table

각 토큰이 목표 특성을 얼마나 특징짓는지를 하나의 점수로 정량화한다.

  1. One-vs-rest log-odds (Dirichlet prior). 목표 특성 코퍼스와 나머지 전체를 대조해 각 토큰의 log-odds 비를 계산한다. "이 토큰이 목표 특성에서 얼마나 더 두드러지는가"를 나타내는 방향성 신호다. 저빈도 토큰의 수치 불안정성은 코퍼스 통합 빈도로 만든 Dirichlet prior(α)로 스무딩한다.
  2. Z-score 정규화. raw log-odds는 저빈도 토큰에서 노이즈에 취약하다. 추정 분산의 제곱근(표준오차 역할)으로 나눠 표준화하여, 단지 빈번할 뿐 아니라 통계적으로 신뢰할 수 있는 distinctive 토큰을 우선한다.

Phase 2 — Denoising-time Steering

각 디노이징 step에서 다음을 수행한다.

  1. Global steering bias. 점수 벡터에 element-wise clipping과 스케일링을 적용해 bias 벡터 $\mathbf{b}_k = \lambda \cdot \mathrm{clip}(\mathbf{s}_k, -\tau, \tau)$를 만든다 ($\lambda$=steering 강도, $\tau$=clipping 임계값). 이 벡터는 한 번 계산해 전체 디노이징 과정 내내 재사용한다.
  2. 모든 마스킹 위치에 주입. active block의 모든 마스킹 위치 logit에 동일한 bias를 더한다: $\tilde{\mathbf{z}}_i = \mathbf{z}_i + \mathbf{b}_k$. 그 보정된 분포에서 후보 토큰을 샘플링한다.