LoRA와 DoRA를 활용한 NVIDIA Cosmos Predict 2.5 로봇 비디오 생성 파인튜닝

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

NVIDIA Cosmos Predict 2.5는 물리적 세계의 역학을 이해하고 비디오를 생성하는 대규모 월드 모델이다. 특정 도메인 적응을 위해 LoRA와 DoRA를 활용한 파라미터 효율적 파인튜닝을 수행하여 메모리 요구량을 낮춘다. 학습 과정에서 DiT 가중치를 고정하고 어댑터 모듈만 최적화하여 단일 GPU에서도 학습이 가능하다. 실험 결과, 파인튜닝은 Sampson Error를 낮추고 물리적 타당성과 지시 이행 점수를 높여 로봇 조작 작업에 적합한 합성 데이터를 생성한다.

배경

Python 3.10+, PyTorch 2.5+, CUDA 환경, Diffusers 및 Accelerate 라이브러리

대상 독자

로봇 학습을 위한 합성 데이터 생성 및 LLM 파인튜닝에 관심 있는 AI 엔지니어

의미 / 영향

이 기술은 로봇 학습에 필요한 고비용의 실제 데이터 수집을 대체할 수 있는 확장 가능한 합성 데이터 생성 경로를 제공한다. 특히 제한된 GPU 자원으로도 특정 도메인에 최적화된 월드 모델을 구축할 수 있게 하여 로봇 연구의 진입 장벽을 낮춘다.

섹션별 상세

Cosmos Predict 2.5는 텍스트와 이미지를 입력받아 물리적으로 타당한 비디오를 생성하는 월드 모델이다.

LoRA와 DoRA 어댑터를 DiT의 어텐션 및 피드포워드 레이어에 주입하여 전체 파라미터 학습 없이 효율적으로 도메인 적응을 수행한다.

python

lora_config = LoraConfig(
    r=args.lora_rank,
    lora_alpha=args.lora_alpha,
    target_modules=['to_q', 'to_k', 'to_v', 'to_out.0', 'ff.net.0.proj', 'ff.net.2'],
    use_dora=args.use_dora, # set True to switch to DoRA
)
dit.add_adapter(lora_config)
cast_training_params(dit, dtype=torch.float32) # LoRA params in fp32

Cosmos Predict 2.5 모델에 LoRA 또는 DoRA 어댑터를 설정하고 주입하는 코드.

학습 루프는 Rectified Flow 공식을 사용하여 노이즈에서 데이터로의 속도를 예측하도록 최적화된다.

python

# Forward
pred_velocity = dit(
    hidden_states=xt,
    condition_mask=cond_mask,
    timestep=in_timestep,
    encoder_hidden_states=prompt_embeds,
    padding_mask=padding_mask,
    return_dict=False,
)[0]
# MSE loss is computed only on the non-conditioned frames
target_velocity = noise - clean_latent
pred_velocity = target_velocity * cond_mask + pred_velocity * (1 - cond_mask)
loss = F.mse_loss(pred_velocity.float(), target_velocity.float())

Rectified Flow 공식을 사용하여 학습 손실(MSE Loss)을 계산하는 핵심 루프.

Sampson Error를 통해 프레임 간 기하학적 일관성을 측정하고 LLM-as-a-Judge를 통해 물리적 타당성과 지시 이행 능력을 검증한다.

Temporal 및 Cross-View Sampson Error 비교 차트. — Chart파인튜닝 전후의 기하학적 오차를 비교한다. LoRA와 DoRA 적용 시 Sampson Error가 감소하여 프레임 간 및 뷰 간 일관성이 향상되었음을 보여준다.

물리적 타당성(Adherence to Physical Commonsense) 점수 비교 차트. — Chart파인튜닝된 모델이 기본 모델보다 물리적 상식을 더 잘 준수함을 보여준다. LoRA와 DoRA 모두 물리적 타당성 점수를 향상시킨다.

지시 이행(Adherence to Prompt) 점수 비교 차트. — Chart파인튜닝된 모델이 프롬프트 지시를 더 정확하게 수행함을 보여준다. 특히 rank 32 설정에서 더 높은 점수를 기록한다.

실험 결과, rank 32 설정이 지시 이행 능력을 높이며, DoRA와 LoRA는 유사한 성능을 보이나 DoRA는 학습 안정성에 기여할 수 있다.

근거

학습은 100 에포크(8x H100 기준 2.5시간)만으로도 모든 지표를 크게 개선한다. — Conclusion 섹션
rank 32 설정은 지시 이행 능력을 높이지만 기하학적 일관성이나 물리적 타당성은 크게 개선하지 않는다. — Conclusion 섹션

용어 해설

Rectified Flow: — 노이즈 샘플을 데이터로 변환하는 경로를 직선화하여 학습하는 생성 모델링 방식. 모델은 노이즈와 데이터 사이의 속도를 예측하도록 학습되며, 샘플링 효율성과 품질을 높인다.
Sampson Error: — 매칭된 키포인트와 에피폴라 라인 사이의 거리를 측정하는 기하학적 오차 지표. 비디오 생성에서 프레임 간 혹은 뷰 간의 기하학적 일관성을 평가하는 데 사용된다.
World Model: — 물리적 세계의 역학을 이해하고 미래 상태를 예측하는 AI 모델. 로봇 조작이나 자율 주행 등 물리적 상호작용이 필요한 작업에서 합성 데이터를 생성하는 데 활용된다.
PEFT: — 대규모 모델의 전체 가중치를 학습하지 않고 일부 파라미터나 어댑터 모듈만 학습하여 메모리 사용량을 줄이는 기법. LoRA와 DoRA가 대표적이다.

언급된 리소스

튜토리얼Cosmos Cookbook

GitHubCosmos Models and Datasets