I made a Riemannian line search optimizer with strong Wolfe conditions.

It's not yet perfect. I think it makes some redundant calls to closure during stepping, and when it's close to a local minimum it suffers from numerical errors and can sometimes take strange steps

Huray, we are ready to start.

Image from here

Interesting reading I've done so far:

Hyperbolic Networks Hyperbolic Entailment Cones for Learning Hierarchical Embeddings Poincaré GloVe: Hyperbolic Word Embeddings

Some implementation takeouts (mostly from here):

• Project results of all operations in the ball of radius 1 − eps, where eps = 10^{-5}
• Numerical errors also appear when hyperbolic vectors get closer to 0, perturb with eps = 10^{-15}
• Pass clipped to [-15, 15] input to tanh and clip tanh^{-1} to [-1+eps, 1-eps]

CC @leuchine!

Hi,

I subclassed the ProductManifold and created a StereographicProductManifold. Here I added the functions dist2plane, expmap0and mobius_add by calling the respective functions in the underlying Stereographic manifolds. I added some test functions for this. Additionally, I added wrapped_normal to Stereographic as an alternative random function and I added scipy to the requirements in setup.py as mentioned in #161

Currently the c parameter in PoincareBall defines negative curvature though it is most natural to assume c to be just curvature. That is unconventional, misleading, and is likely to cause inconsistencies some time later

I currently use the following piece of code:

if self.init_id:
    init = torch.eye(input_dims[0]).view((input_dims[0], input_dims[0], 1, 1))
else:
    init = torch.randn(input_dims[0], input_dims[0])# + input_dims[0]*torch.eye(input_dims[0])
self.orth_w = geoopt.ManifoldParameter(init, manifold=stiefel_man)

but if init is random, it won't be on the stiefel-manifold! So i tried self.orth_w.proj_(), but it tells me:

RuntimeError: a leaf Variable that requires grad has been used in an in-place operation.

I can guard the .proj_() with a with torch.no_grad(), but I don't know whether this is the right approach.

I just want to do a straightforward matrix-multiplication on the forward pass and want geoopt to not leave the stiefel-manifold during my training.

There are some numerical issues, and I'm starting to think that there is no easy solution for it at this point. Maybe something like this paper should be investigated more.

• [x] Add curvature
• [x] Add tests
• [x] Add an example with hgnn/attention [example with basic usage]
• [x] Add docs

Example with optimization should be added on the separate PR

I am working with a model similar to the example below:

class Model(nn.Module):
    def __init__(self, word2vec):
        super(Model, self).__init__()
        self.word_lut = gt.ManifoldParameter(word2vec, manifold=gt.PoincareBall())

When I move the model to CUDA, and run it there is a problem during the optimization (See full traceback below [1]): RuntimeError: Expected object of backend CUDA but got backend CPU for argument #2 'other'

The reason of the error is that when I move the model to CUDA, the tensor in the ManifoldParameter is moved, but its manifold is not, therefore the attribute tensor self.c of the manifold is still allocated in the CPU.

I solved it doing this:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")    
for p in model.parameters():
    if isinstance(p, (ManifoldParameter, ManifoldTensor)):
        p.manifold.to(device)

I don't know if issue #49 was an attempt to solve this and I didn't moved the model in the proper way to CUDA or this is not completely documented. In case that the explanation on how to do this properly exists and I couldn't find it, my apologies for creating this issue.

[1] Full traceback

