\begin{table}[ht]
    \centering
    \begin{tabular}{c|c|l}
    \toprule
    Level & Kernel Size & Dilations \\
    \hline
    $4$    & -  & -  \\
    $8$    & 7 & 1 \\
    $16$   & 7 & 1,2\\ 
    $32$   & 7 & 1,2,4\\ 
    $64$   & 7 & 1,2,4,8\\ 
    $128$  & 7 & 1,2,4,8,16\\ 
    $256$  & 7 & 1,2,4,8,16,32\\ 
    $512$  & 7 & 1,2,4,8,16,32,64\\ 
    $1024$ & 7 & 1,2,4,8,16,32,64,128\\
    \bottomrule
    \end{tabular}
    \caption{Example of progressive dilation with 8 heads, referred to ``pyramid dilation.''}
    \label{tab:pyramid_model_arch}
\end{table}
\author{
    Steven Walton\textsuperscript{1},
    Ali Hassani\textsuperscript{1},
    Xingqian Xu\textsuperscript{1},
    Zhangyang Wang\textsuperscript{2,3},
    Humphrey Shi\textsuperscript{1,3} \\
{\small \textsuperscript{1}SHI Lab @ U of Oregon \& UIUC, \textsuperscript{2}UT Austin, \textsuperscript{3}Picsart AI Research (PAIR)}\\
{\small \textbf{\url{https://github.com/SHI-Labs/StyleNAT}}}
}
\twocolumn[{
\renewcommand\twocolumn[1][]{#1}
\maketitle

\begin{itemize}
    \item 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.
    \item 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.
     \item We achieve new state of the arts on FFHQ-256 (among all GANs) and FFHQ-1024 (among all transformer-based GANs) with our StyleNAT.
\end{itemize}

\section{Related Works}
\label{sec:related}
In this section, we will briefly introduce attention modules introduced in Swin Transformer~\cite{liu2021swin}, NAT~\cite{hassani2022neighborhood}, and DiNAT~\cite{hassani2022dilated}.
We will then move on to style-based generative adversarial models, particularly StyleGAN~\cite{karras2019a} and StyleSwin~\cite{zhang2022styleswin}.

\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. 

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).

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.

\definecolor{codegreen}{rgb}{0,0.6,0}
\definecolor{codepurple}{rgb}{0.58,0,0.82}
\definecolor{codeblue}{rgb}{0, 0, 0.6}
\definecolor{codered}{rgb}{0.6, 0, 0}
\definecolor{codeyellow}{rgb}{0.6, 0.6, 0}
\definecolor{background}{rgb}{0.9, 0.9, 0.9}

\lstdefinelanguage{pytorch}
{
    morekeywords={class, def, super, len, for, range, len, if, else, zip, apply, return, __init__, append},
    keywordstyle=\color{codepurple}\textbf,
    keywords=[2]{self,},
    keywordstyle=[2]\color{codegreen}\textbf,
    keywords=[3]{in, True, False, None, or, \%, },
    keywordstyle=[3]\color{codeblue}\textbf,
    keywords=[4]{nn, torch, Module, Linear, cat, Parameter, ParameterList, chunk, shape, zeros, Dropout, squeeze, reshape, permute, trunc_normal_, softmax, forward, dim},
    keywordstyle=[4]\color{codered}\textbf,
    keywords=[5]{natten2dqkrpb,natten2dav,HydraNeighborhoodAttention},
    keywordstyle=[5]\color{codeyellow}\textbf,
}
{"value": "4", "type": "Other"} {"value": "8", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"} {"value": "1", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Dilations", "dataset": "FFHQ"} {"value": "16", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"} {"value": "32", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"} {"value": "64", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"} {"value": "128", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"} {"value": "256", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"} {"value": "512", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"} {"value": "1024", "type": "Other"} {"value": "7", "type": "Hyper-parameter/Architecture", "model": "StyleNAT", "parameter/architecture name": "Kernel Size", "dataset": "FFHQ"}