S tr AE : Autoencoding for Pre-Trained Embeddings using Explicit Structure (2023)

Abstract: This work presents StrAE: a Structured Autoencoder framework that through strict adherence to explicit structure, and use of a novel contrastive objective over tree-structured representations, enables effective learning of multi-level representations. Through comparison over different forms of structure, we verify that our results are directly attributable to the informativeness of the structure provided as input, and show that this is not the case for existing tree models. We then further extend StrAE to allow the model to define its own compositions using a simple localised-merge algorithm. This variant, called Self-StrAE, outperforms baselines that don’t involve explicit hierarchical compositions, and is comparable to models given informative structure (e.g. constituency parses). Our experiments are conducted in a data-constrained (circa 10M tokens) setting to help tease apart the contribution of the inductive bias to effective learning. However, we find that this framework can be robust to scale, and when extended to a much larger dataset (circa 100M tokens), our 430 parameter model performs comparably to a 6-layer RoBERTa many orders of magnitude larger in size. Our findings support the utility of incorporating explicit composition as an inductive bias for effective representation learning.

## introduction human understanding of natural language is generally attributed to the understanding of composition. the theory is that smaller constituents (words, tokens) recursively combine into larger constituents (phrases, sentences) in a hierarchically structured manner, and that it is our knowledge of said structure that drives understanding @xcite @xcite . compositionality can help drive efficient and effective learning of semantics as it presumes knowledge only of the immediate constituents of a phrase and the process of their composition. however, the models that have come to dominate nlp in recent years generally do not explicitly take this property into account. transformers @xcite have the capacity to model hierarchical compositions through the self-attention mechanism, whereby tokens can come to represent varying degrees of the surrounding context. however, should such behaviour occur, it is incidental to the training process, as there is no strict requirement for tokens to represent higher-order objects, and tokens are never explicitly merged and compressed with one another. whether transformers are able to acquire knowledge of syntax, as understood from the linguistics perspective, remains unclear @xcite . however, there is evidence that these models do become increasingly tree-like with sufficient scale and training steps @xcite . this raises an interesting question. to what extent is this drive towards tree-likeness responsible for their representation-learning capability? furthermore, evidence from lake et al. @xcite suggests that incorporating appropriate inductive biases towards composition can help bridge the stark disparity between the quantity of data required for learning between ml models and humans, and allow computational models to generalise better. if it is indeed an inductive bias towards hierarchical composition that enables humans to acquire multilevel semantics efficiently, and nlp architectures generally aren't explicitly tasked with modelling such structure, what happens when they are? to investigate this, we develop strae, a structured autoencoder framework. it takes as input a tree structure that represents a hierarchical composition process, and uses it to directly specify the flow of information from leaves to root (encoder) and back from root to leaves (decoder). to disentangle the effect of structure from other confounding factors, we constrain strae to only have access to the information immediately provided to it by the input tree. this means that for a given node the model can only use information from the nodes immediately connected to it (constituents): a property we refer to as faithfulness. following this, we investigate which training objectives are best suited to enable structured multi-level representation learning, and present a novel application of the contrastive loss to hierarchical tree-like structure. to investigate the utility of compositional structure for representation learning, we employ a dataconstrained setting (≈10 million tokens), and evaluate on a series of tasks that measure semantics at both the sentence and word level. we compare the representations learned by strae to a series of baselines consisting of other tree-models which do not enforce faithfulness as rigidly as strae and a series of baselines that do not use explicit structure at all. we also verify that our performance is attributable to the form of composition expressed by the tree structures by comparing results across a range of different input structures. finally, to investigate how useful a simple bias towards hierarchical composition is, we extend strae to allow it to define its own compositional hierarchy. this variant, called self-strae, utilises the learned representations and the encoder to define its own "merge" sequence, which is then employed as the tree structure for the decoder. our results indicate that knowledge of structure is indeed beneficial, and that surprisingly even a simple bias for hierarchical composition leads to promising results. in light of these findings, we then extended our experiments to the 100 million token range and analyse how well self-strae performs in a significantly larger setting. even more surprisingly, despite solely having 430 non-embedding parameters, self-strae is able to achieve comparable performance to a 6 layer roberta @xcite model with 3.95 million parameters. ## model we develop a framework (strae) that processes a given sentence to generate multi-level embeddings by faithfully conforming to the given structure. intuitively, it involves embedding the tokens of a sentence onto the leaves of the structure, composing these embeddings while traversing up the structure to generate the full-sentence embedding, and then traversing back down the structure while decomposing embeddings from parent to children, to finally recover the sentence itself-effectively a structured autoencoder. while this generates embeddings for multiple levels of the sentence as dictated by the structure, it also generates embeddings for a node in two directions-composing upwards and decomposing downwards. the composition embeddings represent the local context for a node, while the decomposition embeddings represent its full context given the sequence. denoting embeddings e @xmath0, we distinguish embeddings formed traversing upwards and downwards as ēi and ēi respectively. encodings for the tokens are denoted as the vertices w @xmath1, in a v -simplex for vocabulary size v . note that while the token encodings w i are effectively onehot, reconstructions ŵi can be any point on the simplex-interpreted as a distribution over vertices, and thus the vocabulary, with an argmax retrieving the appropriate token. we define the four core components of strae as figure 2 : contrastive objective over structure: corresponding node (green) is pulled closer, and other nodes (red) are pushed away. ## experimental setup we divide our experiments into two separate sections, but both share the same overall setup for pretraining data and evaluation. for pre-training data we take a 500k sentence (≈10m tokens) subset of english wikipedia and 40k sentence development set. we restrict our pre-training data to this scale in order to measure our efficiency hypothesis presented in the introduction. for evaluation, we assess on three categories of tasks: word level semantics, sentence level semantics, and sentence pair classification, with each aiming to capture separate areas of semantic understanding. on the word level we use simlex @xcite , wordsim s and wordsim r @xcite . on the sentence level, we use three tasks from the sts suite @xcite @xcite , the sick relatedness dataset @xcite , and for sentence pair classification tasks we use rte and mrpc taken from the glue benchmark @xcite . all tasks apart from the final two classification tasks are measured using the spearman correlation of the cosine similarity of model embeddings for each pair and human judgements; classification tasks are measured using accuracy. to emphasise the impact of the pre-training with structure, we do not fine-tune any models for the evaluation tasks, instead keeping them frozen. where a classifier is required, we fine-tune only a task-specific classification head consisting of a ffn with a 512 dimensional hidden layer and an intermediate tanh activation function. we choose this setup to match the glue benchmark. for all experiments, we pretrain the models across 5 random seeds and present the averaged performance. downstream classifiers are themselves also trained across 5 random seeds and the average reported. additionally, for all experiments we set the embedding dimension to 100. finally, for all models trained on the word level we filter the vocabulary to exclude words occurring fewer than two times in the data. ## comparing tree models here, we compare strae with two existing treearchitectures. these are the iornn @xcite @xcite and the tree-lstm @xcite . both architectures take structure as input and are able to traverse through the tree to learn representations. however, they differ in the constraints they impose on information flow. we conduct our experiments along two different axes: how well do the representations perform, and to what extent do the models discriminate between input structures? to achieve these evaluations, we parse our pretraining set into three kinds of structure. the first are constituency trees extracted from corenlp @xcite and binarised using nltk @xcite . the two other kinds are purely right-branching and balanced binary trees, which we extract using standard algorithms. the resulting structures are then converted to dgl graphs @xcite . hyper-parameter descriptions for each model can be found in appendix a. ## comparison to unstructured baselines here, we compare strae to a series of baselines that are not provided any form of parse tree as input. the aim is to evaluate the usefulness of explicit structure against other inductive biases. we also introduce a variant of strae called self-strae that is not provided structure as input but must learn its own mode of composition. it serves as a measure of the utility of an explicit compositional bottleneck. ## million tokens this is also indicated by its comparative lower wordsim r performance. on both other lexical semantics tasks strae and self-strae both outperform it. fasttext also struggles on the sentence level. similarly, the bi-lstm performs well on lexical semantics, but struggles at capturing higher levels. this is evidenced both by the sts results and lower lexical performance of the bpe bi-lstm compared with self-strae. roberta performs comparably to strae on sts, but struggles on the word level. this might simply be because transformers aren't designed to learn static lexical embeddings, but could also be the result of our data-constrained pretraining setting, as there was significant variability between seeds. consequently, we conducted a final experiment using a significantly larger pre-training set to assess the impact of scale. 100 million tokens for this experiment, we turned to the wikitext-103 benchmark dataset @xcite . wikitext-103 consists of 103 million tokens, with an average sequence length of 118. each input sequence corresponds to an article rather than a sentence, as in our original dataset. we set the maximum sequence length to 512 and train a subword tokeniser on the training set with a maximum vocabulary of 25000. this vocabulary is used for both roberta and self-strae. as shown in ## discussion and future work we establish two findings. firstly, defining representation similarity through composition is useful, and secondly, asking a model to arrange its own composition is a powerful inductive bias. neither of these findings are limited in their application to the architecture presented in this paper. the requirements are: an explicit merge operation and an objective that optimises representations across all levels. as long as these conditions are met, the findings are in theory applicable to any number of architectures. for future development, the natural next step is to examine what happens when we allow for a significantly more flexible models to dictate their own compositions. transformers can be naturally adapted to incorporate such a bias, and we believe this holds considerable promise for future work, especially in light of recent findings that incorporating compression bottlenecks in transformers is beneficial @xcite . extensions need not be limited to the transformer architecture, either. recurrent and recursive neural networks share many similar features, and given the recent resurgence of rnns @xcite there may also be promise in extending research in this direction.