On Geodesic Distances and Contextual Embedding Compression for Text Classification (2021)

Abstract: In some memory-constrained settings like IoT devices and over-the-network data pipelines, it can be advantageous to have smaller contextual embeddings. We investigate the efficacy of projecting contextual embedding data (BERT) onto a manifold, and using nonlinear dimensionality reduction techniques to compress these embeddings. In particular, we propose a novel post-processing approach, applying a combination of Isomap and PCA. We find that the geodesic distance estimations, estimates of the shortest path on a Riemannian manifold, from Isomap’s k-Nearest Neighbors graph bolstered the performance of the compressed embeddings to be comparable to the original BERT embeddings. On one dataset, we find that despite a 12-fold dimensionality reduction, the compressed embeddings performed within 0.1% of the original BERT embeddings on a downstream classification task. In addition, we find that this approach works particularly well on tasks reliant on syntactic data, when compared with linear dimensionality reduction. These results show promise for a novel geometric approach to achieve lower dimensional text embeddings from existing transformers and pave the way for data-specific and application-specific embedding compressions.

## introduction contextual embeddings, like those bert @xcite generates, improve on non-contextual word embeddings by providing contextual semantics to the real-valued representation of a text. although these models have been shown to achieve state-of-the-art performance on most nlp tasks, they are notably expensive to train. to help combat this, as mentioned by @xcite , model compression techniques like data quantization @xcite , model pruning @xcite , and * equal contribution knowledge distillation @xcite have been developed. however, at 768 dimensions, the embeddings themselves can be prohibitively large for some tasks and settings. smaller embeddings both enable more compact data sizes in storage-constrained settings and overthe-air data pipelines, and help lower the requisite memory for using the embeddings for downstream tasks. for non-contextual word embeddings, @xcite note that loading matrices can take multiple gigabytes of memory, a prohibitively large amount for some phones and iot devices. while contextual embeddings are smaller, downstream models will face similar headwind for large corpora. although there has been more extensive study in the efficacy of compressing non-contextual word embeddings @xcite , to the best of our knowledge few contextual embedding compression postprocessing approaches have been proposed @xcite . in their work, @xcite propose the variational information bottleneck, an autoencoder to create smaller, task specific embeddings for different languages. while effective, the computational expense of additional training loops is not appropriate for some memory constrained applications. our approach more closely mirrors the work of @xcite who propose a principal component analysis (pca)-based post-processing algorithm to lower the dimensionality of non-contextual word embeddings. they find that they can replicate, or, in some cases, increase the performance of the original embeddings. one limitation to this approach is the lack of support for nonlinear data patterns. nonlinear dimensionality reductions, like the isomap shown in figure 1 , can pick up on latent textual features that evade linear algorithms like pca. to achieve this nonlinearity, we extend this approach to contextual embeddings, adding in figure 1 : visualization of two-dimensional pca and isomap compressions based on bert embeddings for the sms-spam dataset @xcite . spam is represented by a blue dot and ham by an orange x. we see that for this dataset in two dimensions, the isomap compression appears more linearly separable than the pca compression, making classification easier for the former. the goal of this paper is not to compete with state-of-the-art models, but, rather, (1) to show that 12-fold dimensionality reductions of contextual embeddings can, in some settings, conserve much of the original performance, (2) to illustrate the efficacy of geodesic similarity metrics in improving the downstream performance of contextual embedding compressions, and (3) propose the creation of more efficient, geodesic-distance-based transformer architectures. in particular, our main result is showing that a 64-dimensional concatenation of compressed pca and isomap embeddings are comparable to the original bert embeddings and outperform our pca baseline. we attribute this success to the locality data preserved by the k-nearest neighbors (k-nn) graph generated by the isomap algorithm. ## related work as best we know, there is very little literature regarding the intersection of contextual embedding compression and geodesic distances. most of the existing work in related spaces deals with noncontextual word embeddings. despite the rapid growth in the popularity of transformers, these embeddings still retain popularity. for non-contextual word embeddings, @xcite propose a post-processing algorithm that projects embedded data away from the dominant principal components, in order to greater differentiate the data. ## method with the goal of reducing the contextual embedding dimensionality, we first processed our data using a pre-trained, uncased bert base model. then, we compressed the data to a lower dimension using both pca and isomap as described in section 3.2. this method aims to capture as much information as possible from the original bert embeddings while preserving graphical locality information and nonlinearities in the final contextual embeddings. ## experiments and results we assess the results of these compression techniques on two text classification datasets. we provide the code for our experiments 1 . ## conclusions and future work we present a novel approach for compressing bert embeddings into effective lower dimension representations. our method shows promise for the inclusion of geodesic locality information in transformers and future compression methods. we hope our results lead to more work investigating the geometric structure of transformer embeddings and developing more computationally efficient nlp training pipelines. to further this work, we plan to investigate the efficacy of (1) other graph dimensionality reduction techniques, (2) non-euclidean distance metrics, and (3) our approach on different transformers. in addition, we would like to investigate whether datasets for other tasks can be effectively projected onto a manifold.