Box Embeddings: Paper Overviews

Tamanna Hossain-Kay / 2021-11-09

Geometric embeddings can be used to represent transitive asymmetric relationships. Box-embeddings are n-dimensional hyperrectangles that can represent such relationships by containment.

Source: Box Embeddings: An open-source library for representation learning
using geometric structures — **Source:** Box Embeddings: An open-source library for representation learning using geometric structures

Formally, a “box” is defined as a Cartesian product of closed intervals, $\begin{aligned} B (θ) & = \prod_{i = 1}^{n} [z_{i} (θ), Z_{i} (θ)] \\ = [z_{1} (θ), Z_{1} (θ)] \times \dots \times [z_{n} (θ), Z_{n} (θ)] \\ = [x_{m, 1}, x_{M, 1}] \times \dots \times [x_{m, n}, x_{M, n}] \end{aligned}$

Parameterization: $θ$ may be the output from a neural network

Python package: Box Embeddings: An open-source library for representation learning using geometric structures

In order to learn box-embeddings, intersections and volumes of boxes need to be computed. This can be done in a few different ways:

Hard Intersection + Hard Volume: Probabilistic Embedding of Knowledge Graphs with Box Lattice Measures (Vilnis et al., 2018)
Hard Intersection + Soft Volume: Smoothing the Geometry of Probabilistic Box Embeddings (Li et al., 2018)
Gumbel Intersection + Bessel Approx. Volume: Improving Local Identifiability in Probabilistic Box Embeddings (Dasgupta et al., 2020)

Read on for brief overviews of each method.

Probabilistic Embedding of Knowledge Graphs with Box Lattice Measures

Authors: Luke Vilnis, Xiang Li, Shikhar Murty, Andrew McCallum; University of Massachusetts Amherst

Method: Hard Intersection, Hard Volume

Defining Box operations

$\begin{array}{l} x \land y =⊥ if x and y disjoint, else \\ x \land y = \prod_{i} [max (x_{m, i}, y_{m, i}), min (x_{M, i}, y_{M, i})] \\ x \lor y = \prod_{i} [min (x_{m, i}, y_{m, i}), max (x_{M, i}, y_{M, i})] \end{array}$

Note,

While boxes are closed under intersection they’re not closed under union.
$(x \land y) = (x \cap y$ ), but $(x \lor y) > (x \cup y)$

Volume can be computed as,

$V o l (B (θ)) = \prod_{i} max (x_{M, i} - x_{m, i}, 0) = \prod_{i} m_{h} (x_{M, i} - x_{m, i})$

where $m_{h}$ is the standard hinge function.

Volume can be interpreted as probability. For a concept $c$ associated with box $B (θ)$ ,

$P (c) = V o l (B (θ))$

Learning WordNet

The transitive reduction of the WordNet noun hierarchy contains 82,114 entities and 84,363 edges.

The learning task is an edge classification task. Given a pair of nodes $(h, t)$ , the probability of existence of the edge $h_{i} \to t_{i}$ is computed as $P (h_{i} \to t_{i}) = P (h_{i} | t_{i}) = \frac{P (h_{i} \cap t_{i})}{P (t_{i})} = \frac{Vol (B (θ_{h_{i}}) \cap B (θ_{t_{i}}))}{Vol (B (θ_{t_{i}}))}$

Gold label: Existence $(h_{i} \to t_{i}) ⟺ P (h_{i} | t_{i}) = 1$

Loss: Cross-entropy

Surrogate Intersection: When boxes do not intersect, the joint probability is 0 so parameters cannot be updated. In this case, a lower bound is used:

$\begin{aligned} p (a \cap b) & = p (a) + p (b) - p (a \cup b) \\ \geq p (a) + p (b) - p (a \lor b) \end{aligned}$

WordNet Results

Box-embeddings perform better than baselines in terms of classification accuracy. Refer to paper for details on baseline models.

Smoothing the Geometry of Probabilistic Box Embeddings

Authors: Xiang Li, Luke Vilnis,Dongxu Zhang, Michael Boratko, Andrew McCallum; University of Massachusetts Amherst

Method: Hard Intersection, Soft Volume

