생성 모델은 공간을 안다: Video Generation Model의 implicit 3D prior를 3D Scene Understanding에 활용하기

# VEGA-3D 핵심 로직 의사코드
import torch
import torch.nn as nn

class AdaptiveGatedFusion(nn.Module):
    def __init__(self, d_llm):
        super().__init__()
        self.gate = nn.Linear(d_llm * 2, 1)
        self.ln_gen = nn.LayerNorm(d_llm)
        self.ln_sem = nn.LayerNorm(d_llm)

    def forward(self, F_gen, F_sem):
        # F_gen, F_sem: [T, N, D_llm]
        concat = torch.cat([self.ln_gen(F_gen), self.ln_sem(F_sem)], dim=-1)
        g = torch.sigmoid(self.gate(concat))  # [T, N, 1]
        return (1 - g) * F_gen + g * F_sem


def extract_generative_prior(video_frames, wan_model, timestep_k=300, layer_idx=20):
    """
    frozen Wan2.1-T2V에서 중간 피처 추출
    video_frames: [T, H, W, 3]
    """
    with torch.no_grad():
        # 1. VAE 인코딩
        z0 = wan_model.vae.encode(video_frames)  # clean latent
        
        # 2. 중간 타임스텝에서 노이즈 주입 (Flow Matching)
        tk = timestep_k / 1000.0
        eps = torch.randn_like(z0)
        zk = (1 - tk) * z0 + tk * eps  # noised latent
        
        # 3. 빈 텍스트 프롬프트로 DiT 피처 추출 (의미 정보 최소화)
        features = wan_model.dit.get_intermediate_feature(
            zk, timestep=timestep_k, text_cond="", layer=layer_idx
        )  # [T, N, D_gen]
    
    return features


# 학습 루프 개요
for frames, query, answer in dataloader:
    # Semantic branch (SigLIP 등)
    sem_feats = semantic_encoder(frames)       # [T, N, D]
    F_sem = sem_projector(sem_feats)           # [T, N, D_llm]
    
    # Generative branch (frozen Wan2.1)
    gen_feats = extract_generative_prior(frames, wan_model)
    F_gen = gen_projector(gen_feats)           # [T, N, D_llm]
    
    # Adaptive Gated Fusion
    fused = gated_fusion(F_gen, F_sem)         # [T, N, D_llm]
    
    # LLM 입력
    loss = llm(query, visual_tokens=fused, target=answer)