Traceback (most recent call last):
  File "./train.py", line 122, in <module>
    main()
  File "./train.py", line 113, in main
    coach.train()
  File "/home/lopezfo/projects/hyfi/hyfi/Coach.py", line 43, in train
    train_loss = self.train_epoch(epoch)
  File "/home/lopezfo/projects/hyfi/hyfi/Coach.py", line 99, in train_epoch
    self.optim.step()
  File "/home/lopezfo/anaconda3/envs/hyfi/lib/python3.6/site-packages/geoopt/optim/radam.py", line 145, in step
    state["exp_avg_sq"],
  File "/home/lopezfo/anaconda3/envs/hyfi/lib/python3.6/site-packages/geoopt/optim/tracing.py", line 34, in partial
    step(manifold, *args, **kwargs)
  File "/home/lopezfo/anaconda3/envs/hyfi/lib/python3.6/site-packages/geoopt/optim/radam.py", line 191, in perform_step
    point, -step_size * direction, exp_avg
  File "/home/lopezfo/anaconda3/envs/hyfi/lib/python3.6/site-packages/geoopt/manifolds/poincare/__init__.py", line 110, in retr_transp
    y = self.retr(x, u)
  File "/home/lopezfo/anaconda3/envs/hyfi/lib/python3.6/site-packages/geoopt/manifolds/poincare/__init__.py", line 69, in retr
    return math.project(approx, c=self.c)
  File "/home/lopezfo/anaconda3/envs/hyfi/lib/python3.6/site-packages/geoopt/manifolds/poincare/math.py", line 78, in project
    return _project(x, c, dim, eps)
  File "/home/lopezfo/anaconda3/envs/hyfi/lib/python3.6/site-packages/geoopt/manifolds/poincare/math.py", line 86, in _project
    cond = norm > maxnorm
RuntimeError: Expected object of backend CUDA but got backend CPU for argument #2 'other'

I'm trying to implement the SPD manifolds, which is widely used in many papers.

Current commits have implemented all abstract methods of base.Manifold, while more method implementations are on the way. Some test has been done locally and constructions of the testing module are also on the way.

Progress

• [x] move symmetric matrix operations to batch_linalg
• [x] keepdims functionality for inner and _stein_metric.
• [x] mention in documentation about SPD manifolds.
• [x] implementation for random and origin
• [x] manifold test module.
  • [x] shape case for a symmetric positive-definite matrix.
  • [x] test basic operation of SPD manifolds.
  • [x] simple optimization problem on SPD manifolds.
• [ ] mention the PR in the CHANGELOG.

Reference Implementation

• psd.py in pymanopt
• spd.py in matrix-manifolds
• psd.py in geotorch
• spd_matrices.py in geotorch

Some paper using SPD Manifolds

• Computationally Tractable Riemannian Manifolds for Graph Embeddings [arxiv]
• A Riemannian Network for SPD Matrix Learning [arxiv]

All suggestions will be accepted with an open mind. Hoping this pull request will be merged when all works have done.

First of all, thank you for this library!

Description of the bug

When training with Distributed Data Parallel (DDP), the gradient between different devices is not correctly synchronized when using RiemannianSGD (or RiemannianAdam). Replacing it with a standard torch.optim.SGD works well. Note that when using DDP the gradient is synchronized during .backprop() (see this link).

To Reproduce

Simple code training on ImageNet:

import os

import geoopt
import torch
import torch.distributed
import torch.multiprocessing as mp
import torchvision
import torchvision.models as models
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils import data
from torchvision import transforms


def process_ddp(master_port, local_rank, world_size):

    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = str(master_port)
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    torch.distributed.init_process_group("nccl", rank=local_rank, world_size=world_size, init_method='env://')
    assert world_size == torch.distributed.get_world_size()

    return device