As mentioned in the previous paper (Vilnis et al., 2018), the gradient cannot flow when boxes disjoint. Vilnis et al., 2018 address this issue using a surrogate intersection probability. Here an alternative is proposed using using a Gaussian convolution with the original boxes to smooth the box edges into density functions.

Define boxes as indicator functions of intervals. Volume of boxes can then be expressed as integrals over the products of the indicator functions. Consider the 1D case for ease of exposition and results will generalize to higher dimensions.

The volume of the intersection of two boxes representing the intervals $x = [a, b]$ and $y = [c, d]$ (representing the joint probability between intervals) can be expressed as,

$p (x \land y) = \int_{R} 1_{[a, b]} (x) 1_{[c, d]} (x) d x$

The smoothed version is,

$p_{ϕ} (x \land y) = \int_{R} f (x; a, b, σ_{1}^{2}) g (x; c, d, σ_{2}^{2}) d x$

where,

$f (x; a, b, σ^{2}) = 1_{[a, b]} (x) * ϕ (x; σ^{2}) = \int_{R} 1_{[a, b]} (z) ϕ (x - z; σ^{2}) d z = \int_{a}^{b} ϕ (x - z; σ^{2}) d z$

and $ϕ (z; σ^{2}) = \frac{1}{σ \sqrt{2 π}} e^{\frac{- z^{2}}{2 σ^{2}}} (Gaussian PDF)$

Integral Solution:

$\begin{aligned} p_{ϕ} (x \land y) & = σ (m_{Φ} (\frac{b - c}{σ}) + m_{Φ} (\frac{a - d}{σ}) - m_{Φ} (\frac{b - d}{σ}) - m_{Φ} (\frac{a - c}{σ})) \end{aligned}$ where $σ = \sqrt{σ_{1}^{2} + σ_{2}^{2}}$ , $soft (x) = \log (1 + \exp (x))$ is the softplus function, the integral of the logistic sigmoid, and $ρ = \frac{σ}{1.702}$

Proof

Let $τ := \frac{1}{\sqrt{σ_{1}^{2} + σ_{2}^{2}}}$ , and note the identity from (Jebara et al., 2004; Vilnis & McCallum, 2015):

$\int_{R} ϕ (x - μ_{1}; σ_{1}^{2}) ϕ (x - μ_{2}; σ_{2}^{2}) d x = ϕ (μ_{1} - μ_{2}; σ_{1}^{2} + σ_{2}^{2})$

Then, $\begin{aligned} p_{ϕ} (x \land y) & = \int_{R} \int_{a}^{b} ϕ (x - y; σ_{1}^{2}) d y \int_{c}^{d} ϕ (x - z; σ_{2}^{2}) d z d x Definition f, g \\ = \int_{a}^{b} \int_{c}^{d} \int_{R} ϕ (x - y; σ_{1}^{2}) ϕ (x - z; σ_{2}^{2}) d x d y d z Fubini's Theorem \\ = \int_{c}^{d} \int_{a}^{b} ϕ (y - z; τ^{- 2}) d y d z (Jebara et al., 2004; Vilnis & McCallum, 2015) \\ = \int_{c}^{d} \int_{a}^{b} Φ^{'} (τ (y - z)) τ d y d z Chain Rule; Φ is the standard normal CDF \\ = \int_{c}^{d} Φ (τ (b - z)) - Φ (τ (a - z)) d z \\ = \frac{- 1}{τ} (m_{Φ} (τ (b - d)) - m_{Φ} (τ (a - d)) - m_{Φ} (τ (b - c)) + m_{Φ} (τ (a - c))) \end{aligned}$ Then with $σ = τ^{- 1}$ , $p_{ϕ} (x \land y) = σ (m_{Φ} (\frac{b - c}{σ}) + m_{Φ} (\frac{a - d}{σ}) - m_{Φ} (\frac{b - d}{σ}) - m_{Φ} (\frac{a - c}{σ}))$

Approximation

Multiplication of Gaussian-smoothed indicators have the property $f^{2} \neq f$ so for cases $ρ > 0$ ,

Idempotency requirement of bounded lattices is violated: $a \land a = a \lor a = a$
Practically then if we are to treat the outputs of $p_{ϕ}$ as probabilities, the consequence is $p_{ϕ} (x ∣ x) = \frac{p_{ϕ} (x, x)}{p_{ϕ} (x)} = \frac{p_{ϕ} (x \land x)}{p_{ϕ} (x)} \neq 1$

