Box Embeddings: An open-source library for representation learning using geometric structures (2021)

Abstract: A fundamental component to the success of modern representation learning is the ease of performing various vector operations. Recently, objects with more geometric structure (eg. distributions, complex or hyperbolic vectors, or regions such as cones, disks, or boxes) have been explored for their alternative inductive biases and additional representational capacity. In this work, we introduce Box Embeddings, a Python library that enables researchers to easily apply and extend probabilistic box embeddings. Fundamental geometric operations on boxes are implemented in a numerically stable way, as are modern approaches to training boxes which mitigate gradient sparsity. The library is fully open source, and compatible with both PyTorch and TensorFlow, which allows existing neural network layers to be replaced with or transformed into boxes easily. In this work, we present the implementation details of the fundamental components of the library, and the concepts required to use box representations alongside existing neural network architectures.

## introduction much of the success of modern deep learning rests on the ability to learn representations of data compatible with the structure of deep architectures used for training and inference @xcite . vectors are the most common choice of representation, as linear transformations are well understood and element-wise non-linearities offer increased representational capacity while be-ing straightforward to implement. recently, various alternatives to vector representations have been explored, each with different inductive biases or capabilities. @xcite represent words using gaussian distributions, which can be thought of as a vector representation with an explicit parameterization of variance. this variance was demonstrated to be capable of capturing the generality of concepts, and kl-divergence provides a natural asymmetric operation between distributions, ideas which were expanded upon in @xcite . @xcite , on the other hand, change the embedding space itself from euclidean to hyperbolic space, where the negative curvature has been shown to provide a natural inductive bias toward modeling tree-like graphs @xcite @xcite . a subset of these alternative approaches explores region-based representations, where entities are not represented by a single point in space but rather explicitly parameterized regions whose volumes and intersections are easily calculated. order embeddings @xcite represent elements using infinite cones in r n + and demonstrate their efficacy of modeling partial orders. @xcite endow order embeddings with probabilistic semantics by integrating the space under a negative exponential measure, allowing the calculation of arbitrary marginal, joint, and conditional probabilities. cone representations are not particularly flexible, however -for instance, the resulting probability model cannot represent negative correlation -motivating the development of probabilistic box embeddings @xcite , where entities are represented by n-dimensional rectangles (i.e. cartesian products of intervals) in euclidean space. probabilistic box embeddings have undergone several rounds of methodological improvements. the original model used a surrogate function to pull disjoint boxes together, which was improved upon in @xcite via gaussian convolution of box indicator functions, resulting in a smoother loss landscape and better performance as a result. dasgupta et al. (2020) improved box training further by using a latent random variable approach, where the corners of boxes are modeled using gumbel random variables. these latter models lacked valid probabilistic semantics, however, a fact rectified in @xcite . while each methodological improvement demonstrated better performance on various modeling tasks, the implementations grew more complex, bringing with it various challenges related to performance and numerical stability. various applications of probabilistic box embeddings (eg. modeling joint-hierarchies @xcite , uncertain knowledge graph representation @xcite , or fine-grained entity typing @xcite ) have relied on bespoke implementations, adding unnecessary difficulty and differences in implementation when applying box embeddings to new tasks. to mitigate this issue and make applying and extending box embeddings easier, we saw the need to introduce a reusable, unified, stable library that provides the basic functionalities needed in studying box embeddings. to this end, we introduce "box embeddings", a fully open-source python library hosted on pypi. the contributions of this work are as follows: ## applications in this section, we demonstrate the box embeddings library by using it to implement models for two real-world tasks: a representation learning task of hierarchical graph modeling @xcite , and the nlp task of natural language inference @xcite . we first demonstrate the intuition behind the containment-based loss function used to train these models using a toy example involving two 2-dimensional boxes. ## conclusion in this paper, we have introduced box embeddings, the first python library focused on allowing regionbased representations to be used with deep learning libraries. our library implements proposed training methods and geometric operations on probabilistic box embeddings in a well-tested and numericallystable fashion. we described the concepts needed to understand and apply this library to novel tasks, and applied the library to graph modeling and natural language inference, demonstrating both shallow and deep contextualized box representations. we hope the release of this package will aid researchers in using region-based representations in their work, and that the well-documented codebase will facilitate additional methodological extensions to probabilistic box embedding models.