[feat] add casual vqvae ✨

[qqingzheng]

TODO

• Codebase
• Train a test version
• Do more validation

Casual VQVAE

The structure optimization of vqvae in VideoGPT is as follows:

1. Replace regular Conv with CausalConv. (Repeat the first frame k-1 times, where k represents the temporal dimension size of the convolutional kernel, thereby independently isolating the first frame as an image.)
2. Use interpolation plus convolution for upsampling.
3. Separate temporal compression from spatial compression. Spatial compression is achieved with 2D convolution, while temporal compression is achieved with 3D convolution (padding the space to ensure the spatial size remains unchanged).
4. Incorporate Axial Attention on the temporal dimension with a Causal Mask to ensure that the first frame and other frames are not associated, and that each frame is only related to the current and previous frames.

Examples（4000 steps ckpt）

Image Reconstruction

Small Model	Large Model

Video Reconstruction

Small Model	Large Model

[alvinliu0]

Hi,

Many thanks for your great contribution! May I ask the reasons of step1 structure optimization of video vqvae:

[Replace regular Conv with CausalConv. (Repeat the first frame k-1 times, where k represents the temporal dimension size of the convolutional kernel, thereby independently isolating the first frame as an image.)]

I understand that with such padding, the temporal conv kernel will only see replicates of the first frame in the first temporal convolution operation, but what's the significance for doing so? For example, if the temporal conv kernel size is 4, and frame number is 4, stride is 2, the convoluted sequence would be like:

replicated frame id: 1, 1, 1, 1, 2, 3, 4, pad
convoluted frame id: [1, 1, 1, 1], [1, 1, 2, 3], [2, 3, 4, pad]

which is unclear to me why the casual video vae should work like this. Besides, how do you extend the pretrained spatial vqvae into this casual video one?

Thanks for your explanation in advance!

Best

[qqingzheng]

For the first question, if the first frame is copied k-1 times and concatenated, then no matter how many convolutional layers it goes through, the first frame will exist independently of the other frames and will not blend with them. In this way, the first frame is treated as an independent image.

For example:

Input: (1,2,3,4) -> (1,1,1,1,2,3,4,0)

Conv1: ([1,1,1,1], [1,1,2,3], [2,3,4,0]) is denoted as （a,b,c) -> (a,a,a,a,b,c)

Conv2: ([a,a,a,a], [a,a,b,c])

Despite undergoing two layers of convolution, it can be observed that the computation from 1 to "a" is actually independent of other frames, and this holds true regardless of the number of layers of convolution.

[qqingzheng]

For the second question, you can see #168.

[alvinliu0]

Many thanks for your explanation, which makes much sense to me!

May I ask a following question: with such designs, how do you enable the joint image-video training, or what's the significance of an independent first frame.

For example, if you have both an image and video dataset, and sample i_0 from image dataset, and i_1, ..., i_4 from video dataset. Do you mean the training batch is replicated as: [i_0, i_0, i_0, i_0, i_1, i_2, i_3, i_4]? I guess not, since in this way, the convolution of [i_0, i_0, i_1, i_2] would be meaningless as they are from different datasets.

If the first frame just comes from the first frame of video sequence, then we can not benefit from the large-scale image dataset.

Many thanks!

[qqingzheng]

Many thanks for your explanation, which makes much sense to me!

May I ask a following question: with such designs, how do you enable the joint image-video training, or what's the significance of an independent first frame.

For example, if you have both an image and video dataset, and sample i_0 from image dataset, and i_1, ..., i_4 from video dataset. Do you mean the training batch is replicated as: [i_0, i_0, i_0, i_0, i_1, i_2, i_3, i_4]? I guess not, since in this way, the convolution of [i_0, i_0, i_1, i_2] would be meaningless as they are from different datasets.

If the first frame just comes from the first frame of video sequence, then we can not benefit from the large-scale image dataset.

Many thanks!

Currently, we have not yet attempted to conduct mixed training of images and videos in AE. At present, we still treat the first frame of the video as an image and use video datasets for training. However, we have adopted a mixed form of images and videos in the training of DiT, and you can find the relevant code in opensora/train/train_t2v.py.

[ivylilili]

