flow — task-free decoder, aggregate-prior encoder

Conditional generation with a task-free decoder: all conditioning lives in an encoder that maps a prompt x and noise n to a latent w, whose aggregate (marginal) distribution is constrained to be standard Gaussian, W ~ N(0, I). The decoder is a single deterministic transport step (a frozen rectified flow) that takes no conditioning input. Inference is one encoder pass plus one Euler step.

Training is two-phase on a single joint law p(X, Y):

• Phase 1 trains and freezes the decoder T as a flow N(0, I) → p(Y).
• Phase 2 trains the encoder E and a train-time inverse encoder Q (which couples n to y), with a coherence loss ||T(w) − y||², an I(N; X) independence penalty, and SIGReg on n. Q is discarded at inference.

The full theory, derivations, and experiment write-ups are in latex/flow.tex (compiled: latex/flow.pdf).

Repository layout

oneflow/
├── latex/                  paper (flow.tex, flow.bib, flow.pdf)
├── experiments/
│   ├── exp01_toy2d_obsolete/   early 2-D toy (superseded; kept for history)
│   ├── exp02_templates/        template toy, initial run (exposed a sampler bug)
│   ├── exp03_templates/        corrected template toy + diagnostics (S1/S2 lessons)
│   └── exp04_coco_clip/        real text-to-image: COCO + SD-VAE + CLIP-L (RTX 4090)
├── README.md               ← you are here
└── .gitignore

Experiments

Dir	What	Status
exp01_toy2d_obsolete	2-D toy, first pass	obsolete
exp02_templates	template toy (binary-mask mixture)	done; sampler-bug lesson
exp03_templates	corrected template toy, full diagnostic battery	done; motivates S1 (JSD) + S2 (||T(w)−y||²)
exp04_coco_clip	real T2I on MS-COCO 256² latents + CLIP-L text	ready to run on GPU

The toy experiments (exp02, exp03) are single-file (main.py, env-var configured) and run on CPU/MPS. exp04_coco_clip is the GPU target and has its own detailed guide: experiments/exp04_coco_clip/README.md.

Quickstart (GPU server, exp04)

cd experiments/exp04_coco_clip
python3.11 -m venv .venv && source .venv/bin/activate
pip install --upgrade pip && pip install -r requirements.txt

bash setup_download.sh        # MS-COCO 2017 (~20 GB)
python encode.py --split train && python encode.py --split val
python encode.py --verify

SMOKE=1 python main.py        # ~45 min end-to-end smoke test, then full run

See the exp04 README for the full Phase 1 / 2A / 2B / eval command sequence, logging, resume/continue, and reproducibility notes.

Not in the repo (gitignored)

These are large and machine-specific — regenerate them locally:

• .venv/ — Python virtualenv. Recreate with the exp04 requirements.txt (tested: Python 3.11, CUDA 12.1, PyTorch 2.4).
• *.pt / *.ckpt / *.safetensors — model checkpoints and cached encoded tensors. Produced by training (main.py) and by encode.py.
• **/coco2017/ — raw dataset. Re-download with setup_download.sh.
• Pretrained weights (SD-VAE, CLIP-L) auto-download to ~/.cache/huggingface/ on first use; nothing to commit.
• LaTeX build artifacts (*.aux, *.log, …); flow.tex/flow.bib/flow.pdf are kept.