\begin{table}[!ht]
    \centering
    \begin{tabular}{l|c|c|r}
    \toprule
    \textbf{Partitions} & \textbf{Min Heads} & \textbf{FID} $\downarrow$ & \textbf{Diff} $\downarrow$\\
    \hline
    2 & 4 & 23.33 & --\phantom{00} \\
    4 & 4 & 6.08 &  \gain{-17.25}\\
    6 & 8 & 5.50 & \gain{-0.58}\\
    8 & 8 & 3.40 & \gain{-2.10} \\
    \bottomrule
    \end{tabular}
    \caption{Comparison for number of head partitions when learning LSUN Church. Min heads represents the minimum number of heads in our transformer.}
    \label{tab:church_heads}
\end{table}
\begin{abstract}
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 \href{https://github.com/SHI-Labs/StyleNAT}{https://github.com/SHI-Labs/StyleNAT}.
\end{abstract}

\subsubsection{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~\cite{Karras2018ProgressiveGO}.
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~\cite{karras2019a,Karras2020AnalyzingAI,Karras2020TrainingGA,Karras2021AliasFreeGA}, 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~\cite{Abdal2021StyleFlowAE,humayun2022magnet,humayun2022polarity,karnewar2020msg}. 
\subsubsection{StyleSwin}
StyleSwin~\cite{zhang2022styleswin} 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~\cite{gal2021swagan} and also used by StyleGAN3~\cite{Karras2021AliasFreeGA}, 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\times4$, wherein they use a window size of 4.
We follow a similar structure except replace the initial layer with MHSA for efficiency. 

The \textbf{key design philosophy} of StyleNAT is then conceptually simple yet novel to our best knowledge: \textit{partitioning our attention heads into different groups which can utilize different dilated forms of localized attention}. 
Multiple-headed 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~\cref{fig:architecture}.
While the figure shows three partitions, any number may be used as long as it does not exceed the number of heads.

To accurately judge the quality and usability of StyleNAT we investigate the performance on FFHQ (\cref{sec:ffhq}) and LSUN Churches (\cref{sec:church}). 
We also perform an ablation study compared to StyleSwin (\cref{sec:ablation}) 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)~\cite{Heusel2017GANsTB} training method, have a discriminator learning rate of $2\times 10^{-4}$, use Balanced Consistency Regularization (bCR)~\cite{Zhao2021ImprovedCR}, r1 regularization~\cite{Mescheder2018WhichTM}, and a learning rate decay (specified per dataset).

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\footnote{This time should not be used to calculate throughput, but can be used as a rough estimate. Proper throughput calculations only measure inference as a means to standardize against disk speeds. See our implementation for more details.}.
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~\cref{fig:ffhq256}.

Our second experiment is with LSUN~\cite{yu2015lsun} 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~\cref{tab:church}, but note that not all models from~\cref{table:ffhq} 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.

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~\cref{tab:church_heads} for ablation on the number of partitions. 
An uneven partition results in the tail heads evenly absorbing the remainder.

This flexibility is important for adapting to different datasets with different complexities.
We demonstrate this with our two ablations.
In our ablation on FFHQ, \cref{tab:ablation}, we do not see continual improvements from added partitions (+ Prog Di (4)).
But in our LSUN Church ablations, \cref{tab:church_heads}, 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.
{"value": "2", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Partitions", "dataset": "LSUN Church"} {"value": "4", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Min Heads", "dataset": "LSUN Church"} {"value": "23.33", "type": "Result", "task": "Image Generation", "metric": "FID", "training data/set": "LSUN Church", "test data/set": "LSUN Church", "experimental settings": {"xx": "yy"}, "model": "StyleNAT", "model settings": {"Partitions": "2", "Min Heads": "4"}} {"value": "4", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Partitions", "dataset": "LSUN Church"} {"value": "4", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Min Heads", "dataset": "LSUN Church"} {"value": "6.08", "type": "Result", "task": "Image Generation", "metric": "FID", "training data/set": "LSUN Church", "test data/set": "LSUN Church", "experimental settings": {"xx": "yy"}, "model": "StyleNAT", "model settings": {"Partitions": "4", "Min Heads": "4"}} {"value": "-17.25", "type": "Other"} {"value": "6", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Partitions", "dataset": "LSUN Church"} {"value": "8", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Min Heads", "dataset": "LSUN Church"} {"value": "5.50", "type": "Result", "task": "Image Generation", "metric": "FID", "training data/set": "LSUN Church", "test data/set": "LSUN Church", "experimental settings": {"xx": "yy"}, "model": "StyleNAT", "model settings": {"Partitions": "6", "Min Heads": "8"}} {"value": "-0.58", "type": "Other"} {"value": "8", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Partitions", "dataset": "LSUN Church"} {"value": "8", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Min Heads", "dataset": "LSUN Church"} {"value": "3.40", "type": "Result", "task": "Image Generation", "metric": "FID", "training data/set": "LSUN Church", "test data/set": "LSUN Church", "experimental settings": {"xx": "yy"}, "model": "StyleNAT", "model settings": {"Partitions": "8", "Min Heads": "8"}} {"value": "-2.10", "type": "Other"}