So we will approximate to get a version that is idempotent.

$\begin{aligned} p_{ϕ} (x \land y) & \approx (ρ soft (\frac{b - c}{ρ}) + ρ soft (\frac{a - d}{ρ})) - (ρ soft (\frac{b - d}{ρ}) + ρ soft (\frac{a - c}{ρ})) \\ \approx (soft (b - c) \lor soft (a - d)) \land (soft (b - d) \lor soft (a - c)) = soft (b \land d - a \lor c) \\ = soft (b \land d - a \lor c) \end{aligned}$

The first approximation follows from the approximation of $Φ$ by a logistic sigmoid given in Bowling et al. (2009).

However, it is still not idempotent except in the zero-temperature limit. Since as

ρ

goes to zero, we recover the original formula with the hinge function which is idempotent,

\begin{aligned} p_{ϕ} (x \land y) & = lim_{ρ \to 0} (ρ soft (\frac{b - c}{ρ}) + ρ soft (\frac{a - d}{ρ})) - (ρ soft (\frac{b - d}{ρ}) + ρ soft (\frac{a - c}{ρ})) \\ = (m_{h} (b - c) + m_{h} (a - d)) - (m_{h} (b - d) + m_{h} (a - c)) \\ = m_{h} (b \land d - a \lor c) hinge function on an interval \end{aligned}

The limit

(From Wolfram Alpha) Two-sided limit does not exist.

Limit from the left: $lim_{ρ \to 0^{-}} ρ \log (1 + \exp (\frac{y}{ρ})) = 0$ for $y > 0$

Limit from the right: $lim_{ρ \to 0^{+}} ρ \log (1 + \exp (\frac{y}{ρ})) = y$ for $y > 0$

The second approximation is needed to get a similar idempotent property when using the softplus fuinction as when using the hinge function,

$(m_{h} (b - c) + m_{h} (a - d)) - (m_{h} (b - d) + m_{h} (a - c)) = m_{h} (b \land d - a \lor c)$

$(soft (b - c) \lor soft (a - d)) \land (soft (b - d) \lor soft (a - c)) = soft (b \land d - a \lor c)$

Normalization

But finally, because softplus upper-bounds the hinge function, it is capable of outputting values that are greater than 1 , and therefore must be normalized.

Thus, in the multi-dimensional case we have:

$\begin{aligned} p (x) & \approx \prod_{i} m_{s o f t}^{(i)} (x_{M, i} - x_{m, i}) \\ p (x, y) & \approx \prod_{i} m_{s o f t}^{(i)} (x_{M, i} \land y_{M, i} - x_{m, i} \lor y_{m, i}) \end{aligned}$

where,

$m_{s o f t}^{(i)} (x) = \frac{soft (\frac{x}{ρ})}{soft (\frac{G_{m} - G_{m}}{ρ})}$

and $(G_{m}^{(i)}, G_{M}^{(i)})$ = Global minimum and maximum in $i^{t h}$ dimension.

WordNet Results

Sparsity is increased (positive:negative ratio increased) in the data to create more disjoint cases. In sparser regimes, the smoothed box greatly outperforms the baselines (original hard box, and OE).

Improving Local Identifiability in Probabilistic Box Embeddings

Authors: Shib Sankar Dasgupta, Michael Boratko, Dongxu Zhang, Luke Vilnis, Xiang Lorraine Li, Andrew McCallum; University of Massachusetts, Amherst

Method: Gumbel Intersection, Bessel Approx. Volume

Both of the above proposed methods can have problems with local indentifiability since many different configurations of boxes can lead to the same joint probability value. This can lead to pathological parameter settings during train (Figure 2c; Figure 3c).

The method proposed here defines box edges as random variable, which leads to probabilities as expected volumes. This densifies the gradient of joint probabilities.

Gumbel Distribution

Since computing intersections of boxes involves computing minima and maxima, a min-max stable distribution is chosen: Gumbel distribution. This means that the maximum of multiple Gumbel random variables also follows a Gumbel distribution (MaxGumbel), and the minimum of multiple Gumbel random variables also follows a Gumbel distribution (MinGumbel).

