PT-DiT: A More Efficient Paradigm for Text-to-Any Generation

360 AI Research
2024-10-17 145 views
PT-DiT: A More Efficient Paradigm for Text-to-Any Generation

The performance is not inferior to the SOTA model, but the computational overhead is lower——

Sun Yat-sen University and 360 AI Research jointly launched PT-DiT. Under the same parameter scale, the calculation amount is only 51.4% of DiT and 17.5% of Lumina-Next.

Specifically, PT-DiT is based on the Proxy token mechanism and can be used for various tasks such as text-to-image (T2I) (Qihoo-T2I), text-to-video (T2V) (Qihoo-T2V) and text-to-multi-view (T2MV) (Qihoo-T2MV).

(Qihoo-T2X refers to text to arbitrary vision tasks)

Picture

Without further ado, let’s look directly at some final generated effects, text-to-image (T2I)be like:

Picture

Next is the hot video generation this year, the prompt is as follows:

Sunset cityscape with spires, buildings, clouds, warm glow, and trees.

, duration 00:03

Finally, there is multi-view generation. Any materials such as swords and little yellow ducks can be converted into 3D effects.

Picture

At present, the research has opened its papers, project homepage and code repository, and will soon be open source.

Research motivation

Currently, models based on Diffusion Transformer (Sora, Vidu, Flux, etc.) are able to generate high-fidelity images or videos with strong consistency with text instructions, which greatly promotes the progress of visual generation.

However, the quadratic complexity of global self-attention with respect to sequence length increases the computational overhead of Diffusion Transformer, resulting in longer generation time and higher training cost in practical applications.

This problem also hinders the application of Diffusion Transformer in high-quality and long-term video generation.

For example, 3D full attention, which is better than 2D spatial attention + 1D temporal attention, is difficult to explore for higher resolution and longer video generation due to limitations in computational overhead.

Some research in the field of visual understanding and recognition has found that global self-attention has a certain redundancy due to the sparse and repetitive nature of visual information.

The research team discovered through visual attention maps:

**Different tokens in the same window pay similar attention to tokens that are far away in space, and pay different attention to tokens that are close in space. **

Picture

As shown in the figure above, the research team analyzed the attention map in self-attention of PixArt-𝓪 at a resolution of 512x512.

Then the attention maps of tokens located in the same spatial window are combined together, as shown on the right side of the figure, where the vertical axis represents different tokens in the window, and the horizontal axis represents the degree of association between the token in the window and all tokens.

It is obvious that the attention of different tokens in the same window is almost the same for tokens that are far apart in space, that is, at the same horizontal position, the vertical values ​​​​are almost the same; while spatially adjacent tokens show different attention.

This shows that computing attention for all tokens is redundant, while computing attention for spatially adjacent tokens is crucial.

Therefore, the research team proposed a proxy token-based sparse attention strategy to sample limited proxy tokens from each window to execute self-attention, thereby reducing redundancy and reducing complexity.

Qihoo-T2X method

As shown in the figure below, PT-DiT proposed by the research team introduces proxy tokens to reduce the number of tokens involved in calculating global self-attention and efficiently establish the association of global visual information.

The two core modules included in PT-DiT are:

  • Global Information Interaction Module (GIIM)

  • Texture Complement Module (TCM)

Picture

Among them, GIIM uses a sparse proxy token mechanism to promote efficient interaction between all potential codes, and the correlation between spatially adjacent tokens cannot be ignored, especially for image generation tasks with high requirements for detailed textures.

For this purpose, the research team designed TCM, which further refines local details through window attention and shift window attention.

These two parts are described in detail below:

Global Information Interaction Module

Given a series of latent tokens, a series of proxy tokens are first obtained by calculating the average token according to the spatial and temporal priors (that is, located in the same spatial window).

Each proxy token represents the information of a local area within the image or video, and interacts with proxy tokens in other local areas through self-attention to establish global visual association.

Subsequently, the information contained in proxy tokens is propagated to all latent tokens through cross-attention with latent tokens, thereby achieving efficient global visual information interaction.

Texture Complement Module

Due to the interactive characteristics of sparse proxy tokens and the lack of correlation between spatially adjacent tokens, the generation model has limited ability to model detailed textures and is difficult to meet the high-quality requirements of the generation task.

