Lorentz Direct Concatenation for Stable Training in Hyperbolic Neural Networks
Abstract
Hyperbolic neural networks have achieved considerable success in extracting representation from hierarchical or tree-like data.
However, they are known to suffer from numerical instability, which makes it difficult to build hyperbolic neural networks with deep hyperbolic layers, no matter whether the Poincar´e or Lorentz coordinate system is used.
In this note, we study the crucial operation of concatenating hyperbolic representations.
We propose the
Lorentz direct concatenation and illustrate that it is much more stable than concatenating in the tangent space.
We provide some insights and show superiority of performing direct concatenation in real tasks.
1. Introduction
Recent development of geometric deep learning has witnessed a boost of neural network models specifically designed for data lying in non-Euclidean domains.
For data that show a hierarchical or tree-like structure, hyperbolic neural networks have been effective and successfully applied to, for instance, product recommendation (Wang et al., 2019), drug discovery (Wu et al., 2021), and action recognition (Peng et al., 2020).
One possible approach to building hyperbolic neural networks is performing neural operations in the tangent space.
The earliest such model was proposed by Ganea et al.
(2018), which got refined recently by Shimizu et al.
If the dataset has a graph structure, it is also possible to combine hyperbolic operations in the tangent space with message passing layers in graph convolution (Chami et al., 2019; Liu et al., 2019; Bachmann et al., 2020).
Although tangent spaces are effective ways to approximate the hyperbolic domain, applying exponential and logarithmic operations usually causes numerical instability, especially within the Poincar´e ball model (Nickel and Kiela, 2018; Chami et al., 2019).
Recently, Chen et al.
(2021) proposed to use fully hyperbolic layers, which avoid going back and forth to the tangent spaces.
However, their focus was on linear layers.
More complex neural networks usually include different types of layers which, even combined with fully hyperbolic linear layers, may suffer from numerical instability.
In this note, we propose a concatenation layer within the Lorentz model and illustrate its numerical stability.
Moreover, we analyze the stability and use it to achieve a hyperbolic generative network competitive for real tasks.
We provide some preliminaries of hyperbolic geometry in Appendix A.
2. Motivation and Definition of Lorentz Direct Concatenation
In the Poincar´e model, Shimizu et al.
(2021) proposed Poincar´e β -concatenation and β -split, both of which are numerically unstable in deep networks.
As for the Lorentz model, we remark that one could also define operations in the tangent space similarly to the Poincar´e β -concatenation and β -split.
More specifically, if we want to concatenate the input vectors { x i } Ni =1 where each x i ∈ L n i K , we could follow a “Lorentz tangent concatenation”: first lift each x
i to the tangent space of the origin o : v i
= log K
o ( x i ) = (cid:2) v it v is (cid:3) ∈ R n i +1 , and then perform the Euclidean concatenation to get v := (cid:0) 0 , v ⊤ 1 s , . . .
, v ⊤ N s (cid:1) ⊤ .
Finally, we would get y =
exp K o ( v ) as a concatenated vector in the hyperlolic space.
Similarly, we could perform the “Lorentz tangent split” on an input x i ∈ L nK with split sub-dimensions (cid:80) Ni =1 n i = n to get v = log K
o ( x ) = (cid:0) 0 , v ⊤ 1 s ∈ R n 1 , . .
. , v ⊤ N s ∈ R n N (cid:1) ⊤ , v i = (cid:2)
v is (cid:3) ∈
T o L n i K , and the split vectors
y i =
exp K
o
( v i ) successively.
Unfortunately, both the Lorentz tangent concatenation and the Lorentz tangent split are not “regularized”, which means that the norm of the spatial component will increase after concatenation, and decrease after split.
This will make the hidden embeddings numerically unstable.
While this problem could be solved by adding regulations, a bigger issue with the
Lorentz tangent concatenation and split is that if we use them in a deep neural network, there would be too many exponential and logarithmic maps.
Moreover, the tangent space is chosen at o .
If the points to concatenate are not close to o , their hyperbolic relation may not be captured very well.
Therefore, we abandon the use of the tangent space and propose more direct and numerically stable operations, which we call the “Lorentz direct concatenation and split” and define as follows: given the input vectors { x i } Ni =1 where each x i ∈ L n i K and M = (cid:80) Ni =1 n i , the Lorentz direct concatenation of { x i } Ni =1 is defined to be
a vector y ∈
L Mk given by
[Equation]
Note that each x i s is the spatial component of x i .
If we consider x i ∈
L n i K as a point in
R n i +1 , the projection of x i onto the Euclidean subspace { 0 }× R n , or the closest point there, is x
i s .
The Lorentz direct concatenation can thus be considered as Euclidean concatenation of the projected points.
Lastly, the Euclidean concatenated point is mapped back to L MK by the inverse map of the projection.
We remark that this concatenation directly inherits from the Lorentz model, which is one advantage that the Poincar´e model does not have.
We also define the Lorentz Split for completeness, though our main focus is on concatenation:
Given an input x ∈ L nK , the Lorentz Direct Split of x , with sub-dimensions n 1 , · · · , n N where (cid:80) Ni =1 n i = n , will be { y i } N i =1 , where each y i
∈
L n i K is given by first splitting x in the space dimension as x = (cid:2) x t , y T1 s , · · · , y T N s (cid:3) T , and then calculating the corresponding time dimension as y i = (cid:104) √ ∥ y is ∥ 2 − 1 /K
y is (cid:105) .
3. Numerical Stability
We use the following simple experiment to show the advantage of our Lorentz Direct Concatenation over the Lorentz tangent concatenation.
The hyperbolic neural network in this simple experiment consists of a cascading of L blocks, and the architecture of each block is as follows: for l = 0 , · · · , L − 1,
[Equation]
In our test, we take d = 64.
We sample input and output data from two wrapped normal distributions (Nagano et al., 2019) with different means (input: origin o , output: E2H( 1 64 )) and variances (input: diag( 1 64 ), output: 3 × diag( 1 64 )).
Taking the input as x (0) , we fit x ( L ) to the output data.
We record the average gradient norm of the three hyperbolic linear layers in each block.
The results for L = 64 blocks and L = 128 blocks are shown in Figure 1.
Clearly, for the first 20 blocks, the Lorentz tangent concatenation leads to significantly larger gradient norms.
This difference in norms is clearer when the network is deeper.
The gradients from the Lorentz direct concatenation are more stable.
[Figure]
4. Concatenation and Hyperbolic Distances
In this section, we perform additional analysis of Lorentz Direct Concatenation and Lorentz tangent concatenation, particularly their effect on hyperbolic distances. First, we study the hyperbolic distances to the hyperbolic origin for both concatenation methods. Suppose we have x ∈ L nK and y ∈ L mK . Let z
= HCat( x , y ) ∈
L n + m − 1 K and z ′ =
HTCat( x , y ) ∈ L n + m − 1 K be their hyperbolic direct concatenation and hyperbolic tangent concatenation, respectively. We compare the difference between d L ( z , o ) and d L ( z ′ , o ) as follows. Note that the distance between an arbitrary point x ∈ L nK and the origin only depends on the time component: [Equation]
Hence, the distance information is completely contained in the time component.
After the concatenation, the time component is (cid:112) x 2 t + y 2 t
+ 1 /K .
Consequently, for Lorentz Direct
[Figure]
n = 3 , 6 , 64, respectively.
Column 4: | d L ( x c , y c ) − d L ( x , y ) | − | d L ( x
[Equation]
c , y
[Equation]
c )
[Equation]
L ( x , y ) | .
Concatenation, the distance is
[Equation]
For Lorentz tangent concatenation, since both the logarithmic and exponential maps reserve the distances, one has
[Equation]
Although the hyperbolic distance d L ( z , o ) is not the squared sum of d L ( x , o ) and d L ( y , o ), d L ( z , o ) is larger than each of d L ( x , o ) and d L ( y , o ).
On the other hand, after concatenation, d 2 L ( z ′ , o ) =
d 2 L ( x , o )+ d 2 L ( y , o ).
This relation agrees with the Euclidean concatenation.
However, norm-preservation is not why concatenation works in the Euclidean domain.
Therefore, we don’t consider this as an advantage of the Lorentz tangent concatenation.
The Lorentz Direct Concatenation is more efficient and stable, and no information is lost during concatenation.
Therefore, it is still preferred as a neural layer.
More importantly, we study how concatenation changes the relative distances, which is closely related to stability.
Specifically, we perform the following experiments.
Given x , y , c ∈ L nK , let x c = HCat( x , c ) and y c
=
HCat( y , c ) be the direct-concatenated version of ( x , c ) and ( y , c ), respectively.
Similarly we denote x ′ c =
HTCat( x , c ) and y ′ c =
HTCat( y , c ) for Lorentz tangent concatenations.
Since the same vector c is attached to x and y , we naturally hope d L ( x c , y c ) and d L ( x ′ c , y ′ c ) do not deviate much from d L ( x , y ).
We describe our experiments as follows.
Take K = − 1.
We randomly sample three points independently from L nK as x , y and c respectively.
We have two scenarios for sampling the points: (1) “spatial normal”: the points are sampled so that their spatial components follow the standard normal distribution; (2) “wrapped normal”: the points are sampled from the wrapped normal distribution with unit variance.
In each scenario, for n ∈ { 3 , 16 , 64 } , we do the experiments for 10,000 times.
We report the distances | d L ( x c , y c ) − d L ( x , y ) | and | d L ( x ′ c , y ′ c ) − d L ( x , y ) | as well as their differences in Figure 2.
Our experiments clearly show that, especially for large dimensions, the distance between d L ( x c , y c ) and d L ( x , y ) is smaller than the distance between d L ( x ′ c , y ′ c ) and d L ( x , y ).
In particular, in many cases, | d L ( x c , y c ) − d L ( x , y ) | is around zero.
On the other hand, | d L ( x ′ c , y ′ c ) − d L ( x , y ) | tend to be large when n = 16 , 64, especially when samples follow the wrapped normal distribution.
From this result, the Lorentz Direct Concatenation should be preferred to the Lorentz tangent concatenation.
In particular, the significant expansion of distance when concatenating with the same vector, in the case n = 64, may be one cause of numerical instability.
5. Applications of Lorentz Direct Concatenation
In this section, we show that Lorentz direct concatenation can be used in deep and complex neural network architecture with numerical stability.
We consider the task of molecular generation using the MOSES dataset (Polykovskiy et al., 2020).
Our model contains a tree decoder and a graph decoder are used in the hyperbolic space (the detailed network structure is presented in Appendix B).
In both decoders the hyperbolic features need to be concatenated and we consider the following methods for concatenation: β -concatenation (Shimizu et al., 2021), the Lorentz tangent concatenation and the Lorentz direct concatenation.
The results of molecular generation are presented in Table 1.
While Lorentz direct concatenation saturates the important metrics of validity, uniqueness and novelty, the other two concatenation methods both suffer from numerical instability, even if a fully hyperbolic approach is adopted in all the methods.
[Table]
Future work involves more theoretical analysis on the numerical stability of hyperbolic models [Table] as well as models that enjoy both expressivity and stability [Table]
The research results of this article are sponsored by the Kunshan Municipal Government research funding.
Appendix A. Preliminaries of Hyperbolic Geometry
Hyperbolic geometry is a special kind of Riemannian geometry with a constant negative curvature (Cannon et al., 1997; Anderson, 2006).
There are five models (coordinate systems) of the hyperbolic space: the Lorentz model, the Poincar´e ball model, the Hemisphere model, the Klein model, and the Poincar´e half-space model.
The models are isometric and the relation between their coordinates is surveyed by Dai et al.
We work with the
Lorentz model for its numerical stability.
The Lorentz Model
The Lorentz model L nK =
( L , g K ) of an n dimensional hyperbolic space with constant negative curvature K is an n -dimensional manifold L embedded in the ( n + 1)-dimensional Minkowski space, together with the Riemannian metric tensor g K = diag([ − 1 , 1 ⊤ n ]), where 1 n denotes the n -dimensional vector whose entries are all 1’s.
Every point in L nK is represented by x = (cid:20) x
t x s (cid:21) , x t > 0 , x s ∈ R n and satisfies ⟨ x , x ⟩ L = 1 /K , where ⟨· , ·⟩ L is the Lorentz inner product induced by g K :
[Equation]
Geodesics and Distances
Geodesics are shortest paths in a manifold, which generalize the notion of “straight lines” in Euclidean geometry.
In particular, the length of a geodesic in L nK (the “distance”) between x , y ∈ L nK is given by
[Equation]
Tangent Space
For each point x ∈ L nK , the tangent space at x is T x L nK := { y ∈ R n +1 | ⟨ y , x ⟩ L = 0 } .
It is a first order approximation of the hyperbolic manifold around a point x and is a subspace of R n +1 .
We denote ∥ v ∥ L = (cid:112) ⟨ v , v ⟩ L as the norm of v ∈ T x L nK .
Exponential and Logarithmic Maps
The exponential and logarithmic maps are maps between hyperbolic spaces and their tangent spaces.
For x , y ∈ L nK and v ∈ T x
L nK , the exponential map exp K x
( v ) :
T x L nK → L nK maps tangent vectors to hyperbolic spaces by assigning v to the point exp K x ( v ) :
= γ (1), where γ is the geodesic satisfying γ (0) = x and γ ′ (0) = v .
Specifically,
[Equation]
The logarithmic map log K x ( y ) :
L nK → T x L nK is the inverse map that satisfies log K x (exp K x ( v ))
= v .
Specifically,
[Equation]
Appendix B. Network Details in Section
Our network contains a hyperbolic auto-encoder that learns the hyperbolic embedding of molecules, and a hyperbolic generative network that learns to sample de novo latent embeddings.
B.1. Details of Hyperbolic Auto-Encoder Hyperparameters
• Manifold curvature: K = − 1 .•
For all hyperbolic linear layers: – Dropout: 0.
– Use bias: True • Optimizer: Riemannian Adam ( β 1 = 0 . 0 , β 2 = 0 .
999) • Learning rate: 5e- • Learning rate scheduler: StepLR (step = 20000, γ = 0 . 5) • Batch size:• Number of epochs:
Graph Encoder
• Input: graph node features dimension:• Map features to hyperbolic space: R 35 → L 35 K • Hyperbolic GCN layers: – Input dimension:– Hidden dimension:– Depth:– Output dimension:• Hyperbolic centroid on all vertices • Output: graph embedding in L 256 K
Tree Encoder
• Input: junction tree features dimension• Hyperbolic embedding layer: R 828 → L 256 K • Hyperbolic GCN layers: – Input dimension:– Hidden dimension:– Depth:– Output dimension:• Hyperbolic centroid on all vertices • Output: tree embedding in L 256 K
Tree Decoder
• Input: tree embedding in L 256 K • Message passing RNN: – Input: node feature of current tree node, inward messages – Hyperbolic linear layer on inward messages:
L 256 K → L 256 K – Hyperbolic centroid on inward messages – Hyperbolic embedding layer on node feature: R 828 → L 256 K – Concatenation on node feature and inward message: L 256 K → L 512 K – Hyperbolic linear layer: L 512 K → L 256 K – Output dimension:• Topological Prediction: – Input: tree embedding, node feature of current tree node, inward messages – Hyperbolic linear layer on inward messages:
L 256 K → L 256 K – Hyperbolic centroid on inward messages – Hyperbolic embedding layer on tree feature: R 828 → L 256 K – Concatenation on node feature, inward message, and tree embedding: L 256 K → L 768 K – Hyperbolic linear layer: L 768 K → L 256 K – Hyperbolic centroid distance layer:
L 256 K → R– Softmax on output – Output dimension:• Label Prediction: – Input: tree embedding, outward messages – Concatenation on outward message, and tree feature: L 256 K
→ L 512 K – Hyperbolic linear layer: L 512 K → L 256 K – Hyperbolic centroid distance layer:
L 256 K → R– Softmax on output – Output dimension:• Output: junction tree
Graph Decoder
• Input: junction tree, tree message, and graph embedding • Construction candidate subgraphs • Hyperbolic graph convolution layers on all subgraphs: – Input dimension:– Hidden dimension:– Depth:– Output dimension:• Hyperbolic centroid on vertices of all subgraphs • Concatenation on subgraph embedding and graph embedding: L 256 K → L 512 K • Hyperbolic linear layer:
L 512 K → L 256 K • Hyperbolic centroid distance layer:
L 256 K
→ R • Use subgraph score to construct molecular graph • Output: molecular graph
B.2. Details of Hyperbolic Generative Adversarial Network Hyperparameters
• Manifold curvature: K = − 1 .•
Gradient penalty coefficient: λ =• For all hyperbolic linear layers: – Dropout: 0.
– Use bias: True • Optimizer: Riemannian Adam ( β 1 = 0 , β 2 = 0 . 9) • Learning Rate: 1e- • Batch size:• Number of epochs:• Gradient penalty λ :
Generator
• Input: points sampled from wrapped normal distribution G ( o , diag( 1 128 )) in L 128 K • Hyperbolic linear layers for graph embedding: – Input dimension:– Hidden dimension:– Depth:– Output dimension:• Hyperbolic linear layers for tree embedding: – Input dimension:– Hidden dimension:– Depth:– Output dimension:• Output: graph embedding and tree embedding in L 128 K Discriminator • Input: graph embedding and tree embedding in L 128 K • Hyperbolic linear layers for graph embedding: – Input dimension:– Hidden dimension:– Depth:– Output dimension:• Hyperbolic linear layers for tree embedding: – Input dimension:– Hidden dimension:– Depth:– Output dimension:• Lorentz Direct concatenation on graph embedding and tree embedding:
L 256 K → L 512 K • Hyperbolic linear layer: L 512 K → L 256 K • Hyperbolic centroid distance layer:
L 256 K → R • Output: score in R