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]).