StrAE: Autoencoding for Pre-Trained Embeddings using Explicit Structure

TL;DR

StrAE demonstrates that enforcing explicit hierarchical structure as an inductive bias improves multi-level representation learning, outperforming unstructured baselines and existing tree models in a data-constrained setting. It introduces a contrastive objective over tree-structured embeddings and a Self-StrAE variant that learns its own merge strategy, achieving strong results even with far fewer parameters than large transformers. The work shows the practical viability of explicit composition for efficient, scalable representation learning and suggests promising avenues for integrating hierarchical structure into broader architectures like Transformers.

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.

Figures (7)

• Figure 1: StrAE operation to encode and decode a sentence. Shared colours indicate shared parameters.
• Figure 2: Contrastive objective over structure: corresponding node (green) is pulled closer, and other nodes (red) are pushed away.
• Figure 3: Average performance for models on different task areas by structure type, higher is better.
• Figure 4: Self-StrAE composition. Red dotted arrows indicate cosine similarity between nodes.
• Figure 5: Comparison of trees: Self-StrAE vs. CoreNLP