Hi，I have a question about the structure design in the CausalVAE. TimeUpsample2x since it both can accept the image input and video input, for image, the 'TimeUpsample2x' will not upsample time dimension. But for video input, if the frames number is 4 (since the time compression is 4), then the latent gets 1 channel in the time axis, which will be treated as 'image' when going through 'TimeUpsamle2x' and still output 1 channel in the time dimension?

[qqingzheng]

Hi，I have a question about the structure design in the CausalVAE. TimeUpsample2x since it both can accept the image input and video input, for image, the 'TimeUpsample2x' will not upsample time dimension. But for video input, if the frames number is 4 (since the time compression is 4), then the latent gets 1 channel in the time axis, which will be treated as 'image' when going through 'TimeUpsamle2x' and still output 1 channel in the time dimension?

Yes, currently this is a not very good feature of our CausalVAE. If you have better suggestions, welcome PR.

[Birdylx]

@qqingzheng hi, may I ask a question about time compress and loss funciton for the latest CausalVAE

1. there is no parameters in TimeDownsample2x and TimeUpsample2x, is that enough to learn the time compression?
2. it seems you only compute the perceputal loss and recon loss, no gan loss involved?

[qqingzheng]

@qqingzheng hi, may I ask a question about time compress and loss funciton for the latest CausalVAE

1. there is no parameters in TimeDownsample2x and TimeUpsample2x, is that enough to learn the time compression?
2. it seems you only compute the perceputal loss and recon loss, no gan loss involved?

Thank you for your question. Your question is very crucial.

For the first question, we also believe that it is insufficient. This is a compromise solution we adopted to maximize the utilization of the pretrain weights of the SD VAE. In our subsequent training, we also tried some methods to solve this problem, such as using special initialization methods to initialize the 3D convolution kernel.

For the second question, GAN may be difficult to train, and we hope to iteratively train the model from simple to complex. Therefore, we have temporarily removed GAN to accelerate our experimental speed and cost.

[Sutongtong233]

For the second question, you can see #168.

The encoder and decoder in this repository differ from yours. Perhaps you could provide your implementation of the inflate 2DVAE? Alternatively, I can contribute my implementation.

[Birdylx]

@qqingzheng hi, may I ask a question about time compress and loss funciton for the latest CausalVAE

1. there is no parameters in TimeDownsample2x and TimeUpsample2x, is that enough to learn the time compression?
2. it seems you only compute the perceputal loss and recon loss, no gan loss involved?

Thank you for your question. Your question is very crucial.

For the first question, we also believe that it is insufficient. This is a compromise solution we adopted to maximize the utilization of the pretrain weights of the SD VAE. In our subsequent training, we also tried some methods to solve this problem, such as using special initialization methods to initialize the 3D convolution kernel.

For the second question, GAN may be difficult to train, and we hope to iteratively train the model from simple to complex. Therefore, we have temporarily removed GAN to accelerate our experimental speed and cost.

@qqingzheng I add an identity init in CausalConv3d which acts like Identity function. BTW, I add gan loss in training, it actually works well, if only perceptual loss used, there are many grid effect in reconstructed image, gan loss will remove these grid and generate more realistic image

[qqingzheng]

@qqingzheng hi, may I ask a question about time compress and loss funciton for the latest CausalVAE

1. there is no parameters in TimeDownsample2x and TimeUpsample2x, is that enough to learn the time compression?
2. it seems you only compute the perceputal loss and recon loss, no gan loss involved?

Thank you for your question. Your question is very crucial.
For the first question, we also believe that it is insufficient. This is a compromise solution we adopted to maximize the utilization of the pretrain weights of the SD VAE. In our subsequent training, we also tried some methods to solve this problem, such as using special initialization methods to initialize the 3D convolution kernel.
For the second question, GAN may be difficult to train, and we hope to iteratively train the model from simple to complex. Therefore, we have temporarily removed GAN to accelerate our experimental speed and cost.

@qqingzheng I add an identity init in CausalConv3d which acts like Identity function. BTW, I add gan loss in training, it actually works well, if only perceptual loss used, there are many grid effect in reconstructed image, gan loss will remove these grid and generate more realistic image

Thank you for your work! We have also discovered the importance of GAN Loss in our subsequent experiments and will synchronize our latest code to the repository soon.