def main(local_rank, world_size):

    path_dataset = '/path/to/ImageNet'  # Any other dataset should result in a similar behavior
    master_port = 9999

    device = process_ddp(master_port, local_rank, world_size)

    model = models.resnet18()
    model = model.to(device)

    # optimizer = geoopt.optim.RiemannianSGD(model.parameters(), lr=0.1, stabilize=10)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

    # Data parallelization
    model = DDP(model, device_ids=[local_rank], output_device=local_rank)

    # Prepare dataset
    transform = transforms.Compose([
        transforms.CenterCrop(size=256),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    dataset = torchvision.datasets.ImageNet(split='train', root=path_dataset, transform=transform)
    sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=True)
    data_loader = torch.utils.data.DataLoader(dataset, batch_size=4, sampler=sampler, shuffle=False, num_workers=8)

    # Train part of the first epoch
    model.train()

    for idx, (images, labels) in enumerate(data_loader):
        if idx >= 10:
            break
        images = images.to(device)
        labels = labels.to(device)

        with torch.set_grad_enabled(True):
            features = model(images)
            loss = torch.nn.functional.cross_entropy(features, labels)

        loss.backward()
        print(f'grad iteration {idx} on gpu {device}: {model.module.conv1.weight.grad.mean()}', flush=True)
        optimizer.step()
        optimizer.zero_grad()
        print(f'weight iteration {idx} on gpu {device}: {model.module.conv1.weight.mean()}', flush=True)

    # cleanup
    torch.distributed.destroy_process_group()


if __name__ == '__main__':
    world_size_main = torch.cuda.device_count()
    mp.spawn(main,
             args=(world_size_main,),
             nprocs=world_size_main,
             join=True)

In order to run, use: CUDA_VISIBLE_DEVICES=0,1 python run.py

Expected behavior

The expected behavior is the one that occurs when the line optimizer = torch.optim.SGD(model.parameters(), lr=0.1) is uncommented, and the line optimizer = geoopt.optim.RiemannianSGD(model.parameters(), lr=0.1, , stabilize=10) is commented. In that case, the output is:

grad iteration 0 on gpu cuda:1: 0.011769304051995277
grad iteration 0 on gpu cuda:0: 0.011769304051995277
weight iteration 0 on gpu cuda:1: -0.001249525579623878
weight iteration 0 on gpu cuda:0: -0.001249525579623878
grad iteration 1 on gpu cuda:0: -0.015764284878969193
grad iteration 1 on gpu cuda:1: -0.015764284878969193
weight iteration 1 on gpu cuda:1: 0.0003269027511123568
weight iteration 1 on gpu cuda:0: 0.0003269027511123568
grad iteration 2 on gpu cuda:1: -0.006310341879725456
grad iteration 2 on gpu cuda:0: -0.006310341879725456
weight iteration 2 on gpu cuda:1: 0.000957937038037926
weight iteration 2 on gpu cuda:0: 0.000957937038037926
grad iteration 3 on gpu cuda:0: 0.0021547293290495872
grad iteration 3 on gpu cuda:1: 0.0021547293290495872
weight iteration 3 on gpu cuda:1: 0.000742464151699096
weight iteration 3 on gpu cuda:0: 0.000742464151699096
grad iteration 4 on gpu cuda:1: -0.002606849418953061
grad iteration 4 on gpu cuda:0: -0.002606849418953061
weight iteration 4 on gpu cuda:1: 0.001003148965537548
weight iteration 4 on gpu cuda:0: 0.001003148965537548
grad iteration 5 on gpu cuda:1: 0.00043087091762572527
grad iteration 5 on gpu cuda:0: 0.00043087091762572527
weight iteration 5 on gpu cuda:1: 0.0009600619086995721
weight iteration 5 on gpu cuda:0: 0.0009600619086995721
grad iteration 6 on gpu cuda:0: 0.00014396056940313429
grad iteration 6 on gpu cuda:1: 0.00014396056940313429
weight iteration 6 on gpu cuda:1: 0.0009456658735871315
weight iteration 6 on gpu cuda:0: 0.0009456658735871315
grad iteration 7 on gpu cuda:1: -0.002603260101750493
grad iteration 7 on gpu cuda:0: -0.002603260101750493
weight iteration 7 on gpu cuda:1: 0.001205991953611374
weight iteration 7 on gpu cuda:0: 0.001205991953611374
grad iteration 8 on gpu cuda:0: 0.000458348571555689
grad iteration 8 on gpu cuda:1: 0.000458348571555689
weight iteration 8 on gpu cuda:0: 0.0011601571459323168
weight iteration 8 on gpu cuda:1: 0.0011601571459323168
grad iteration 9 on gpu cuda:1: -0.0004215179360471666
grad iteration 9 on gpu cuda:0: -0.0004215179360471666
weight iteration 9 on gpu cuda:1: 0.0012023089220747352
weight iteration 9 on gpu cuda:0: 0.0012023089220747352

The gradients in the two GPUs are correctly synchronized. However, when using RiemannianSGD, the output is:

grad iteration 0 on gpu cuda:1: 0.0035285688936710358
grad iteration 0 on gpu cuda:0: 0.0035285688936710358
weight iteration 0 on gpu cuda:0: -7.928906597953755e-06
weight iteration 0 on gpu cuda:1: -7.928906597953755e-06
grad iteration 1 on gpu cuda:0: -0.04444637894630432
grad iteration 1 on gpu cuda:1: 0.002550020581111312
weight iteration 1 on gpu cuda:0: 0.00018905977776739746
weight iteration 1 on gpu cuda:1: -2.1470928913913667e-05
grad iteration 2 on gpu cuda:0: 0.009863540530204773
grad iteration 2 on gpu cuda:1: 0.0026304360944777727
weight iteration 2 on gpu cuda:0: 0.00013691956701222807
weight iteration 2 on gpu cuda:1: -3.7374120438471437e-05
grad iteration 3 on gpu cuda:0: -0.0161017756909132
grad iteration 3 on gpu cuda:1: -0.0023103044368326664
weight iteration 3 on gpu cuda:0: 0.0002405263076070696
weight iteration 3 on gpu cuda:1: -2.4383840354857966e-05
grad iteration 4 on gpu cuda:0: 0.010763526894152164
grad iteration 4 on gpu cuda:1: -0.017034146934747696
weight iteration 4 on gpu cuda:0: 0.00017218326684087515
weight iteration 4 on gpu cuda:1: 7.665574958082289e-05
grad iteration 5 on gpu cuda:0: 0.008465449325740337
weight iteration 5 on gpu cuda:0: 0.00012061965389875695
grad iteration 5 on gpu cuda:1: 0.0011690922547131777
weight iteration 5 on gpu cuda:1: 6.942617619642988e-05
grad iteration 6 on gpu cuda:0: 0.0013559082290157676
weight iteration 6 on gpu cuda:0: 0.00011242596519878134
grad iteration 6 on gpu cuda:1: 0.0008932517375797033
weight iteration 6 on gpu cuda:1: 6.43157254671678e-05
grad iteration 7 on gpu cuda:0: 0.02651313878595829
weight iteration 7 on gpu cuda:0: -1.233588955074083e-05
grad iteration 7 on gpu cuda:1: -0.007853103801608086
weight iteration 7 on gpu cuda:1: 0.00010782096069306135
grad iteration 8 on gpu cuda:0: 0.009321866557002068
weight iteration 8 on gpu cuda:0: -7.130965968826786e-05
grad iteration 8 on gpu cuda:1: -0.0039948648773133755
grad iteration 9 on gpu cuda:0: -0.0119229881092906
weight iteration 8 on gpu cuda:1: 0.00013168319128453732
weight iteration 9 on gpu cuda:0: -1.202533780997328e-06
grad iteration 9 on gpu cuda:1: 0.002446404891088605
weight iteration 9 on gpu cuda:1: 0.00011651107342913747

There is some problem with the gradient synchronization, which causes the weights in the two devices to diverge.

Library version information:

• python -c 'import torch;print("torch:", torch.version.__version__, end=" ");print("cuda:", torch.version.cuda)' torch: 1.8.1 cuda: 11.1

• the way you installed geoopt, github, pip pip

• OS Ubuntu 18.04.5 LTS

EDIT: I simplified a little bit the code by removing mixed precision.

Hi,

I'm building neural network in pytorch which has to learn to make regression from vectors in Euclidean Space to some vectors in Poincarè Disk Model, so I think that the usage of RiemannianSGD can be a good choice as the optimizer.

I'm trying to use the library but I have some question:

• When and how I have to cast tensors? Now I transform the target tensors with ManifoldTensor with Y = ManifoldTensor(Y, manifold = geoopt.manifolds.PoincareBall(), I think that even the prediction has to be a Manifold Tensor, but if I transform the prediction outside the model the ManifoldTensor does not maintain the grad value

• Have I to build each layer of the network with ManifoldTensor or torch.nn.Linear can be used?

• Have I to specify to the optimizer RiemannianSGD the manifold? From the documentation I think that I don't have to do it

Anyway thanks for your work on this project ; - )

Hi,

I recently updated the pytorch environment library and so I installed again geoopt via : pip install git+https://github.com/geoopt/geoopt.git

When I try to import the library this is the error that is reported to me:

  import geoopt
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/geoopt/__init__.py", line 1, in <module>
    from . import manifolds
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/geoopt/manifolds/__init__.py", line 3, in <module>
    from .stiefel import Stiefel, EuclideanStiefel, CanonicalStiefel, EuclideanStiefelExact
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/geoopt/manifolds/stiefel.py", line 4, in <module>
    from .. import linalg
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/geoopt/linalg/__init__.py", line 1, in <module>
    from .batch_linalg import svd, qr, sym, extract_diag, matrix_rank, expm, block_matrix
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/geoopt/linalg/batch_linalg.py", line 3, in <module>
    from . import _expm
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/geoopt/linalg/_expm.py", line 8, in <module>
    @torch.jit.script
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/torch/jit/__init__.py", line 364, in script
    graph = _script_graph(fn, _frames_up=_frames_up + 1)
  File "/home/vmanuel/.conda/envs/MTNCI/lib/python3.6/site-packages/torch/jit/__init__.py", line 360, in _script_graph
    return _jit_script_compile(ast, rcb)
RuntimeError: 
builtin cannot be used as a value:
        33522128640.0,
        1323241920.0,
        40840800.0,
        960960.0,
        16380.0,
        182.0,
        1.0,
    )

    ident = torch.eye(A.shape[1], dtype=A.dtype, device=A.device)
                      ~~~~~~~ <--- HERE
    A2 = torch.matmul(A, A)
    A4 = torch.matmul(A2, A2)
    A6 = torch.matmul(A4, A2)
    U = torch.matmul(
        A,
        torch.matmul(A6, b13 * A6 + b11 * A4 + b9 * A2)
        + b7 * A6
        + b5 * A4
        + b3 * A2

Can you help me?

Describe the bug/To Reproduce The following warning appears when I use Distance2PoincareHyperplanes as the classifier for torchvision.models.efficientnet_v2_s.

~/miniconda3/envs/a5000/lib/python3.10/site-packages/geoopt/manifolds/stereographic/math.py:1562: UserWarning: operator() profile_node %301 : int = prim::profile_ivalue(%299)
 does not have profile information (Triggered internally at  /opt/conda/conda-bld/pytorch_1656352645774/work/torch/csrc/jit/codegen/cuda/graph_fuser.cpp:104.)
  return _dist2plane(
~/miniconda3/envs/a5000/lib/python3.10/site-packages/geoopt/manifolds/stereographic/math.py:1562: UserWarning: FALLBACK path has been taken inside: compileCudaFusionGroup. This is an indication that codegen Failed for some reason.
To debug try disable codegen fallback path via setting the env variable `export PYTORCH_NVFUSER_DISABLE=fallback`
To report the issue, try enable logging via setting the envvariable ` export PYTORCH_JIT_LOG_LEVEL=manager.cpp`
 (Triggered internally at  /opt/conda/conda-bld/pytorch_1656352645774/work/torch/csrc/jit/codegen/cuda/manager.cpp:237.)
  return _dist2plane(

Expected behavior Warning appears during model training.

Please complete the following information:

• torch: 1.12.0, cuda: 11.3, geoopt 0.4.1
• The way you installed geoopt: pip
• OS: Ubuntu 20.04.4 LTS (GNU/Linux 5.4.0-109-generic x86_64)

Additional context NIL

Describe the bug With certain matrices, the use of the projx under the Lorentz Manifold fails with_check_point_on_manifold

To Reproduce Steps to reproduce the behavior:

• I provided an example problem tensor for this case
• Call projx on the tensor with a k of 5.0
• Check if the result is on the manifold

Expected behavior With float 64 and a defined error value, check on manifold should result in True

Screenshots

• It may be a training/testing loss, other information for more context

Please complete the following information:

• torch: 1.11.0+cu113 cuda: 11.3
• Installed through Github
• Windows 11 trouble_tensor_projection.zip

Thanks for your useful library. However, I found that Lorentz model has not implemented any mobius methods such as mobius_add() or mobius_matvec(). I hope you will update these methods to this library soon.

following #142

• [ ] mobius_add

• [ ] mobius_matvec

• [ ] mobius_scalar_mul

• [ ] mobius_pointwise_mul

• [ ] mobius_fn_apply

• [ ] test new functions

• [ ] re-review Rasul's implementation

Hi,

Thanks for the useful tool you've developed, and I really appreciate it. When using geoopt, I found that the mobius_add() and mobius_matvec() methods for Lorentz manifold are missing, and there's an implementation at https://www.github.com/HazyResearch/hgcn. Could those methods be added into the package soon? Thanks a lot!