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 fine-tune technique to help BERT models to learn & capture form and content information on textual data (without any form of structural parsing needed). They key addition to the classic BERT model is the introduction of the R and S embeddings. R &S are supposed to learn the information in text that is traditionally represented as the structural positions and the content-bearing  symbols in those positions.