Specifically,

If $X_{i} \sim MaxGumbel (μ_{x_{i}}, β)$ then $\begin{aligned} max (X_{i}) & \sim MaxGumbel (β \ln (\sum_{i} e^{\frac{μ_{x_{i}}}{β}}), β) \\ = MaxGumbel (β LogSumExp (\frac{μ_{x_{i}}}{β})) \end{aligned}$ and if $X_{i} \sim MinGumbel (μ_{x}, β)$ , then $\begin{aligned} min (X_{i}) & \sim MinGumbel (- β \ln (\sum_{i} e^{- \frac{μ_{x_{i}}}{β}}), β) \\ = MinGumbel (- β LogSumExp (- \frac{μ_{x_{i}}}{β})) \end{aligned}$

Gumbel-box process

For a probability space

(Ω_{Box}, E, P_{Box})

, with

Ω_{B o x} \subseteq R^{d}

, the Gumbel-box process is generated as

\begin{aligned} β & \in R_{+}, μ^{i, \lor} \in Ω_{B o x} μ^{i, \land} \in Ω_{B o x} \\ X_{j}^{i, \land} & \sim MaxGumbel (μ_{j}^{i, \land}, β), X_{j}^{i, \lor} \sim MinGumbel (μ_{j}^{i, \lor}, β) \\ Box (X^{i}) & = \prod_{j = 1}^{d} [X_{j}^{i, \land}, X_{j}^{i, \lor}], X^{i} = 1_{Box (X^{i})} \\ x^{1}, \dots, x^{n} & \sim P (X^{1}, \dots, X^{n}) \end{aligned}

For a uniform base measure on

R^{d}

, each

X^{i}

is distributed according to

X^{i} \sim Bernoulli (\prod_{j = 1}^{d} (X_{j}^{i, \lor} - X_{j}^{i, \land}))

Intersection

$\begin{aligned} x \land y & = \prod_{i} [μ_{max_{X^{t} \in T} X_{j}^{t, \lor}}, μ_{min_{X^{t} \in T} X_{j}^{t, \land}}] \\ = \prod_{i} [β LogSumExp (\frac{μ_{j}^{t, \lor}}{β}), - β {LogSumExp}_{X^{t} \in T} (- \frac{μ_{j}^{t, \land}}{β})] \end{aligned}$

Volume Calculation

The expected length of an interval $[X, Y]$ where $X \sim MaxGumbel (μ_{x}, β)$ and $Y \sim MinGumbel (μ_{y}, β)$ is given by $E [max (Y - X, 0)] = 2 β K_{0} (2 e^{- \frac{μ_{y} - μ_{x}}{2 β}})$ where $K_{0}$ is the modified Bessel function of second kind, order $0 .$

Proof

