Regarding the specific number of training steps for each subset of data.

[wangxuanx]

Thank you for your excellent work. I have a small question after reading the original paper. The paper mentions training the IMIG-100k dataset on four A100 images with a batch size of 16 using 5k steps. My question is, is this 5k steps for all three subsets or for each individual subset? Thank you very much for your answer!

[nenhang]

Thank you for recognition and concern. We trained the model on the mixture of all three subsets. With LoRA Rank <= 512, 5k training steps are sufficient to allow the model to converge roughly, while scaling up the batch size and training steps would allow the model to see more different data, which may further enhance generalization ability.

[Whitney-9]

The total batch size is 16 and the total training steps are 5k. This means that the model has only seen 80,000 samples in total. Is the 100k dataset too large for this task? Will the training effect deteriorate if only some of the data are used?

[nenhang]

The total batch size is 16 and the total training steps are 5k. This means that the model has only seen 80,000 samples in total. Is the 100k dataset too large for this task? Will the training effect deteriorate if only some of the data are used?

The 5k steps in the paper were the training setting we used at that stage, with data being randomly sampled during training rather than traversed in a fixed order, mainly for the 512-resolution SFT version. Due to time limitations, we did not continue that version much further at that time. Under that setup (SFT, Prodigy optimizer, default LR), the model was close to convergence, before the later preference-tuning stage, but its generalization was not necessarily the best. So your intuition is reasonable: if training stops around that stage, the model may not fully benefit from the full 100k data distribution, especially for broader generalization.

The released checkpoint here:
https://huggingface.co/ruihangxu/ContextGen/blob/main/transformer_25000.safetensors
was trained for much longer (around 25k steps in total), and in our experience it generalizes better across different resolutions and scene types.