프로브, 컨디셔닝 채널, GRPO, DiT, dLLM 각각에 대해 의사코드와 학습 데이터 구조를 함께 살펴본다.
프로브
프로브는 모델의 특정 층에서 활성화값을 뽑아내고, 그 위에 작은 분류기를 얹어서 학습시키는 구조다. 의사코드로 보면 다음과 같다.
# 1단계: 활성화값 수집
def get_activation(model, input_text, layer_idx):
activations = {}
def hook(module, input, output):
activations['value'] = output.detach()
handle = model.layers[layer_idx].register_forward_hook(hook)
model.forward(input_text)
handle.remove()
return activations['value'] # shape: [hidden_dim]
# 2단계: 프로브 정의 (아주 작은 MLP 또는 선형 분류기)
class Probe(nn.Module):
def __init__(self, hidden_dim):
super().__init__()
self.linear = nn.Linear(hidden_dim, 1) # 정직/기만 이진 분류
def forward(self, activation):
return torch.sigmoid(self.linear(activation))
# 3단계: 프로브 학습
for batch in probe_dataset:
activation = get_activation(model, batch.text, layer_idx=20)
pred = probe(activation)
loss = binary_cross_entropy(pred, batch.label)
loss.backward()
probe_optimizer.step()
# 4단계: 훈련 중 보상에 반영
def reward_with_probe(response, activation):
task_reward = check_correctness(response)
honesty_score = probe(activation) # 0에 가까우면 기만 의심
return task_reward - penalty_weight * (1 - honesty_score)
학습 데이터 구조는 활성화값과 정직 여부 라벨의 쌍으로 이루어진다.
{
"example_id": "0001",
"prompt": "이 코드가 모든 테스트를 통과하는지 확인하고 수정하라",
"model_response": "def solve(): return hardcoded_answer",
"activation_layer20": [0.023, -1.14, 0.87, ...], # hidden_dim 크기 벡터
"label": "deceptive", # 또는 "honest"
"ground_truth_check": "test case hardcoded, actual logic broken"
}
라벨은 사람이 직접 코드를 검토해서 실제로 문제를 풀었는지, 아니면 테스트만 통과하도록 편법을 썼는지 판단해서 붙인다. 이 라벨과 그 순간의 활성화값을 짝지어 놓은 데이터셋으로 프로브를 학습시키면, 나중에는 라벨 없이 활성화값만 보고도 기만 여부를 예측할 수 있게 된다.
컨디셔닝 채널
컨디셔닝 채널은 여러 조건 정보를 같은 크기의 텐서로 맞춘 다음, 채널 축으로 이어 붙여서 모델에 입력하는 구조다.
# 각 조건을 동일한 공간 해상도의 텐서로 인코딩
noise_latent = get_noise_latent() # shape: [C, H, W]
text_embed = text_encoder(prompt) # shape: [C_text] -> 브로드캐스트하여 [C_text, H, W]
ref_image_feat = image_encoder(ref_image) # shape: [C_img, H, W]
pose_map = pose_estimator(pose_input) # shape: [C_pose, H, W]
audio_feat = audio_encoder(audio_clip) # shape: [C_audio] -> [C_audio, H, W]로 확장
# 채널 축으로 이어붙이기
combined_input = concat([noise_latent, text_embed, ref_image_feat, pose_map, audio_feat], dim=channel_axis)
# shape: [C_total, H, W]
# 기존 DiT 백본의 입력 채널 수만 늘려서 그대로 학습
output = diffusion_transformer(combined_input)
loss = diffusion_loss(output, target_video)
기존 모델이 입력 채널 하나(노이즈)만 받던 것을, 여러 조건 채널을 더해서 받도록 첫 번째 레이어의 입력 차원만 넓혀주면 나머지 트랜스포머 구조는 그대로 재사용할 수 있다.
학습 데이터 구조는 여러 조건이 하나의 샘플 안에 함께 묶여 있는 형태다.
{
"sample_id": "hoivg_00042",
"text_prompt": "한 사람이 상자를 들어 옮긴다",
"reference_image": "person_box_ref.png",
"pose_sequence": "pose_00042.npy", # 프레임별 관절 좌표
"audio_clip": "footstep_lift.wav",
"target_video": "target_00042.mp4" # 학습 정답이 되는 실제 영상
}
한 샘플 안에 텍스트, 참조 이미지, 포즈, 오디오, 정답 영상이 모두 짝지어져 있어야 하고, 학습 시에는 이 모든 조건을 인코딩해서 채널로 이어 붙인 다음 정답 영상과의 차이를 손실로 계산한다.
GRPO
GRPO는 같은 질문에 대해 여러 응답을 한꺼번에 샘플링하고, 그 응답들 사이의 상대적 보상 차이로 정책을 업데이트하는 구조다.
for prompt in training_prompts:
# 1단계: 같은 프롬프트로 여러 응답 샘플링 (그룹)
responses = [model.generate(prompt) for _ in range(group_size)] # 예: 8개
# 2단계: 각 응답에 보상 부여
rewards = [reward_function(prompt, r) for r in responses]
# 3단계: 그룹 내 평균과 표준편차로 정규화 (상대적 우열 계산)
mean_r = mean(rewards)
std_r = std(rewards)
advantages = [(r - mean_r) / (std_r + eps) for r in rewards]
# 4단계: 정책 업데이트 (상대적으로 좋은 응답의 확률은 높이고 나쁜 응답은 낮춤)
for response, advantage in zip(responses, advantages):
log_prob = model.log_prob(response, prompt)
loss = -advantage * log_prob
loss += kl_penalty(model, reference_model, prompt, response) # 원래 모델에서 너무 벗어나지 않도록 정칙화
loss.backward()
optimizer.step()
별도의 가치 신경망 없이, 그룹 안에서의 평균과 표준편차만으로 각 응답이 상대적으로 얼마나 좋은지를 계산하는 것이 핵심이다.
학습 데이터 구조는 정답이 하나로 고정되어 있지 않고, 프롬프트와 그 프롬프트에 대한 평가 기준만 있으면 된다.
{
"prompt": "최근 알츠하이머 치료제 연구 동향을 조사해 정리하라",
"rubric": [
{"criterion": "최근 2년 내 임상시험 결과를 포함했는가", "weight": 0.3},
{"criterion": "상반된 연구 결과를 균형있게 다뤘는가", "weight": 0.3},
{"criterion": "출처를 명확히 인용했는가", "weight": 0.2},
{"criterion": "결론이 근거에 기반해 신중하게 서술되었는가", "weight": 0.2}
]
}
정답 텍스트가 미리 정해져 있는 것이 아니라, 응답을 채점할 루브릭만 있으면 된다. 모델이 같은 질문에 대해 여러 답변을 만들면, 이 루브릭으로 각 답변을 채점해서 상대적 점수를 매기고, 그 점수 차이를 이용해 정책을 업데이트한다. DR Tulu는 이 루브릭 자체도 훈련이 진행되면서 점점 더 정교해지도록 함께 학습시킨다는 점이 추가된 특징이다.
DiT
DiT는 디퓨전 모델, 즉 노이즈를 조금씩 걷어내면서 이미지나 영상을 만들어내는 생성 방식에서, 노이즈를 예측하는 핵심 신경망을 U-Net 대신 트랜스포머로 바꾼 구조다. 원래 디퓨전 모델은 U-Net을 많이 썼는데, U-Net은 이미지를 점점 작게 압축했다가 다시 키우는 구조라 합성곱 연산에 의존한다. DiT는 이미지를 작은 패치로 잘게 잘라서 마치 문장의 단어처럼 하나의 토큰 시퀀스로 취급하고, 그 토큰들을 트랜스포머의 셀프 어텐션으로 처리한다. 비유하자면 그림을 그릴 때 U-Net은 전체 캔버스를 계속 확대 축소하며 다듬는 화가에 가깝고, DiT는 그림을 작은 조각으로 잘라 붙인 모자이크 타일들을 서로 비교하고 조율하며 전체 그림을 완성해가는 방식에 가깝다.
DiT가 기존 U-Net 방식보다 좋은 이유는 두 가지다. 첫째, 트랜스포머는 이미지 크기나 데이터 양이 늘어날수록 성능이 꾸준히 좋아지는 확장성이 검증된 구조라서, 모델과 데이터를 키울수록 U-Net보다 더 안정적으로 품질이 좋아진다. 둘째, 토큰 시퀀스라는 형태를 쓰기 때문에 앞서 설명한 컨디셔닝 채널처럼 텍스트, 이미지, 포즈 같은 다른 종류의 조건 정보를 같은 토큰 시퀀스 안에 자연스럽게 섞어 넣기가 쉽다.
예시
앞서 다룬 Causal Forcing과 OmniShow가 모두 이 DiT 구조를 베이스 모델로 쓴다. OmniShow의 경우 영상의 각 프레임을 패치 단위로 잘라 토큰으로 만들고, 여기에 텍스트와 참조 이미지, 포즈, 오디오에서 나온 조건 토큰들을 함께 섞어서 트랜스포머에 입력한다. 노이즈가 낀 영상 토큰과 조건 토큰들이 셀프 어텐션 안에서 서로 영향을 주고받으면서, 조건에 맞는 방향으로 노이즈를 걷어내는 법을 학습한다.
의사코드
# 1단계: 이미지 또는 영상 프레임을 패치로 분할하고 토큰화
def patchify(image, patch_size=16):
patches = split_into_patches(image, patch_size) # 예: 256x256 -> 16x16 패치 256개
tokens = linear_projection(patches) # 각 패치를 하나의 벡터 토큰으로 변환
tokens = tokens + positional_embedding
return tokens # shape: [num_patches, hidden_dim]
# 2단계: 노이즈 스텝 정보를 임베딩으로 만들어 조건으로 사용
def timestep_embedding(t):
return sinusoidal_encode(t) # 스칼라 t를 hidden_dim 크기 벡터로 변환
# 3단계: DiT 블록 (트랜스포머 블록 + 조건 주입)
class DiTBlock(nn.Module):
def forward(self, tokens, condition_embed):
# 조건 임베딩으로 정규화 파라미터를 조절 (adaptive layer norm)
scale, shift, gate = condition_to_scale_shift_gate(condition_embed)
x = adaptive_layer_norm(tokens, scale, shift)
x = x + gate * self.self_attention(x)
x = adaptive_layer_norm(x, scale, shift)
x = x + gate * self.mlp(x)
return x
# 4단계: 전체 forward, 노이즈 예측
def dit_forward(noisy_tokens, t, text_condition):
condition_embed = timestep_embedding(t) + text_encoder(text_condition)
x = noisy_tokens
for block in dit_blocks: # 예: 28개 블록
x = block(x, condition_embed)
predicted_noise = unpatchify(final_linear(x))
return predicted_noise
# 5단계: 학습 루프
for batch in dataset:
clean_tokens = patchify(batch.image)
t = random_timestep()
noise = sample_gaussian_noise()
noisy_tokens = add_noise(clean_tokens, noise, t) # 노이즈 스케줄에 따라 섞기
predicted_noise = dit_forward(noisy_tokens, t, batch.text_prompt)
loss = mse_loss(predicted_noise, noise)
loss.backward()
optimizer.step()
핵심은 이미지를 패치 토큰으로 바꾼 다음, 노이즈 스텝 정보를 트랜스포머 블록 안의 정규화 파라미터를 조절하는 방식으로 주입한다는 점이다. 이렇게 하면 같은 트랜스포머 블록이 노이즈가 많이 낀 초기 단계와 노이즈가 거의 없는 마지막 단계 모두에서 서로 다르게 동작하도록 조절된다.
학습 데이터 구조
{
"sample_id": "video_gen_00187",
"text_prompt": "노을이 지는 해변을 걷는 사람",
"clean_video_frames": "target_00187.mp4", # 노이즈가 없는 정답 영상
"resolution": [256, 256],
"num_frames": 16,
"patch_size": 16
}
학습 시에는 정답 영상을 패치로 잘라 토큰으로 만든 다음, 무작위로 고른 노이즈 스텝만큼 인위적으로 노이즈를 섞고, 모델이 그 노이즈를 예측하도록 학습시킨다. 즉 데이터셋 자체에는 노이즈가 없는 깨끗한 영상과 텍스트 프롬프트만 있으면 되고, 노이즈는 학습 과정에서 그때그때 만들어서 씌운다. 조건이 여러 개인 OmniShow 같은 경우는 여기에 참조 이미지, 포즈, 오디오 항목이 앞서 설명한 컨디셔닝 채널 구조로 함께 추가된다.
dLLM
기존 LLM은 문장을 왼쪽에서 오른쪽으로 한 단어씩 순서대로 예측하는 자기회귀 방식을 쓴다. 반면 dLLM은 이미지에 노이즈를 씌웠다가 걷어내는 디퓨전의 원리를 텍스트에 그대로 적용한다. 문장의 토큰들을 무작위로 마스킹, 즉 가려버린 다음, 모델이 가려진 토큰들이 무엇이었는지 한꺼번에 맞히도록 학습시킨다. 비유하자면 기존 LLM이 문장을 왼쪽부터 한 글자씩 순서대로 받아쓰는 방식이라면, dLLM은 빈칸 뚫기 문제지를 받아서 뚫린 빈칸들을 한 번에, 그리고 여러 번 다시 고쳐가며 채워나가는 방식이다.
구조적으로 보면 dLLM은 이미지 패치 토큰 대신 텍스트 토큰을 입력으로 쓴다는 점만 다를 뿐, 트랜스포머 블록 자체는 DiT와 거의 동일하다. 다만 이미지처럼 노이즈를 연속적인 값으로 섞는 대신, 텍스트는 이산적인 토큰이기 때문에 토큰을 마스크 토큰으로 치환하는 방식으로 노이즈를 흉내낸다. 마스킹 비율이 노이즈 스텝 역할을 한다. 마스킹 비율이 높으면 노이즈가 많이 낀 초기 단계이고, 마스킹 비율이 낮으면 노이즈가 거의 걷힌 마지막 단계다.
dLLM이 기존 자기회귀 LLM보다 나은 점은, 문장을 왼쪽부터 순서대로 만들 필요가 없어서 여러 토큰을 동시에 병렬로 생성할 수 있다는 것이다. 앞서 다룬 Any-Order GPT as Masked Diffusion Model 논문에서 언급했듯이 이 병렬 생성 덕분에 추론 속도가 최대 25배까지 빨라질 수 있다. 다만 앞서 다룬 The Flexibility Trap 논문에서 지적했듯이, 어떤 순서로든 생성할 수 있다는 유연성이 오히려 모델을 추론에 중요한 어려운 토큰을 회피하게 만드는 부작용도 있다.
예시
수학 문제를 푸는 상황을 생각해보면, 기존 자기회귀 LLM은 "2 더하기 3은" 다음에 "5" 하나를 예측하고, 그다음 토큰을 예측하는 식으로 왼쪽부터 순서대로 답을 완성해간다. dLLM은 전체 답변 자리에 마스크 토큰을 전부 채워놓은 상태에서 시작해서, 처음에는 확신이 높은 몇몇 자리, 예를 들어 숫자 "5"부터 채우고, 그다음 라운드에서 나머지 마스크된 자리들을 점점 채워나가며 여러 번에 걸쳐 문장 전체를 완성해간다.
의사코드
# 1단계: 텍스트를 토큰화하고 일부를 마스크 토큰으로 치환
def mask_tokens(token_ids, mask_ratio):
num_to_mask = int(len(token_ids) * mask_ratio)
mask_positions = random_sample(range(len(token_ids)), num_to_mask)
noisy_tokens = token_ids.clone()
for pos in mask_positions:
noisy_tokens[pos] = MASK_TOKEN_ID
return noisy_tokens, mask_positions
# 2단계: 마스킹 비율을 노이즈 스텝처럼 사용해 조건으로 임베딩
def mask_ratio_embedding(mask_ratio):
return sinusoidal_encode(mask_ratio) # DiT의 timestep_embedding과 동일한 방식
# 3단계: dLLM 블록 (DiT 블록과 구조 동일, 입력만 텍스트 토큰)
class DLLMBlock(nn.Module):
def forward(self, tokens, condition_embed):
scale, shift, gate = condition_to_scale_shift_gate(condition_embed)
x = adaptive_layer_norm(tokens, scale, shift)
x = x + gate * self.self_attention(x) # 양방향 어텐션, 인과적 마스크 없음
x = adaptive_layer_norm(x, scale, shift)
x = x + gate * self.mlp(x)
return x
# 4단계: 전체 forward, 마스크된 토큰 예측
def dllm_forward(noisy_tokens, mask_ratio):
condition_embed = mask_ratio_embedding(mask_ratio)
x = token_embedding(noisy_tokens)
for block in dllm_blocks:
x = block(x, condition_embed)
predicted_token_logits = final_linear(x) # 각 위치마다 vocab 크기의 분포
return predicted_token_logits
# 5단계: 학습 루프
for batch in text_dataset:
clean_tokens = tokenize(batch.text)
mask_ratio = random_uniform(0, 1) # 이미지의 random_timestep과 동일한 역할
noisy_tokens, mask_positions = mask_tokens(clean_tokens, mask_ratio)
predicted_logits = dllm_forward(noisy_tokens, mask_ratio)
# 마스킹된 위치에서만 손실 계산
loss = cross_entropy(predicted_logits[mask_positions], clean_tokens[mask_positions])
loss.backward()
optimizer.step()
# 6단계: 생성 시 여러 라운드에 걸쳐 마스크를 점점 채워나감
def generate(prompt_tokens, answer_length, num_rounds=8):
tokens = concat([prompt_tokens, [MASK_TOKEN_ID] * answer_length])
for round in range(num_rounds):
mask_ratio = 1.0 - (round / num_rounds)
logits = dllm_forward(tokens, mask_ratio)
probs, predicted = softmax_and_argmax(logits)
# 확신이 높은 자리부터 확정하고 나머지는 다음 라운드로
confident_positions = top_k_by_confidence(probs, k=answer_length // num_rounds)
for pos in confident_positions:
tokens[pos] = predicted[pos]
return tokens
DiT와 비교했을 때 가장 큰 구조적 차이는 셀프 어텐션이 인과적 마스크 없이 양방향으로 열려 있다는 점이다. 기존 자기회귀 LLM은 지금 토큰이 뒤에 나올 토큰을 미리 보지 못하도록 마스킹하지만, dLLM은 마스킹된 자리 이외의 모든 토큰을 양방향으로 참고할 수 있어야 빈칸을 제대로 채울 수 있기 때문이다.
학습 데이터 구조
{
"sample_id": "text_00512",
"clean_text": "2 더하기 3은 5이고, 5에 4를 더하면 9이다.",
"token_ids": [1023, 892, 55, 1104, ...],
"vocab_size": 32000,
"max_seq_len": 512
}
DiT의 학습 데이터가 노이즈 없는 깨끗한 영상만 있으면 되는 것처럼, dLLM의 학습 데이터도 마스킹 없는 원본 텍스트만 있으면 된다. 마스킹은 DiT에서 노이즈를 학습 중에 그때그때 씌우는 것과 마찬가지로, 학습 루프 안에서 무작위 비율로 그때그때 만들어서 적용한다. 즉 데이터셋 준비 단계에서는 기존 자기회귀 LLM과 동일한 순수 텍스트 코퍼스를 그대로 쓸 수 있고, 학습 방식만 마스킹과 병렬 예측으로 바뀌는 것이다.