$\begin{aligned} E [max (Y - X, 0)] = & \iint_{{y > x}} (y - x) f_{\land} (y; μ_{y}) f_{\lor} (x; μ_{x}) d x d y \\ = & \int_{- \infty}^{\infty} \int_{- \infty}^{y} y f_{\land} (y; μ_{y}) f_{\lor} (x; μ_{x}) d x d y - \\ \int_{- \infty}^{\infty} \int_{x}^{\infty} x f_{\land} (y; μ_{y}) f_{\lor} (x; μ_{x}) d y d x \\ = & \int_{- \infty}^{\infty} y f_{\land} (y; μ_{y}) F_{\lor} (y; μ_{x}) d y - \\ \int_{- \infty}^{\infty} x f_{\lor} (x; μ_{x}) (1 - F_{\land} (x; μ_{y})) d x \\ \int_{- \infty}^{\infty} y f_{\land} (y; μ_{y}) F_{\lor} (y; μ_{x}) d y - \\ = & \int_{- \infty}^{\infty} x f_{\lor} (x; μ_{x}) F_{\lor} (- x; - μ_{y}) d x \end{aligned}$ We note that this is equal to $E_{Y} [Y \cdot F_{\lor} (Y; μ_{x})] - E_{X} [X \cdot F_{\lor} (- X; - μ_{y})]$ however this calculation can be continued, combining the integrals, as $= \int_{- \infty}^{\infty} z f_{\land} (z; μ_{y}) F_{\lor} (z; μ_{x}) - z f_{V} (z; μ_{x}) F_{V} (- z; - μ_{y}) d z$ We integrate by parts, letting $u = z, d v = \frac{1}{β} (e^{\frac{π - μ_{μ}}{β}} - e^{- \frac{\tilde{-} μ_{-}}{β}}) \exp (- e^{\frac{π - μ_{ν}}{β}} - e^{- \frac{\tilde{-} μ_{2}}{β}}) d z$ , $d u = d z, v = - \exp (- e^{\frac{x - μ_{μ}}{β}} - e^{- \frac{x - μ π}{β}})$ where the $d v$ portion made use of the substitution $w = e^{\frac{\tilde{-} μ_{μ}}{β}} + e^{- \frac{- μ - μ}{β}}, d w = \frac{1}{β} (e^{\frac{\tilde{-} μ_{μ}}{β}} - e^{- \frac{\tilde{- μ}}{β}}) d z$ We have, therefore, that (12) becomes $= - {z \exp (- e^{\frac{z - μ_{μ}}{β}} - e^{- \frac{\tilde{-} μ}{β}}) |}_{- \infty}^{\infty} + \int_{- \infty}^{\infty} \exp (- e^{\frac{x - μ_{μ}}{β}} - e^{- \frac{\tilde{-} μ}{β}}) d z$ The first term goes to $0 .$ Setting $u = \frac{z - (μ_{y} + μ_{z}) / 2}{β}$ the second term becomes $β \int_{- \infty}^{\infty} \exp (- e^{\frac{μ_{x} - μ_{u}}{2 \bar{β}}} (e^{u} + e^{- u})) d u = 2 β \int_{0}^{\infty} \exp (- 2 e^{\frac{μ_{x} - μ_{u}}{2 \bar{N}}} \cosh u) d u$ With $z = 2 e^{\frac{μ_{x} - μ_{y}}{2, S}}$ , this is a known integral representation of the modified Bessel function of the second kind of order zero, $K_{0} (z)$ [4, Eq. 10.32.9], which completes the proof of proposition 1 .

Let $m (x) = E [m a x (x, 0)] = 2 β K_{0} (2 e^{- \frac{x}{2 β}}) .$ Then the expected volume is, $\begin{aligned} E [\prod_{i = 1}^{d} m a x (min_{X^{t} \in T} X_{j}^{t, \land} - max_{X^{t} \in T} X_{j}^{t, \lor}, 0)] \\ = \prod_{i = 1}^{d} E [m a x (min_{X^{t} \in T} X_{j}^{t, \land} - max_{X^{t} \in T} X_{j}^{t, \lor}, 0)] rearrange product and expectation \\ = \prod_{i = 1}^{d} m (min_{X^{t} \in T} X_{j}^{t, \land} - max_{X^{t} \in T} X_{j}^{t, \lor}) definition of m \\ = \prod_{i = 1}^{d} 2 β K_{0} (2 e^{\frac{- β {LogSumExp}_{X^{t} \in T} (- \frac{μ_{j}^{t, \land}}{β}) - β LogSumExp (\frac{μ_{j}^{t, \lor}}{β})}{2 β}}) solution for m \end{aligned}$ where $T$ is a set of variables taking the value 1, and $μ^{t, \land}$ and $μ^{t, \lor}$ are the vectors of location parameters for the minimum and maximum coordinates of a $Box (X^{t})$ .

Approximation

Computing conditional probabilities exactly requires the conditional expectation, $E [P (x^{i}, x^{j}, Z) / P (x^{j}, Z)]$ which is intractable since it is a quotient. Instead, we approximate it by its first-order Taylor expansion, which takes the form of a ratio of two simple expectations: $E [P (x^{i}, x^{j}, Z)] / E [P (x^{j}, Z)]$

Softplus Approximation

Due to extreme similarity of the exact volume to the softplus function, the following approximation is used:

$m (x) \approx β \log (1 + \exp (x / β - 2 γ)$

WordNet Results

Gumbel boxes perform better than baselines on WordNet’s noun hierarchy in terms of F1 score. Refer to paper for details on baseline models.