When turning a pre-trained model into an LSC model, how to decide when to start stage 2 training?

[Darkbblue]

Hello! This is a great study. I find HunyuanDiT especially stable compared to DiTs without LSCs, but this model is a bit outdated, so I want to add LSCs to other DiTs myself.
I managed to add some LSC adapters to a pre-trained DiT. Interestingly, I notice that although the generated image is largely corrupted, I can still observe some patterns of the overall object contour.
Anyway, my major concern now is how to decide when to start the stage 2 training, a.k.a. when you train the entire model. Should I wait until the generated image quality is almost restored before going into stage 2?

[guanjiechen118]

Hi! Thank you for your interest and agreement regarding the stability of Skip-DiT.

For the two-stage training of t2v-Latte, we determine the transition to the second stage by visually monitoring the generated samples. The training progression typically follows these stages:

1. Initially, the model generates only garbled outputs.
2. Gradually, it begins to produce recognizable object outlines, though colors remain inaccurate or distorted.
3. After approximately one full day of training (using 8 H100 GPUs), the model generates a more recognizable and properly colored first frame, although motion patterns are unnatural.

We initiate the second stage once the first frame appears visually consistent and realistic. As a gentle suggestion, we recommend also monitoring the increase in the weight and bias norms of the Linear layers in the LSCs to avoid over-training at the first stage.

Maybe you can post the generated images here, and I will try my best to help you.

[Darkbblue]

Hi! Thank you for your interest and agreement regarding the stability of Skip-DiT.

For the two-stage training of t2v-Latte, we determine the transition to the second stage by visually monitoring the generated samples. The training progression typically follows these stages:

1. Initially, the model generates only garbled outputs.
2. Gradually, it begins to produce recognizable object outlines, though colors remain inaccurate or distorted.
3. After approximately one full day of training (using 8 H100 GPUs), the model generates a more recognizable and properly colored first frame, although motion patterns are unnatural.

We initiate the second stage once the first frame appears visually consistent and realistic. As a gentle suggestion, we recommend also monitoring the increase in the weight and bias norms of the Linear layers in the LSCs to avoid over-training at the first stage.

Maybe you can post the generated images here, and I will try my best to help you.

Thanks for your quick reply! I guess I'm still in the early stage, as all outputs are random colors. I'll wait to see if it can converge. At least the loss is decreasing, a good sign.

[Darkbblue]

Hi! Thank you for your interest and agreement regarding the stability of Skip-DiT.

For the two-stage training of t2v-Latte, we determine the transition to the second stage by visually monitoring the generated samples. The training progression typically follows these stages:

1. Initially, the model generates only garbled outputs.
2. Gradually, it begins to produce recognizable object outlines, though colors remain inaccurate or distorted.
3. After approximately one full day of training (using 8 H100 GPUs), the model generates a more recognizable and properly colored first frame, although motion patterns are unnatural.

We initiate the second stage once the first frame appears visually consistent and realistic. As a gentle suggestion, we recommend also monitoring the increase in the weight and bias norms of the Linear layers in the LSCs to avoid over-training at the first stage.

Maybe you can post the generated images here, and I will try my best to help you.

Hey, I wonder how do you set the learning rates in stage 2. Do lsc connector and the backbone share the same lr?

[guanjiechen118]

Yes, we train the long-skip-connections using the same learning rate as the backbone parameters. In practice, we do not perform any learning-rate-selection experiments, as all our training experiments converge smoothly. The same learning rate is consistently used in both Stage 1 and Stage 2. We hope this information is helpful.

[Darkbblue]

Yes, we train the long-skip-connections using the same learning rate as the backbone parameters. In practice, we do not perform any learning-rate-selection experiments, as all our training experiments converge smoothly. The same learning rate is consistently used in both Stage 1 and Stage 2. We hope this information is helpful.

Some updates: I've trained stage 1 smoothly. I can see clearly the validation sample grow from noises to recognizable objects, though the quality is low. Then I started stage 2 training from a checkpoint. But after about 2500 optimizer steps the validation sample still has roughly the same quality as the checkpoint it started from. Is this normal?

[guanjiechen118]

Some updates: I've trained stage 1 smoothly. I can see clearly the validation sample grow from noises to recognizable objects, though the quality is low. Then I started stage 2 training from a checkpoint. But after about 2500 optimizer steps the validation sample still has roughly the same quality as the checkpoint it started from. Is this normal?

I think this is normal. Stage 2 typically requires longer training time compared to Stage 1. I expect that another day of training will yield better results, and the model should converge within approximately 3 to 5 days.

[Darkbblue]

Some updates: I've trained stage 1 smoothly. I can see clearly the validation sample grow from noises to recognizable objects, though the quality is low. Then I started stage 2 training from a checkpoint. But after about 2500 optimizer steps the validation sample still has roughly the same quality as the checkpoint it started from. Is this normal?

I think this is normal. Stage 2 typically requires longer training time compared to Stage 1. I expect that another day of training will yield better results, and the model should converge within approximately 3 to 5 days.

This information is really a relief. I'll wait patiently in the coming weekend then.

[Darkbblue]

Some updates: I've trained stage 1 smoothly. I can see clearly the validation sample grow from noises to recognizable objects, though the quality is low. Then I started stage 2 training from a checkpoint. But after about 2500 optimizer steps the validation sample still has roughly the same quality as the checkpoint it started from. Is this normal?

I think this is normal. Stage 2 typically requires longer training time compared to Stage 1. I expect that another day of training will yield better results, and the model should converge within approximately 3 to 5 days.

This is the original output of PixArt-Sigma. This is where I end stage 1 and start stage 2 training (at 61950 optimizer steps). This is the stage 2 output using lr=1e-5 at 41000 optimizer steps. This is the stage 2 output using lr=1e-4 at 26000 optimizer steps. (I've tried more steps resuming from an earlier stage 1 checkpoint at lr=1e-4, and the output is more recognizable, but still bad.)

I'm a bit confused right now, as it seems really hard to fully restore the generative capability of the original model. What might be the cause? The stage 1 checkpoint being bad? The lr being too large or small? Or simply I need more time?

[guanjiechen118]

Recommendations

1. The generated image from Stage 2 with a learning rate of $1 \times 10^{-5}$ looks reasonable. While further training is needed, there has been some improvement compared to Stage 1. The final result should be promising.

2. In contrast, the result with a learning rate of $1 \times 10^{-4}$ shows clear signs of training collapse and is not viable. We can safely exclude this setting from further consideration.

Training should progress gradually. Any generated samples resembling the last figure—exhibiting artifacts or degradation—are likely indicators of instability or training failure.

Based on the visualizations, we recommend continuing training with a learning rate of $1 \times 10^{-5}$ for an additional 10,000 steps.

Model Training Context

It’s worth noting that PixelArt-Sigma is a weak-to-strong, high-resolution model with a two-stage training pipeline. In contrast, the pre-training stage of Skip-DiT is designed for foundation models. Given these architectural and methodological differences, it is uncertain whether our current training strategy can be directly applied to such a complex model without further adaptation.