PHASE2_NOTES.md

Phase-2 전환 노트 (DiT unfreeze) — 2026-07-16 (최종본, B 방식)

무엇이 바뀌었나

Phase-1: DiT trunk(453M) 완전 freeze, encoder(ViT)+attn-pool+cond만 학습. Phase-2: DiT unfreeze, encoder와 동일 lr로 함께 학습 (dit_lr_scale 없음).

LR 처리 (중요 — B 방식으로 확정)

Phase-1은 step380000에서 **100-epoch warmup의 38% 지점(lr≈9.5e-6)**이었음. Phase-2는 이 지점부터 원래 warmup 스케줄을 그대로 이어감:

• 시작 lr ≈ 9e-6 (step380000 값과 동일), warmup이 끝나며(step 1,000,800) 2.5e-5로 점점 도달.
• ❌ fresh 100ep warmup(0부터, full까지 48h) 아님 — 그건 낭비라 폐기.
• ❌ 3ep 단축 warmup(0→2.5e-5 급상승) 아님 — DiT init에 급격.
• ✅ step380000 LR 지점에서 자연스럽게 이어받아 원래 스케줄대로 상승.

어떻게 (구현)

1. 가중치: init_from이 accelerate ckpt 디렉토리도 받도록 확장 → step380000의 model.safetensors(encoder+pool+DiT+cond) 로드. (loss 연속성으로 검증: 시작 diff/repa가 Phase-1 종료값과 동일)
2. LR 스케줄 위치: trainer에 PHASE2_RESUME_STEP env 추가 → self.steps=380000으로 설정. loaded_steps=-1이라 데이터 fast-forward 없이(38 epoch 재읽기 회피) 스케줄러가 step_update(380000)으로 lr=9e-6 위치. warmup_epochs=100 유지.
3. optimizer: fresh (frozen→unfrozen이면 param group 134M→599M 변경 → accelerate 전체 load_state 불가). momentum만 새로, LR 위치는 정확.
4. ckpt 번호: step380000부터 이어가 step390000, 400000... 연속 저장 (재시작 아님).
5. trainable 134M(P1) → 599M(P2, DiT 포함). FID eval은 step400000(50k 배수)에서 첫 시작.

파일 (전부 이 repo code/ 에 업로드됨)

• spatial_diffuse_slot.py — 모델(SpatialAttnPool + DiTSpatial + init_from 디렉토리 지원)
• spatial_mask.py — spatial-align mask 빌더
• diffusion_trainer_PHASE2_PATCH.txt — trainer의 PHASE2_RESUME_STEP 패치 (적용 위치 명시)
• trainer_utils_create_optimizer.py — dit_lr_scale 지원 create_optimizer (B에선 미사용, 참고)
• tokenizer_l_spatial.yaml (Phase-1) / tokenizer_l_spatial_phase2.yaml (Phase-2)
• train_spatial_l.sh (P1) / train_spatial_l_phase2.sh (P2, PHASE2_RESUME_STEP env)
• hf_ckpt_watcher.py (P1) / hf_ckpt_watcher_phase2.py (P2) / make_loss_plot*.py

HF 네임스페이스 (이 repo)

• Phase-1(frozen, 보존): step360000~380000/, samples_all_steps/, loss_curves.png → step380000/은 Phase-2 init 소스라 반드시 유지 (resume에 필요)
• Phase-2(unfrozen): phase2_step390000+/, samples_all_steps_phase2/, loss_curves_phase2.png, logs_phase2/
• watcher가 최신 2개 유지 + squash(용량 회수). repo는 public(private 용량한도 회피).

resume (새 서버)

1. code/의 파일들 + semanticist repo에 배치, diffusion_trainer.py에 PHASE2 패치 적용.
2. Phase-2 이어가기: 최신 phase2_stepN/(HF)을 로컬에 두고 configs/tokenizer_l_spatial_phase2.yaml의 init_from을 그 경로로, PHASE2_RESUME_STEP=N으로 train_spatial_l_phase2.sh 실행.
3. Phase-2 처음부터: init_from=step380000 + PHASE2_RESUME_STEP=380000 (현재 설정).

Level drop 추가 (2026-07-16, step390000부터)

우리 multi-res 방식의 핵심인 level drop(nested) 를 Phase-2에 추가 (step390000에서 resume).

• 방식: whole-level, coarse-first. keep_levels ~ uniform(1,4) → {1x1},{1x1,2x2},{1x1,2x2,4x4},{all} 각 25%. → 레벨별 유지확률 1x1=100%, 2x2=75%, 4x4=50%, 8x8=25%.
• 왜 25%(레벨균등): 토큰비율(uniform-over-token) 방식은 8x8이 64/85라 coarse-only가 거의 학습 안 됨. 레벨균등은 모든 granularity를 고르게 학습 → 각 레벨이 독립적으로 쓸모있음 (multi-res reasoning 목적).
• 드롭 순서: 8x8(finest, DiT 이미지토큰과 1:1)부터 드롭, 1x1(global)은 항상 유지. 전부 드롭(uncond)은 CFG(uncond_drop_prob)가 따로 담당.
• 구현: LevelNestedSampler (spatial_diffuse_slot.py), config enable_nest: True.
• inference: inference_with_n_slots(토큰 budget)를 whole-level prefix로 매핑 (85→all, 21→8x8 drop, 5→[2x2,1x1], 1→[1x1]).