{'fraction_non_alphanumeric': 0.024500183981356557, 'fraction_numerical': 0.009689684778609101, 'mean_word_length': 4.94476850164054, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 1, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
StyleNAT: Giving Each Head a New Perspective

Image generation has been a long sought-after but challenging task, and performing the generation task in an efficient manner is similarly difficult. Often researchers attempt to create a"one size fits all"generator, where there are few differences in the parameter space for drastically different datasets. Herein, we present a new transformer-based framework, dubbed StyleNAT, targeting high-quality image generation with superior efficiency and flexibility. At the core of our model, is a carefully designed framework that partitions attention heads to capture local and global information, which is achieved through using Neighborhood Attention (NA). With different heads able to pay attention to varying receptive fields, the model is able to better combine this information, and adapt, in a highly flexible manner, to the data at hand. StyleNAT attains a new SOTA FID score on FFHQ-256 with 2.046, beating prior arts with convolutional models such as StyleGAN-XL and transformers such as HIT and StyleSwin, and a new transformer SOTA on FFHQ-1024 with an FID score of 4.174. These results show a 6.4% improvement on FFHQ-256 scores when compared to StyleGAN-XL with a 28% reduction in the number of parameters and 56% improvement in sampling throughput. Code and models will be open-sourced at https://github.com/SHI-Labs/StyleNAT .


Introduction
Adversarially trained generative image modeling has been long dominated by Convolutional Neural Networks (CNNs). More recently, however, there have been strides in shifting towards new architectures, such as Transformerbased GANs [48,49] and Diffusion-based GANs [44]. In addition to this, other generative modeling methods, such as pure Diffusion [5,38] and Variational Autoencoders [3,11], are catching up to the standards set by GANs. Diffusion and VAEs have some advantages in that they are able to approximate the densities of the training distribution, and thus can incorporate more features and can have a higher recall. But these methods also have a major disadvantage in that they need to process substantially more information, compared to GANs, and thus can be quite computationally expensive and difficult to train.
While convolutional models have the advantage of utilizing local receptive fields, transformer-based models, with self attention mechanisms have the advantage of utilizing a global receptive field. Because of this, there has been a push towards integrating transformers into image generation, as self attention would be a natural choice to help increase quality of generated images, thanks to its content-dependent weighting and ability to model long-range dependencies. Transformers have a major drawback, in that they are both computationally intensive and can be difficult to train. Self attention bears with it a quadratically growing memory and FLOPs, for transformers means that it may simply be computationally infeasible to apply a standard transformer to the entirety of the latent structure. These are some of the reasons early transformer-based GANs had difficulties competing with convolutional based ones [14,15,18].
Quite recently, there has been a push for restricted attention mechanisms, which typically do not pay attention to the entire feature space. By doing this, these trans-formers have a reduced computational burden and also put pressure on fixed-size receptive fields, which can help with convergence. The most popular of these, Swin Transformer [28,29], uses a window self attention (WSA) mechanism. This sifted window breaks up the global attention into smaller non-overlapping patches, which results in a linear time and space complexity. Swin also follows every WSA with a shifted variant, which shifts the dividing lines to allow for out-of-patch interactions. StyleSwin [48] was one of the earliest transformer-based GANs that caught up to its convolutional counterparts. Its style-based generator [23] was built upon Swin's WSA attention mechanism. Despite its success, it failed to excel beyond existing state-of-the-art GANs at the time.
Localized attention mechanisms come with other costs though: the inability to attend globally and therefore capture long-range dependencies. This led to works such as Dilated Neighborhood Attention (DiNA) [8], which was proposed to tackle this issue in models based on localized self attention. This work further extends Neighborhood Attention (NA) [9], which introduced an efficient sliding-window attention mechanism that directly localizes self attention for any point to its nearest neighbors. Because of the flexibility sliding windows provide, NA/DiNA can easily be dilated to span the entire feature map and capture more global context. Within this work we seek to extend the power of the attention heads of NA/DiNA by partitioning them and allowing them to incorporate different vantage points across the latent structure for high-resolution image generation. While this method, called Hydra-NA, is more engineered than the full transformer, it greatly extends the capabilities of previous local attention mechanisms and comes at almost no additional computational cost on top of NA/DiNA. Through these partitions and the use of various kernels and dilations, these heads are able to incorporate a much more feature rich landscape into their structure, thus increasing the information gain that they can obtain on a given latent structure. Our Hydra-NA module increases the flexibility of neighborhood attention modules, allowing them to be better adapted to the data they are trying to learn on various generation tasks.
Our paper's contributions are: • We introduce Hydra-NA, which extends NA/DiNA and provides a flexible design to combine local and long-range receptive fields via different attention heads.
• We propose StyleNAT, an efficient and flexible image generation framework that is not only memory and compute efficient but also adaptive to different datasets and tasks via neighborhood attention and its dilated variants with our Hydra-NA design.
• We achieve new state of the arts on FFHQ-256 (among all GANs) and FFHQ-1024 (among all transformerbased GANs) with our StyleNAT.

Attention-based Models
The Transformer [42] is arguably one of the most prevalent architectures in language processing. The architecture simply consists of linear projections and dot product attention, along with normalization layers and skip connections. Other than large language models directly using this architecture [4,34], the work inspired research into models built with dot product self attention [33,35]. Later in 2020, Vision Transformer [6] applied a plain transformer encoder to image classification, which was outperformed existing CNNs at large scale classification. This inspired many researchers to study Transformers as direct competitors to CNNs in different settings [10,39,40] and across different vision tasks [12,26].
One branch in these studies is hierarchical vision transformers [29,43], which stack multiple transformer encoders, and downsample outputs in between, creating a gradual feature map downsampling which CNNs typically follow. This design allows such models to be easily integrated with existing vision frameworks as backbones. However, the transformer encoders are typically modified with restricted self attention mechanisms to avoid the quadratic cost of self attention. One of the most prevalent of these is Swin Transformer [29].

Swin Transformer
Liu et al. [29] proposed Window Self Attention (WSA) and Shifted Window Self Attention (SWSA), both of which partition feature maps into windows of fixed size, and apply self attention to each window separately. The difference between the regular and shifted variants is that the latter shifts the partition boundaries by shifting pixels, therefore allowing out-of-window interactions and receptive field growth. Because of the fixed window size, self attention's quadratic complexity drops to a linear complexity. Through these mechanisms, they propose a hierarchical transformer model, which is a stack of 4 transformer encoders with self attention replaced with WSA and SWSA (every other layer uses the shifted variant). Their model was studied on a variety of vision tasks, beyond image classification, and became the state of the art on COCO object detection and in- Hydra-NA Architecture Figure 2. Architecture of StyleNAT. We follow the standard StyleGAN architecture but replace the convolutions with attention mechanisms. The first step, size 4 × 4, we use a Multi-Headed Self Attention, and for further steps we use a Hydra-Neighborhood Attention mechanism. We show a Hydra architecture to the right were we partition the heads into 3 sets, which are able to have different kernel sizes and different dilations. We demonstrate this with one full NA and two DiNA attention, with different dilations. The heads are then concatenated back together and the output is treated as if we had performed a normal attention mechanism. The number of partitions, kernel sizes, and dilations are all hyper-parameters within our network and those shown are arbitrary. More detailed architecture implementations are shown in Appendix A.
stance segmentation, as well as on ADE20K semantic segmentation. It also inspired follow up works in using their architecture and WSA/SWSA attention patterns in image restoration [27], masked image modeling [45], video classification [31], and generation [48].

Neighborhood Attention Transformer
Neighborhood Attention (NA) [9] was proposed as a direct restriction of self attention to local windows. The key difference between NA and WSA is that the restriction is pixelwise, leading to each pixel attending to only its nearestneighboring pixels. The resulting attention spans would be in theory similar to how convolutions apply weights, with the exception of how cornering pixels are handled. Compared to WSA, in addition to introducing locality [29], NA also maintains translational equivariance. NA also approaches self attention itself as its window size grows, and unlike WSA, would not need pixel shifts as it is a dynamic operation. Similar operations, in which self attention is restricted in a token-wise manner, had been investigated prior to this work [35], but were less actively studied due to implementation difficulties [29,35,41]. To that end, Neighbor-hood Attention Extension (NATTEN) [9] was created as an extension to PyTorch, with efficient CUDA kernels, which allow NA to run even faster than WSA, while using less memory. The model built with NA, Neighborhood Attention Transformer (NAT) [9], was shown to have superior classification performance compared to Swin, and competitive object detection, instance segmentation, and semantic segmentation performance.

Dilated Neighborhood Attention Transformer
Dilated Neighborhood Attention Transformer (DiNAT) [8] followed NAT [9] by extending NA to Dilated NA (DiNA), which allows models to use extended receptive fields and capture more global context, all with no additional computational cost. DiNA was added to the existing NAT-TEN package and the underlying CUDA kernels, which allowed the easy utilization of this attention pattern. By simply stacking NA and DiNA layers in the same manner that Swin stacks WSA and SWSA, DiNAT models significantly outperformed Swin Transformer [29], as well as NAT [9] across multiple vision tasks, especially downstream tasks.

Style-based GANs
Style based generators have long been the de facto architecture choice for generative modeling. In this section we discuss different relevant generators and their relationship to our network.

StyleGAN
Karras et al. first introduced the idea of using progressive generating using convolutions where the one would begin with a small latent representation, process each layer, and then scale up and repeat the process until the desired image size is achieved [20]. This architecture has major advantages in that it reduces the size of the latent representation, and thus the complexity of data that the model needs to learn. This comes with the obvious drawbacks that we make a trade-off of diversity for fidelity and computational simplicity. Karras et al. later produced several papers that built on this idea [21][22][23][24], improving the progressive network, but also introducing a sub-network, called the "style network", that allows one to control the latent representations at each of these levels. This innovation dramatically increased the usability of this type of network architecture because it allows for simpler feature interpolation, which is often difficult in implicit density models. Because of the high fidelity, flexibility, and the ability to perform feature interpolation, this style of network has become the de facto choice for many researchers [1,16,17,19].

StyleSwin
StyleSwin [48] introduced Swin based transformers into the StyleGAN framework, beating other transformer based models, but still not achieving any state of the art results. This work did show the power of Swin and how localized attention mechanisms are able to out perform many CNN based architectures and did extend the state of the art for style based networks. StyleSwin achieved this not only by replacing the CNN layers with Swin layers, but also partitioned the transformer heads into two, so that each half could operate with a different kernel. This was key to improving their generative quality, allowing the framework to increase flexibility and pay attention to distinct parts of the scene at the same time. StyleSwin also included a wavelet based discriminator, as was introduced in SWAGAN [7] and also used by StyleGAN3 [22], which all show improvements for addressing high-frequency content that causes artifacts. StyleSwin uses a constant window size of 8 for all latents larger than the initial 4 × 4, wherein they use a window size of 4. We follow a similar structure except replace the initial layer with MHSA for efficiency.

Methodology
Our StyleNAT network architecture closely follows that of StyleSwin [48], which closely follows the architecture of StyleGAN2. Herein we described the differences in architecture and the roles that they play in generating our images. We describe how our network can be adapted to many different tasks and how to choose the suitable configurations.

Motivation
Current CNN based GANs have some limitations that make it difficult for them to consistently produce high quality and convincing images. Most notably, CNNs have difficulties in capturing long range features but do have strong local inductive biases and are equivariant to translations and rotations. Thus we need to incorporate architectures that maintain this equivariance, while having both local inductive biases and global inductive biases.
Transformers on the other hand have strong global inductive biases, are translationally equivariant, and are able to capture long range features across an entire scene. The downside of these is that they are computationally expensive. Building a pure transformer-based GAN would be infeasible, as the memory and computation required to generate attention weights grow quadratically. Thus we need to build a transformer like structure that is able to have the advantages of GANs but also be able to include global information while being efficient to train.

Giving Heads a New Perspective
At the core of StyleNAT is the flexible Hydra-NA architecture. Inspired by the multi-headed design in Vaswani et al. [42], we sought to increase the power of these heads, by giving them new ways to view their landscape. While global attention heads can learn to pay attention to various parts of a scene, the reduced receptive field of localized attention mechanisms means that these perspectives can only be performed on that limed viewpoint.
To extend the power of these attention heads, StyleSwin splits attention heads into halves. One half uses a windowed self-attention (WSA) and the other uses a shifted window self-attention (SWSA). This method has limitations in that it still is unable to perceive the global landscape and is limited in the perspectives these heads can take on, even if given more partitions. Additionally, the partitioned self attention breaks translational equivariance [9,41]. Hydra-NA takes a different approach to achieve a similar goal, by allowing each of the attention heads to have a unique perspective of the latent structure: from local to global, depending on the dilation factor. We hence exploit NA/DiNA as our flexible building block.
The key design philosophy of StyleNAT is then conceptually simple yet novel to our best knowledge: partition- StyleGAN-XL has the closest FID but has low throughput and is much larger. ing our attention heads into different groups which can utilize different dilated forms of localized attention. Multipleheaded attention extends the power of the standard attention mechanism by allowing each head to operate on a different set of learnable features. To extend the capabilities of these heads we create a mechanism in which they can have different vantage points to the latent structures. In our framework we allow heads to utilize different NA kernel sizes and different DiNA dilations. Through this, some heads can focus on features within a small and local receptive field, other heads are able to pay attention to medium-sized receptive fields (with low sparsity), while other heads are able to pay attention to large receptive fields (with proportional sparsity). We depict this structure on the right side of Figure 2. While the figure shows three partitions, any number may be used as long as it does not exceed the number of heads.

StyleNAT
We use the same root style based architecture as Style-GAN to incorporate all the advantages of interpretability but seek to increase the inductive powers of our model through the use of Hydra-NA. For this work we converged on a constant kernel size, of 7 (similar to StyleSwin), and chose to include a progressive dilation scheme.
In this respect our partitions have both dense kernels as well as sparse kernels. With this framework we choose two main designs: a dual partition with a local/dense kernel combined with a global/sparse kernel and a progressive dilation design -increasing sparsity and the receptive field with the number of partitions.

Experiments
To accurately judge the quality and usability of Style-NAT we investigate the performance on FFHQ (Section 4.1) and LSUN Churches (Section 4.2). We also perform an ablation study compared to StyleSwin (Section 4.3) to demonstrate the effects of the improvements we have made. These datasets are common evaluation benchmarks that can help us understand how the generator works on more regular data as well as irregular data. All arguments for experiments are saved alongside the checkpoints for verification. For best comparisons we keep hyper-parameters similar to StyleSwin, unless otherwise specified. This means that we use the Two Time-Scale Update Rule (TTUR) [13] training method, have a discriminator learning rate of 2 × 10 −4 , use Balanced Consistency Regularization (bCR) [50], r1 regularization [32], and a learning rate decay (specified per dataset).
Our experiments were performed on nodes with NVIDIA A100 and A6000 GPUs. Most experiments were performed with a single node. While the flexibility of our architecture allows for adaptation to specific datasets we do not do an exhaustive hyper-parameter search, as the number of permutations is too large for our limited compute. We expect follow-up research to be able to out-perform our re-sults with more clever designs and encourage them to open issues on our GitHub page.

FFHQ
Flickr-Faces-HQ Dataset (FFHQ) [23] is a common face dataset that allows generative models to train on regular typed data. FFHQ has images sized 1024x1024 and represents high resolution imaging which are cropped to be centered on faces, making it uni-modal. Faces are fairly regular in structure with the most irregular structure being features like hair and background. This dataset serves as a substitute for data with highly regular structures and can thus give us an idea of how effective generation is for specific and niche tasks. For these experiments we train on all 70k images and then evaluate from a random sample of 50k as compared to 50k samples from the model. For FFHQ-256 experiments we trained our networks for 1M iterations with a batch size of 8 (per GPU), for a total of 64 million images seen. We also used a LR decay starting at 775k iterations and achieve our best score at 940k iterations with a FID50k of 2.046. For our best result we also include a random horizontal flip, which will double the number of potential images. As we are compute limited, it is possible that these scores could be improved by better hyper-parameter tuning and likely could be trained for longer. For this task we found that it was unnecessary to go beyond 2 partitions, and that we got significant improvements with just two partitions with sparse dilations. We show both of these results in our ablation study, Table 3.
Our results represent the current state of the art for FFHQ-256 image generation.
For FFHQ-1024 we use a batch size of 4 (per GPU), 1M iterations, horizontal flips, and we start our LR-Decay at 500k iterations. Our best score is reached at 900k iterations with FID50k of 4.174. This model still significantly benefits from just two partitions but due to the substantial increase in computation required for these larger images, we have yet to be able to train other architectures (this is the only one we tried). We do believe that a larger partition scheme would likely increase the results here, specifically with more partitions in the later layers. Note that training on a single A100 node takes approximately a month and that writing 50k images to disk (to calculate FID) can take over 5.5hrs alone 1 . Still, our results reflect the state of the art for transformer based GANs. In our FFHQ results we notice that our results frequently capture the long range features that we are seeking. This provides evidence that our hypothesis about including non-local features increases the quality of the results. We show some samples of our FFHQ-256 model in Figure 4.

LSUN Church
Our second experiment is with LSUN [47] Church. This dataset is uni-modal, as all images are of churches, but is less structural than FFHQ where images have much higher variance from one another. There is a higher variability in designs as some churches are quite simple and others are complex and have other scene information like people in the foreground/background. This task demonstrates the ability to perform on diverse and complicated distributions of images. We found that with just the two partition design that our models would converge to a FID of around 20. This shows that the simple two partition does not incorporate enough information to build the more complex structures of this distribution. Introducing the four head partition we found that we could achieve a significant improvement in performance, and this dataset represented the major motivation to expanding the capabilities of our architecture. Looking across generative literature it is clear that there is a pattern that models perform better on more structured data or less structured data. We show our results in Table 2, but note that not all models from Table 1 also tested on Church. Computational complexity and throughput will be identical to those of the FFHQ-256 cases as the image sizes are the same. The number of partitions has almost no affect on throughput, slight due to lack of CUDA optimization, and does not affect MACs.
Our partitioned design with progressive dilation demon-  Table 1. FID50k results based on training with entire 70k FFHQ dataset. A random 50k samples are used for each FID evaluation. We separate convolutional based methods from transformer based using a horizontal line. "Usage Metrics" are evaluated on the 256 × 256 resolution images for best comparison to other models. We also note that none of our StyleNAT results (on any dataset) utilize fidelity increasing techniques such as the truncation trick [25].
strates that we can bridge this gap and have the best of both worlds, being able to achieve high quality results on both highly structured data, like FFHQ, as well as less structured data. We are able to achieve this while maintaining a high throughput and without a large number of parameters. While we do not achieve state of the art results yet with our current limited hyperparameter exploration, we still maintain a competitive performance with an FID of 3.400. Samples from the Church generation can be seen in Figure 5 Arch Model FID ↓

Ablation
Our closest competitor is StyleSwin, where we share a lot of similar architecture choices but with some distinct differences. Because of this, we use StyleSwin as our baseline model and build from there. We will also be focusing on FFHQ256 as this shows how we got our best result.
Starting with StyleSwin, introduce Neighborhood Attention into the model, simply by replacing both WSA and SWSA modules in the generator with Neighborhood At-tention (+NA). Next, we introduce Dilated NA (DiNA), replacing only one split of the attention heads with the dilated variant. In other words, we replace WSA with NA and SWSA with DiNA (+DiNA). Third, we introduce the random horizontal flip to demonstrate the effect of augmentation (+Flips). This is not something StyleSwin tried on their FFHQ training but did with other datasets. Finally, we introduce our Hydra-NA into the model, with four partitions; the first partition with dilation 1 (NA) and the last partition with maximum dilation according to feature map size, and the partitions in between adjusted accordingly (+Prog Di). All results are tested on FFHQ-256 and shown in Table 3 In addition to the above ablation we include one for our search space on LSUN Church in an effort to demonstrate the effectiveness of the number of heads and partitions. Because this data is more complex, a simple two partition design was not enough. In an indirect way, the number of partitions required to achieve good results provides insights into the complexity of a dataset. The number of partitions do not change the computational or memory complexity of the model but there is a slight increase in computation as we handle the complexity. All runs use a progressive style of dilation as described above. Below we report the best scores achieved by the number of partitions. These models were trained until they converged or diverged and in all cases we present the best evaluation. Our best runs with 2 and 4 partitions were achieved with r1 set to 20 and were achieved with fewer than 200k iterations. The difference between 2 and 4 partitions represents the largest difference in FID scores between models, demonstrating that for more complex datasets we need more information. For our 6 partition design we made an additional modification to the generative network. The minimum number of heads in the previous setups was 4, which only applied to sizes >128, and we changed this minimum to 8. This change allowed for our usage of 6 partitions but also represents a slight increase in computation and memory, as complexity is influenced by the number of heads. See Table 4 for ablation on the number of partitions. An uneven partition results in the tail heads evenly absorbing the remainder.

Discussion of Other Configurations
We include a small discussion about our architecture search to help others better learn how to utilize StyleNAT/Hydra-NA and accelerate usage. Initially we had explored using progressively growing kernels in a pure NA based architecture but found that this scaled poorly as well as did not provide significant improvements over other models. Progressive kernels also has the downside of significantly increased computational load.
We also investigated replacement of StyleSwin where individual layers were replaced with pure NA layers, but did not run these to convergence. Each swap resulted in a minor improvement in the initial loss curves but not by significant margins. A version where we replaced all version of swin layers can be seen in our ablation study (Table 3, +NAT) and we note that NAT alone only provides minor improvements from Swin, but that dilation (DiNA) offers significant improvements. We remind the reader here that using two NA partitions with the same kernel and dilation -as was done here -is equivalent to using a single partition.
We also did not find any significant improvements when using one layer with a small kernel and another layer with a larger kernel (on the same level but different steps). While we found NA based runs more consistent than StyleSwin runs, layer replacement runs were often within the margin of variance. This leads us to conclude that the increased performance is coming from NA's higher expressivity.
Once we introduced the dilation into the partitioned head design did we find significant improvements within our generative process. This leads us to conclude that the advantage of StyleNAT is the ability to pay attention to both local and global features. We believe that this architecture extends the power of heads in transformer architectures. In a typical multi-headed transformer model, heads are able to pay attention to different features. In our model each of these heads is not only able to pay attention to different features, but can do so in different ways. In addition, one can think of this hydra structure as also a singular creature that can view the same object from multiple viewpoints and combine the information before making decisions. The advantage our of network is this ability to be able to view different things and from different perspectives. This is because we partition the heads, meaning each kernel/dilation configuration has multiple heads.
This flexibility is important for adapting to different datasets with different complexities. We demonstrate this with our two ablations. In our ablation on FFHQ, Table 3, we do not see continual improvements from added partitions (+ Prog Di (4)). But in our LSUN Church ablations, Table 4, we see that this partitioning is essential. We also show that the gain in fidelity is not due purely to the minimum number of heads as there is a much larger gain when moving to 8 partitions rather than 6 partitions, despite that both models have the same minimums. We will continue expanding upon these studies but with increased flexibility comes a larger solution space.

Conclusion
In this work we present StyleNAT, which shows a highly flexible generative network able to accomplish generation on multiple types of datasets. The Hydra-StyleNAT design allows for different partitions of heads to utilize different kernels and/or dilations. This design allows for a single attention mechanism to combine various viewpoints and combine this knowledge to create a much more powerful architecture. We also demonstrate that this style of architecture is able to achieve state of the art performance on image generation tasks.

Ethical Considerations
There are many ethical dilemmas to consider when discussing generative models, especially models that are able to produce high fidelity images of people. We do not condone the usage of our work for nefarious purposes, including but not limited to: manipulation of others or political gain. This includes usage in bot accounts, using the work to generate false or misleading narratives, or acts of terrorism. We do encourage the usage in artistic manners as well as using these results of these models to supplement training data. We highly encourage the usage of these models for the benefit of all and for using them to advance our scientific understanding of the world. These models have both the potential to do great good as well as great harm. While we cannot control how users of our models and frameworks utilize them, we have a duty to stress the importance of thinking about how products using these models can harm our fellow human beings. We stress that engineers need to think deeply about how their products can be abused and to consider the potential costs and benefits and to encourage these discussions within our communities. We are releasing these model and source code with the hope that they will be used for the benefit of all and not for immoral reasons (including those unmentioned herein).