Transformers with Competitive Ensembles of Independent Mechanisms

1 INTRODUCTION . A major theme throughout the history of deep learning has been the introduction of inductive biases in neural architectures , more recently with a focus on the ability to dynamically keep distinct types of information separated . While an MLP architecture has one large hidden representation at each layer , a convnet keeps different spatial positions ’ representations separated by default . This separation enables more appropriate reuse of parameters , improving generalization ( e.g . compared with a fully connected MLP ) by ensuring that some parts of the hidden representation capturing some aspects of the data can remain unchanged when other aspects are changed . Additionally , it is important to be able to reuse parameters in all situations where the parameters are relevant , and not use parameters in positions where they are irrelevant , and this is where attention mechanisms can be very useful . While dividing information between different positions ( for example time steps or spatial positions ) is already very useful , it has been recognized from the earliest deep learning work on the notion of disentangling ( Bengio , 2009 ; Glorot et al. , 2011 ; Rifai et al. , 2012 ; Mathieu et al. , 2016 ; Achille & Soatto , 2018 ) that other features of the data could advantageously be kept well-separated , even over overlapping sets of positions . This has suggested the idea that a model can be decomposed into multiple components , which are often called modules , each operating on a different set of features . Modularity has been identified as an essential ingredient for generalization in machine learning ( Ronco et al. , 1997 ; Alet et al. , 2018 ; Goyal et al. , 2019 ) . The motivating intuition is that if the relationship between the modules changes between training and evaluation , then a model which keeps these modules sufficiently separate but can adapt how they are combined could be more robust . It can even be robust to changes where the overall data distribution differs between training and evaluation . This has been studied in the causality literature through the notion of “ Independent Mechanisms ” ( Peters et al. , 2018 ; Parascandolo et al. , 2018 ) or causal modules , which can be flexibly re-combined , re-used , and re-purposed . While modularity and independent mechanisms ideas are closely related , the latter has a special focus on the notion that mechanisms should have the ability to remain unchanged when unrelated aspects of the world are changed . In that sense it is a more specific idea which builds on the more general concept of modularity . While the study of independent mechanisms in the context of deep architectures is relatively recent ( Goyal et al. , 2019 ; Mittal et al. , 2020 ) , a few ideas are considered central . One is that mechanisms are separately parameterized ( or dynamically parameterized , with the possibility of separation ) , which means that the function computed by a module remains the same even as other mechanisms need to be changed . Another central idea is specialization between mechanisms , which is the idea that mechanisms should seek to only model some parts of the world . One way to help accomplish this is by forcing the mechanisms to compete to explain different positions ( in time or space ) , such that some mechanisms would not be used by the model on positions where they are less relevant . In this work we explore how the idea of independent mechanisms can be beneficial in the Transformer architecture . Transformers ( Vaswani et al. , 2017 ) are based on information sharing across positions controlled dynamically by a soft-attention mechanism ( Bahdanau et al. , 2014 ) , while still using a fully-connected MLP to process the extracted feature vectors ( concatenated over a set of attention heads ) at each position . An important way in which this improves over convnets is that if this attention becomes sufficiently sparse , then it gains the ability to keep information well-separated between different positions . At the same time , at each position , the Transformer stores a single monolithic hidden representation , over which it applies its entire set of parameters . For example , if we consider a generative model of images of animals in a field , then some of the parameters like those describing how animals have symmetric eyes or a certain number of feet , are only relevant for the positions in the image where the animal is present . A normal Transformer , however , would apply the same parameters to the entire hidden representation at all spatial positions . Additionally , if sources of information need to be accessed over multiple positions , it has no way to keep that information well-separated between parts of the hidden representation , unless a large fraction of the parameters are set to exactly zero . In practice , models tend not to learn these sorts of highly sparse parameter matrices as it is not necessary in order to fit the training set . Thus different underlying factors tend to be freely blended together rather than disentangled : we hypothesize and show empirically that this leads to deteriorated generalization when something about some of these factors changes . Our newly proposed technique , which we call Transformers with Competitive Independent Mechanisms ( TIM ) seeks to address this limitation of the Transformer by dividing the hidden representation and parameters into multiple distinct mechanisms . These mechanisms perform self-attention ( over input elements ) separately , and information is exchanged sparingly between the mechanisms using attention . Thus the model is naturally compelled to keep multiple information signals well-separated , even within a single position . Moreover , only the parameters corresponding to an activated mechanism are called upon , focusing on one aspect of the hidden representation . The process of selectively activating some mechanisms and not others relies on competition between mechanisms , just like in recurrent independent mechanism ( RIMs ) ( Goyal et al. , 2019 ) . We hypothesize and show empirically that this provides an inductive bias encouraging the mechanisms to be more independent and specialized , more robust to changes only affecting other mechanisms . 2 TRANSFORMERS WITH COMPETITIVE INDEPENDENT MECHANISMS . 2.1 PRELIMINARIES . Multihead Self-attention sub-layer The attention mechanism can be formulated as querying a dictionary with key-value pairs ( Bahdanau et al. , 2014 ; Vaswani et al. , 2017 ) , e.g. , Attention ( Q , K , V ) = softmax ( QKT / √ dmodel ) · V , where dmodel is the dimensionality of the hidden representations and Q ( Query ) , K ( Key ) , V ( Value ) are specified as the hidden representations of the previous layer in the so-called self-attention sub-layers in the Transformer architecture . The multi-head variant of attention allows the model to jointly attend to information from different representation subspaces , and is defined as Multihead ( Q , K , V ) = Concat ( head1 , · · · , headH ) WO , with the heads defined as : headk = Attention ( QW Q k , KW K k , V W V k ) where W Q k ∈ Rdmodel×dK , WKk ∈ Rdmodel×dK , WVk ∈ Rdmodel×dV , and WO ∈ RHdV ×dmodel are project parameter matrices , H is the number of heads , and dK and dV are the dimensionalities of Key and Value . Group Linear Layer : It takes multiple hidden representations , and applies a separately parameterized linear transformation to each . This operation can be efficiently implemented using batched-matrix multiplications . We set the numbers of groups ns and define a weight tensor W ∈ Rns×din×dout . If the input h is shaped as h ∈ Rns×din , then we can define the layer as : GroupLinear ( h , W , ns ) = [ hjWj ] ns j=1 2.2 TIM ALGORITHM . We first lay out the parts of a TIM layer and then give more detailed steps in Algorithm 1 . We then give a high-level detail of how to turn a transformer layer into a TIM layer in a typical implementation ( Section 2.3 ) . An illustration of how independent mechanisms differ from heads is given in Figure 1 . 2.2.1 COMPETITION BETWEEN DIFFERENT MECHANISMS . Aside from having separate parameters and only exchanging information via inter-mechanism attention , we wanted to create a stronger inductive bias to encourage the mechanisms to specialize . To do this , we created a competition system in which each mechanism has a layer which outputs a single scalar value ( as a function of the current layer ’ s representation ) , and these are passed through a softmax over the different mechanisms ( this softmax is applied position-wise and separately for each layer ) . The value of this softmax is then used to weight how much each mechanism is allowed to update its representation after the self-attention . This competition score is computed as c = softmax ( GroupLinear ( h , W c , ns ) ) , where we note that each mechanism has its own parameters for the layer ( hence the use of a Group Linear layer instead of a normal linear layer ) . Thus the ns modules have a per-step weighting for how much they are able to read during the later self-attention stage . Thus if one mechanism wants to perform attention on a given position , it suppresses the other mechanisms on that position . We found that this often improved results and that these softmax scores are fairly interpretable as a measure of specialization . Exact equations for this step are given in Step 1 and used in Step 2 in Algorithm 1 in the appendix . 2.2.2 EACH MECHANISM SHARES INFORMATION ACROSS TIME AND PROCESSES INFORMATION . This step allows each mechanism to have its own independent dynamics , which are themselves similar to a normal transformer layer . These independent dynamics allow each mechanism to read information from other time steps using attention and process that information using FFN layers . We modify the self-attention sub-layer and feed-forward sub-layers ( FFN ) to be mechanism-wise as well as position-wise , with separate parameters for each mechanism . Additionally , the layer-normalization is modified to be performed separately for each mechanism . The projections and FFN sub-layers can be modified by replacing the linear layers with group linear layers . When performing the self-attention itself , the mechanisms behave the same as heads , and thus we can use the same type of multi-head attention process , so long as the total number of heads is divisible by the number of mechanisms . One notable property is if TIMs only consisted of this part of the model ( independent dynamics ) by itself , then each TIM would be a completely independent transformer model with its own forward pass and its own parameters . Steps 2 and 4 in the appendix , Algorithm 1 give more detail on this step . 2.2.3 ATTENTION IS USED TO COMMUNICATION INFORMATION BETWEEN DIFFERENT MECHANISMS . Although we allow each TIM to remain independent and process information independently , it is also important to allow the different mechanisms in TIMs to share information between each other ( in case the TIMs are not truly fully independent ) . To do this we use a standard multi-head attention sub-layer to share information between the mechanisms , which is done in a position-wise fashion . We made this attention mechanism relatively small , with just 2 heads with 32 units each . This is because we want the different mechanisms to be as independent as possible , and thus only share small amounts of high level information . This can be thought of as another attention layer , where we treat the different mechanisms as positions , and perform this attention in parallel over the different steps in the sequence . More details on this are given in Step 3 in the appendix ’ s Algorithm 1 .

Summary: This paper proposes an independent mechanism that divides hidden representations and parameters into multiple independent mechanisms. The authors claim that the mechanism benefits the computation of sparse tensors; it does learn better inductive biases than a sizeable monolithic model. This idea is particularly similar to Recurrent Independent Mechanisms (RIM) [1], mentioned in the paper. The main contribution of this work is introducing competition between independent mechanisms. The authors evaluate their models on the image transformer model, speech enhancement, and NLP tasks.