In order to solve this problem, the research team introduced local window attention to supplement the model's detailed modeling and smoothing capabilities.

However, only window attention will lead to a lack of connection between tokens between windows, resulting in an obvious image grid phenomenon.

Therefore, shift window attention is also introduced in TCM to alleviate the problems caused by local window attention.

Picture

Since the number of tokens involved in calculating window attention is small, the computational complexity of the model does not increase on a large scale.

Compression ratio

For the image generation task, the research team found that maintaining the same number of windows at different resolutions is crucial to ensure a consistent semantic hierarchy, which facilitates the training process from low to high resolution.At the same time, the number of windows should be maintained to prevent the semantic information within the window from being too rich, resulting in a single token being insufficient to represent a local area to complete global information modeling.

Therefore, the research team set the compression ratio (𝑝𝑓, 𝑝ℎ, 𝑝𝑤) to (1,2,2), (1,4,4), (1,8,8) and (1,16,16) at 256, 512, 1024 and 2048 resolutions respectively.

When the input is an image, 𝑓 and 𝑝𝑓 will be set to 1.

For the video generation task, the research team set 𝑝𝑓=4 at different resolutions to maintain consistent time compression.

Due to token compression in the frame, height and width dimensions, PT-DiT can train generators for longer videos.

Computational complexity analysis

PT-DiT uses only a small amount of representative token attention, reducing the computational redundancy of the original full tokenself-attention.

The research team further theoretically analyzed the advantages of PT-DiT in terms of computational complexity.

The computational complexity of self-attention is 2N2D and is calculated as follows:

Picture

where N represents the length of the potential tag and D represents the feature dimension.

Similarly, the computational complexity of GIIM and TCM is calculated as follows:

Picture

Obviously, PT-DiT has significant advantages due to the proxy token optimization strategy, especially when the compression ratio (𝑝𝑓, 𝑝ℎ, 𝑝𝑤) is large and the sequence length (N) is long. **

When (𝑝𝑓, 𝑝ℎ, 𝑝𝑤) are (1,2,2), (1,4,4), (1,8,8) and (1,16,16) and the image resolution is 256 (N=256), 512 (N=1024) , 1024 (N=4096) and 2048 (N=16384), the computational complexity of PT-DiT is only **34.3%, 9.7%, 4.7%, and 2.3% of the global self-attention. **

Furthermore, PT-DiT provides greater benefits for video generation tasks with longer sequence lengths.

experiment

The authors conducted qualitative and quantitative experiments to evaluate Qihoo-T2X on T2I, T2V and T2MV tasks.

Qualitative analysis

For Text-to-Image, as shown, Qihoo-T2I is able to generate realistic images that closely match the provided text hints.

Picture

For Text-to-Video, the researchers compared Qihoo-T2V with recently released open source text-to-video models (i.e. EasyAnimateV4 and CogVideoX) at 512 resolution, as shown in the figure and achieved better results.

Picture

Finally, the authors further explore the effectiveness of PT-DiT in the Text to Multi-View (T2MV) task.

The trained Qihoo-T2MV is able to generate 512x512x24 images from various viewpoints based on the provided text instructions, showing strong spatial consistency.

Picture

Quantitative analysis

The research team quantitatively evaluated Qihoo-T2I and Qihoo-T2V on benchmarks such as MS-COCO FID-30K, UCF-101 and MSR-VTT.

The results show that both Qihoo-T2I and Qihoo-T2V can achieve competitive performance, proving the effectiveness of PT-DiT.

Picture

Computational complexity analysis

As shown in the figure, whether it is an image or video generation task, under the same parameter scale, PT-DiT has a substantial computational complexity advantage over the existing Diffusion Transformer method.

At the same time, comparing EasyanimateV4 with 3D full attention modeling, its training memory explodes as the number of frames increases, while the memory of PT-DiT only increases slightly, indicating that PT-DiT has the potential to complete longer video generation tasks.

Picture Picture

The final ablation study also verifies the rationality of the key design in PT-DiT.

Picture

For more details, please refer to the original paper.

Paper address: https://arxiv.org/pdf/2409.04005 Project homepage: https://360cvgroup.github.io/Qihoo-T2X Code repository: https://github.com/360CVGroup/Qihoo-T2X

End