HUBERT Untangles BERT to Improve Transfer across NLP Tasks

1 INTRODUCTION . Built on the Transformer architecture ( Vaswani et al. , 2017 ) , the BERT model ( Devlin et al. , 2018 ) has demonstrated great power for providing general-purpose vector embeddings of natural language : its representations have served as the basis of many successful deep Natural Language Processing ( NLP ) models on a variety of tasks ( e.g. , Liu et al. , 2019a ; b ; Zhang et al. , 2019 ) . Recent studies ( Coenen et al. , 2019 ; Hewitt & Manning , 2019 ; Lin et al. , 2019 ; Tenney et al. , 2019 ) have shown that BERT representations carry considerable information about grammatical structure , which , by design , is a deep and general encapsulation of linguistic information . Symbolic computation over structured symbolic representations such as parse trees has long been used to formalize linguistic knowledge . To strengthen the generality of BERT ’ s representations , we propose to import into its architecture this type of computation . Symbolic linguistic representations support the important distinction between content and form information . The form consists of a structure devoid of content , such as an unlabeled tree , a collection of nodes defined by their structural positions or roles ( Newell , 1980 ) , such as root , left-child-of-root , right-child-of-left-child-of root , etc . In a particular linguistic expression such as “ Kim referred to herself during the speech ” , these purely-structural roles are filled with particular content-bearing symbols , including terminal words like Kim and non-terminal categories like NounPhrase . These role fillers have their own identities , which are preserved as they move from role to role across expressions : Kim retains its referent and its semantic properties whether it fills the subject or the object role in a sentence . Structural roles too maintain their distinguishing properties as their fillers change : the root role dominates the left-child-of-root role regardless of how these roles are filled . Thus it is natural to ask whether BERT ’ s representations can be usefully factored into content × form , i.e. , filler× role , dimensions . To answer this question , we recast it as : can BERT ’ s representations be usefully unpacked into Tensor-Product Representations ( TPRs ) ? A TPR is a collection of constituents , each of which is the binding of a filler to a structural role . Specifically , we let BERT ’ s final-layer vector-encoding of each token of an input string be factored explicitly into a filler bound to a role : both the filler and the role are embedded in a continuous vector space , and they are bound together according to the principles defining TPRs : with the tensor product . This factorization effectively untangles the fillers from their roles , these two dimensions having been fully entangled in the BERT encoding itself . We then see whether disentangling BERT representations into TPRs facilitates their general use in a range of NLP tasks . Concretely , as illustrated in Figure 1 , we create HUBERT by adding a TPR layer on top of BERT ; this layer takes the final-layer BERT embedding of each input token and transforms it into the tensor 1HUBERT is a recursive acronym for HUBERT Untangles BERT . 2Our code and models will be made available after publication . product of a filler embedding-vector and a role embedding-vector . The model learns to separate fillers from roles in an unsupervised fashion , trained end-to-end to perform an NLP task . If the BERT representations truly are general-purpose for NLP , the TPR re-coding should reflect this generality . In particular , the formal , grammatical knowledge we expect to be carried by the roles should be generally useful across a wide range of downstream tasks . We thus examine transfer learning , asking whether the roles learned in the service of one NLP task can facilitate learning when carried over to another task . In brief , overall we find in our experiments on the NLP benchmarks of GLUE ( Wang et al. , 2018 ) and HANS ( McCoy et al. , 2019 ) that HUBERT ’ s recasting of BERT encodings as TPRs does indeed lead to effective knowledge transfer across NLP tasks , while the bare BERT encodings do not . Specifically , after pre-training on the MNLI dataset in GLUE , we observe positive gains ranging from 0.60 % to 12.28 % when subsequently fine-tuning on QNLI , QQP , RTE , SST , and SNLI tasks . This is due to transferring TPR knowledge—in particular the learned roles—relative to transferring just BERT parameters which have gains ranging from minus 0.33 % to positive 2.53 % . Additionally , on average , we gain 5.7 % improvement on the demanding non-entailment class of the HANS challenge dataset . Thus TPR ’ s disentangling of fillers from roles , motivated by the nature of symbolic representations , does yield more general deep linguistic representations as measured by cross-task transfer . The paper is structured as follows . First we discuss the prior work on TPRs in deep learning and its previous applications in Section 2 . We then introduce the model design in Section 3 and present our experimental results in Section 4 . We conclude in Section 5 . 2 RELATED WORK . Building on the successes of symbolic AI and linguistics since the mid-1950s , there has been a long line of work exploiting symbolic and discrete structures in neural networks since the 1990s . Along with Holographic Reduced Representations ( Plate , 1995 ) and Vector-Symbolic Architectures ( Levy & Gayler , 2008 ) , Tensor Product Representations ( TPRs ) provide the capacity to represent the discrete linguistic structure in a continuous , distributed manner , where grammatical form and semantic content can be disentangled ( Smolensky , 1990 ; Smolensky & Legendre , 2006 ) . In Lee et al . ( 2016 ) , TPR-like representations were used to solve the bAbI tasks ( Weston et al. , 2016 ) , achieving close to 100 % accuracy in all but one of these tasks . Schlag & Schmidhuber ( 2018 ) also achieved success on the bAbI tasks , using third-order TPRs to encode and process knowledgegraph triples . In Palangi et al . ( 2018 ) , a new structured recurrent unit ( TPRN ) was proposed to learn grammatically-interpretable representations using weak supervision from ( context , question , answer ) triplets in the SQuAD dataset ( Rajpurkar et al. , 2016 ) . In Huang et al . ( 2018 ) , unbinding operations of TPRs were used to perform image captioning . None of this previous work , however , examined the generality of learned linguistic knowledge through transfer learning . Transfer learning for transformer-based models has been studied recently : Keskar et al . ( 2019 ) and Wang et al . ( 2019 ) report improvements in accuracy over BERT after training on an intermediate task from GLUE ; an approach which has come to be known as Supplementary Training on Intermediate Labeled data Tasks ( STILTs ) . However , as shown in more recent work ( Phang et al. , 2018 ) , the results do not follow a consistent pattern when using different corpora for fine-tuning BERT , and degraded downstream transfer is often observed . Even for data-rich tasks like QNLI , regardless of the intermediate task and multi-tasking strategy , the baseline results do not improve . This calls for new model architectures with better knowledge transfer capability . 3 MODEL DESCRIPTION . Applying the TPR scheme to encode the individual words ( or sub-word tokens ) fed to BERT , a word is represented as the tensor product of a vector embedding its content—its filler ( or symbol ) 3 aspect—and a vector embedding the structural role it plays in the sentence . Given the results of 3We use filler or symbol to indicate the content symbolic representation throughout the paper . Palangi et al . ( 2018 ) , we expect the symbol to capture the semantic contribution of the word while the structural role captures its grammatical role : x ( t ) = s ( t ) ⊗ r ( t ) ( 1 ) Assuming we have nS symbols with dimension dS and nR roles with dimension dR , x ( t ) ∈ RdS×dR is the tensor representation for token t , s ( t ) ∈ RdS is the ( presumably semantic ) symbol representation and r ( t ) ∈ RdR is the ( presumably grammatical ) role representation for token t. s ( t ) may be either the embedding of one symbol or a linear combination of different symbols using a softmax symbol selector , and similarly for r ( t ) . In other words , Eq . 1 can also be represented as x ( t ) = SB ( t ) R > where S ∈ RdS×nS and R ∈ RdR×nR are matrices the columns of which contain the global symbol and role embeddings , common for all tokens , and either learned from scratch or initialized by transferring from other tasks , as explained in Section 4 . B ( t ) ∈ RnS×nR is the binding matrix which selects specific roles and symbols ( embeddings ) from R and S and binds them together . We assume that for a single-word representation , the binding matrix B ( t ) is rank 1 , so we can decompose it into two separate vectors , one soft-selecting a symbol and the other a role , and rewrite equation ( 1 ) as x ( t ) = S ( a ( t ) S a ( t ) > R ) R > where a ( t ) R ∈ RnR and a ( t ) S ∈ RnS can respectively be interpreted as attention weights over different roles ( columns of R ) and symbols ( columns of S ) . For each input token x ( t ) , we get its contextual representations of grammatical role ( a ( t ) R ) and semantic symbol ( a ( t ) S ) by fusing the contextual information from the role and symbol representations of its surrounding tokens . We explore two options for mapping the input token from the current time-step , and the tensor representation from the previous time-step , to a ( t ) R and a ( t ) S : a Long Short-Term Memory ( LSTM ) architecture ( Hochreiter & Schmidhuber , 1997 ) and a one-layer Transformer . All the models share the general architecture depicted in Figure 1 except for BERT and BERT-LSTM , where the TPR layer is absent . Our conclusion based on initial experiments was that the Transformer layer results in better integration and homogeneous combination with the other Transformer layers in BERT , as will be described shortly . The TPR layer with LSTM architecture works as follows ( see also Figure 2 , discussed in Sec . 4.2 ) . We calculate the hidden states ( h ( t ) S , h ( t ) R ) and cell states ( c ( t ) S , c ( t ) R ) and for each time-step according to the following equations : h ( t ) S , c ( t ) S = LSTMS ( vt , ( vec ( x ( t ) ) , c ( t−1 ) S ) ) ; h ( t ) R , c ( t ) R = LSTMR ( vt , ( vec ( x ( t ) ) , c ( t−1 ) R ) ) ( 2 ) where vt is the final-layer BERT embedding of the t-th word , and vec ( . ) flattens the input tensor into a vector . Each LSTM ’ s input cell state is the previous LSTM ’ s cell output state . Each LSTM ’ s input hidden state , however , is calculated by binding the previous cell ’ s role and symbol vectors . In the TPR layer with Transformer architecture , we calculate the output representations ( h ( t ) S , h ( t ) R ) using a Transformer Encoder layer : h ( t ) S = TransformerS ( vt ) ; h ( t ) R = TransformerR ( vt ) ( 3 ) Each Transformer layer consists of a multi-head attention layer , followed by a residual block ( with dropout ) , a layer normalization block , a feed-forward layer , another residual block ( with dropout ) , and a final layer normalization block . ( See Figure 3 , discussed in Sec . 4.2 , and the original Transformer paper , Vaswani et al . ( 2017 ) , for more details . ) Given that each word is usually assigned to a few grammatical roles and semantic concepts ( ideally one ) , an inductive bias is enforced using a softmax temperature ( T ) to make a ( t ) R and a ( t ) S sparse . Note that in the limit of very low temperatures , we will end up with one-hot vectors which pick only one filler and one role.4 a ( t ) S = softmax ( WSh ( t ) S /T ) ; a ( t ) R = softmax ( WRh ( t ) R /T ) ( 4 ) Here WS and WR are linear-layer weights . For the final output of the transformer model , we explored different aggregation strategies to construct the final sentence embedding : P ( c|f ) = softmax ( WfAgg ( x ( 0 ) , x ( 1 ) , ... , x ( N ) ) ) ( 5 ) where P ( c|f ) is a probability distribution over class labels , f is the final sentence representation , Wf is the classifier weight matrix , and N is the maximum sequence length . Agg ( . ) defines the merging strategy . We experimented with different aggregation strategies : max-pooling , meanpooling , masking all but the input-initial [ CLS ] token , and concatenating all tokens and projecting down using a linear layer . In Devlin et al . ( 2018 ) , the final representation for the [ CLS ] token is used as the sentence representation . However , during our experiments , we observed better results when concatenating the final embeddings for all tokens and then projecting down to a smaller dimension , as this exposes the classifier to the full span of token information . The formal symmetry between symbols and roles evident in Eq . 1 is broken in two ways . First , we choose hyper-parameters so that the number of symbols is greater than the number of roles . Thus each role is on average used more often than each symbol , encouraging more general information ( such as structural position ) to be encoded in the roles , and more specific information ( such as word semantics ) to be encoded in the symbols . ( This effect was evident in the analogous TPR learning model of Palangi et al . ( 2018 ) . ) Second , to enable the symbol that fills any given role to be exactly recoverable from a TPR in which it appears along with other symbols , the role vectors should be linearly independent : this expresses the intuition that distinct structural roles carry independent information . Fillers , however , are not expected to be independent in this sense , since many fillers may have similar meanings and be quasi-interchangeable . So for the role matrix R , but not the filler matrix S , we add a regularization term to the training loss which encourages the R matrix to be orthogonal : L = − ∑ c 1 [ c = c∗ ] logP ( c|f ) + λ ( ||RR > − IdR ||2F + ||R > R− InR ||2F ) ( 6 ) HereL indicates the loss function , Ik is the identity matrix with k rows and k columns , and 1 [ . ] is the indicator function : it is 1 when the predicted class c matches the correct class c∗ label , and 0 otherwise . Following the practice in Bansal et al . ( 2018 ) we use double soft orthogonality regularization to handle both over-complete and under-complete matrices R. 4Note that bias parameters are omitted for simplicity of presentation .

Summary: This paper proposes a layer on top of BERT which is motivated by a desire to disentangle content (meaning of the tokens) and form (structural roles of the tokens).  Figure 1 shows this clearly. The paper considers two variants of the disentangling layer (TPR), one with LSTMs (figure 2) and the other with attention (figure 3). The aim in both is to obtain a decomposition of the form x(t) = S a_s(v_t) a_r(v_t) R where S and R are shared matrices of parameters and v is the output of BERT.