Prototype memory and attention mechanisms for few shot image generation

1 INTRODUCTION . Recent neurophysiological findings based on calcium imaging have revealed that many neurons in the superficial layers of V1 are strongly tuned to complex local patterns , rather than simple oriented edges or bars . These complex neurons exhibit far stronger responses ( a 3-5 fold increase ) when exposed to their preferred patterns than when exposed to simple grating and bar-like stimuli . The high degree of specificity in these neurons ’ selectivity suggests that they might serve as specific pattern detectors ( Tang et al. , 2018a ) . Due to the selectivity towards complex stimuli , the population response of these V1 neurons is extremely sparse . Only 4–6 out of roughly 1000 neurons respond strongly to any given pattern or natural stimulus ( Tang et al. , 2018b ) . This finding is reminiscent of an earlier study that found similarly sparse encoding of concepts in the hippocampus ( Quiroga et al. , 2005 ) . Recent data in V4 also indicate the existence of cells with a similarly high degree of sparsity1 . Thus , we theorize such cells should exist in every level of the hierarchical visual system . We called these highly selective sparsely-responding feature detectors `` grandmother neurons '' to highlight their possible explicit encoding of specific prototypes , even though in reality , a prototype is likely represented by a sparse cluster of neurons rather than a single cell . Neurons in different layers of each visual area exhibit different degrees of response sparsity , complementing one another in various functions . The observations of this diverse set of super-sparsely responding feature detectors raise the following question : what are the possible computational benefits and rationales of having such neurons in the early visual cortex ? In this paper , we hypothesize that these `` grandmother neurons '' can serve as a prototype memory prior to modulate the process of image synthesis . Image synthesis is a central theme in a number of hierarchical models of the visual system , including interactive activation and predictive coding ( McClelland & Rumelhart , 2020 ; Grossberg , 1987 ; Mumford , 1992 ; Rao & Ballard , 1999 ; Lee & Mumford , 2003 ) , and is hypothesized to take place through the top-down feedback connections between areas in the visual cortex . These priors allow the synthesis process to look beyond the current spatial context and utilize the prototype memories learned and accumulated over time . Thus , the `` grandmother cells '' serve as structural conceptual priors in a memory attention process during image generation . 1Unfortunately the data is not publicly available yet . To emphasize the importance of having an attention mechanism that operated beyond the current image representation , we name our proposed memory-based attentional process Memory Concept Attention ( MoCA ) . MoCA is a module that can be plugged into any layer of pre-existing generator architectures in the GAN framework . We test our model by performing extensive experiments using the state-of-the-art StyleGAN2 ( Karras et al. , 2020b ) and a newly proposed few-shot image generator FastGAN ( Liu et al. , 2021 ) . Our experiments demonstrate the utilization of prototype information accumulated within the semantic clusters during training can improve few-shot image generation on Animal-Face Dog ( Si & Zhu , 2012 ) , 100-Shot-Obama ( Zhao et al. , 2020 ) , ImageNet100 ( Russakovsky et al. , 2015 ) , COCO-300 ( Lin et al. , 2014 ) , CIFAR10 ( Krizhevsky et al. , 2009 ) and Caltech-UCSD Birds ( CUB ) ( Welinder et al. , 2010 ) . Additionally , we also find that generators with MoCA can resist a certain degree of injected noise corruption during testing time , suggesting that attending to a structured memory prior during generation can improve the robustness of the model . Our goal here is to explore the utility of prototype cells as memory priors in a standard computer vision image generation task , with a view to gaining some insights into the advantages of having these `` grandmother '' neurons in the visual cortex at a functional level . 2 RELATED WORK . Visual Concept Learning Visual concepts , defined as intermediate level semantic features , have been shown to be effective in overcoming misclassification of objects due to occlusion when used in a feedforward voting scheme ( Wang et al. , 2017 ) . Explicit representations of visual concepts using prototype neurons might also serve as effective reconfigurable parts for building compositional machines ( Bienenstock et al. , 1997 ; Geman et al. , 2002 ; Zhu & Mumford , 2007 ) . In this paper , we explore the use of visual concepts in the form of prototype memory priors to provide temporal spatial and temporal contextual modulation with an attention mechanism for the complex compositional task of image generation . Self-Attention The attention mechanism ( Vaswani et al. , 2017 ) in deep learning is popular in Natural Language Processing ( NLP ) and is also known as `` non-local networks '' in the vision community ( Wang et al. , 2018 ) . Zhang et al . ( 2019a ) introduced self-attention into generative models in the Self-Attention GAN . Since then , adding self-attention into Generative Adversarial Networks has become a standard practice . Brock et al . ( 2018 ) and Esser et al . ( 2020 ) demonstrate the benefits of using self-attention in high-fidelity image synthesis models . However , current self-attention mechanisms in GANs only utilize contextual information within the same image to modulate the activation . In this work , we propose the use of a memory cache of intermediate-level visual conceptual prototypes to provide additional modulation in addition to the standard self-attention . In our work , the activation mechanism attends not only to its spatial context within the image itself , but also attends to a cached set of memory prototypes accumulated over time . Prototype Memory Mechanism A memory bank can capture diverse features over a long time period , and has been shown to be effective in other domains . Wu et al . ( 2018 ) utilizes a memory bank in contrastive learning to obtain more diverse negative samples to contrast with . Caron et al . ( 2020 ) also uses a memory queue during training to accumulate representative negative samples . He et al . ( 2020 ) shows that the stability of the features accumulated in the memory bank can be improved using a momentum-updated encoder , a strategy we will also use in learning our prototypes . SimGAN ( Shrivastava et al. , 2017 ) introduces the image pool trick which uses a buffer to store previously generated samples , in order to make the discriminator focus not only on the current training batch , but also improve itself based on memory . The major innovation of our work is the proposal , inspired by recent neurophysiological findings , that we should have memory banks at intermediate levels , and , in principle , every level of the visual hierarchy . While early memory bank work stored images at the instance level for object recognition , which is not useful for image generation , we argue that memory banks at lower levels of the hierarchy can store prototypes of parts and sub-parts which are useful for image generation . Few-shot image generation , in particular , can benefit from memory mechanisms that support flexible composition and decomposition of parts , particularly when data are limited . Few-Shot Prototypes Learning Few-shot learning refers to performing computer vision tasks when training data are very limited . One of the popular ideas in the few-shot learning literature is to form distinct prototypes from the training set and use them during testing ( Snell et al. , 2017 ) . Although MoCA also forms prototypes during the training stage and uses them during inference , there are two important distinctions between our work and Snell et al . ( 2017 ) : ( 1 ) Snell et al . ( 2017 ) forms prototypes at the instance level whereas MoCA ’ s prototypes are generated at the intermediate parts level . ( 2 ) Snell et al . ( 2017 ) simply selects the closest prototype to obtain discrete class prediction , MoCA , on the other hand , employs an attention process to continuously modulate the activation features based on the prototypes , hence can be applied to a broader set of tasks , including image synthesis , which predicts continuous pixel values . Few-Shot Image Generation Few-shot image generation task is a challenging task as GANs are very data-hungry and inefficient . Unconditional image generation is especially difficult among the current few-shot image generation schemes . Different solutions have been developed in the literature . Recent works such as DiffAug ( Zhao et al. , 2020 ) , StyleGAN-Ada ( Karras et al. , 2020a ) propose differentiable augmentation to avoid overfitting in the discriminator . InsGen ( Yang et al. , 2021 ) proposes to use a contrastive learning objective to enhance the adversarial loss in the few-shot generation setting . These works mostly proposed methods to improve the discriminators so that they can provide better error signals to the generator , but do not specifically address the generator architecture . Another line of research proposed generator architectures to avoid mode collapse . The state-of-the-art architecture for few-shot image generation ( hundreds of real images ) is FastGAN ( Liu et al. , 2021 ) . Here , we propose a new architectural change on the generator side and make comparison against FastGAN architecture extensively in the experiment section . 3 METHODS . Our key contribution is the introduction of a novel prototype-based memory modulation module to improve the generator network of a GAN . The activation from the previous layer is modified by two attention processes : 1 ) contextual modulation with Memory Concept Attention ( MoCA ) , and 2 ) spatial contextual modulation within the generated image itself ( Self-Attention ) . Our module takes a feature map in the GAN hierarchy as input combines the results of these two mechanisms to modulate the feature map for further downstream processing . A high-level overview of our model is shown in Figure 1 . Formally , we denote the input to our MoCA layer to be the activation of a particular layer , A ∈ Rn×c×h×w . The output produces a modulation H ∈ Rn×c×h×w that update A into Â . To allow the information to be flexibly modulated , we first transform A by 1× 1 convolutions into a lowerdimensional space via functions θ ( · ) , φ ( · ) , and ψ ( · ) , where { θ ( A ) , φ ( A ) , ψ ( A ) } ∈ Rn×c̃ ×h×w . 3.1 PROTOTYPE MEMORY LEARNING . In this section , we discuss the organization of our prototype memory that is used as one route of modulation . Our prototype concept memory is arranged hierarchically into semantic cells and prototype cells . As shown in Figure 2 , each semantic cell is cluster mean representative of a cluster of prototype cells . Formally , suppose our memory P consists of M semantic cells P = { K1 , K2 , ... , KM } . For each of the semantic cells Ki , there are T of prototype cells { E ( i ) 1 , E ( i ) 2 , E ( i ) 3 , ... , E ( i ) T } stored in the memory cell associated with Ki , where E ( i ) j and Ki both ∈ Rc̃ and Ki = ( ∑T j=1E ( i ) j ) /T is the mean of the stored prototype cells . These prototype cells come from the feature maps in previous iterations after being transformed via a momentum updated context encoder φ̃ ( · ) and are updated in the memory at the end of every training iteration . φ̃ ( · ) is a momentum counterpart of φ ( · ) and its parameter is updated as shown in Equation 1 with the momentum parameter m. φ̃ ( · ) does not change as rapidly as φ ( · ) so that the prototype learned are more stable , accumulating information beyond the current training batch . The rationale is similar to that in ( He et al. , 2020 ) ) . φ̃θ ← φ̃θ ∗ ( 1−m ) + φθ ∗m ( 1 ) After transformed by φ̃ ( · ) to a low-dimensional space , the activation at each hypercolumn ( pixel location ) in the feature map is assigned to its closest semantic cluster and replaces an existing prototype cell in the memory bank of that cluster . We use a random replacement policy to prevent the prototype cells from all collapsing to trivial solutions . The update is done in a batch synchronized fashion , i.e . we update Ki as the cluster mean of the ith cluster of prototypes after the entire memory is updated based on the recent batch .

Summary: In this paper, the authors propose a new “grandmother cell”-like memory mechanism for improving image generation performance in GANs. In short, this method clusters and stores activation vectors observed during training. Then at image generation time, activation vectors are augmented with the sum of the stored memories at the closest cluster. This seems to improve GAN performance for few-shot image generation tasks. The authors also visualize the learned memory clusters, which seems to reveal some semantic clustering.