Package: geotorch
A library for constrained optimization and manifold optimization for deep learning in PyTorchOverviewGeoTorch provides a simple way to perform constrained optimization and optimization on manifolds in PyTorch.
It is compatible out of the box with any optimizer, layer, and model implemented in PyTorch without any boilerplate in the training code. Just state the constraints when you construct the model and you are ready to go!importtorchimporttorch.nnasnnimportgeotorchclassModel(nn.Module):def__init__(self):super().__init__()# One line suffices: Make a linear layer with orthonormal columnsself.linear=nn.Linear(64,128)geotorch.orthogonal(self.linear,"weight")# Works with tensors: Make a CNN with kernels of rank 1self.cnn=nn.Conv2d(16,32,3)geotorch.low_rank(self.cnn,"weight",rank=1)# Weights are initialized to a random value when you put the constraints, but# you may re-initialize them to a different value by assigning to themself.linear.weight=torch.eye(128,64)# And that's all you need to do. The rest is regular PyTorch codedefforward(self,x):# self.linear is orthogonal and every 3x3 kernel in self.cnn is of rank 1# Use the model as you would normally do. Everything just worksmodel=Model().cuda()# Use your optimizer of choice. Any optimizer works out of the box with any parametrizationoptim=torch.optim.Adam(model.parameters(),lr=lr)ConstraintsThe following constraints are implemented and may be used as in the example above:geotorch.symmetric. Symmetric matricesgeotorch.skew. Skew-symmetric matricesgeotorch.sphere. Vectors of norm1geotorch.orthogonal. Matrices with orthogonal columnsgeotorch.grassmannian. Skew-symmetric matricesgeotorch.almost_orthogonal(λ). Matrices with singular values in  the interval[1-λ,1+λ]geotorch.invertible. Invertible matrices with positive determinantgeotorch.low_rank(r). Matrices of rank at mostrgeotorch.fixed_rank(r). Matrices of rankrgeotorch.positive_definite. Positive definite matricesgeotorch.positive_semidefinite. Positive semidefinite matricesgeotorch.positive_semidefinite_low_rank(r). Positive semidefinite matrices of rank at mostrgeotorch.positive_semidefinite_fixed_rank(r). Positive semidefinite matrices of rankrEach of these constraints have some extra parameters which can be used to tailor the
behavior of each constraint to the problem in hand. For more on this, see the documentation.These constraints are a fronted for the families of spaces listed below.Supported SpacesEach constraint in GeoTorch is implemented as a manifold. These give the user more flexibility
on the options that they choose for each parametrization. All these support Riemannian Gradient
Descent by default (more on thishere), but they also support optimization via any other PyTorch
optimizer.GeoTorch currently supports the following spaces:Rn(n):Rⁿ. Unrestricted optimizationSym(n): Vector space of symmetric matricesSkew(n): Vector space of skew-symmetric matricesSphere(n): Sphere inRⁿ.{ x ∈ Rⁿ | ||x|| = 1 } ⊂ RⁿSO(n): Manifold ofn×northogonal matricesSt(n,k): Manifold ofn×kmatrices with orthonormal columnsAlmostOrthogonal(n,k,λ): Manifold ofn×kmatrices with singular values in the interval[1-λ,1+λ]Gr(n,k): Manifold ofk-dimensional subspaces inRⁿGLp(n): Manifold of invertiblen×nmatrices with positive determinantLowRank(n,k,r): Variety ofn×kmatrices of rankror lessFixedRank(n,k,r): Manifold ofn×kmatrices of rankrPSD(n): Cone ofn×nsymmetric positive definite matricesPSSD(n): Cone ofn×nsymmetric positive semi-definite matricesPSSDLowRank(n,r): Variety ofn×nsymmetric positive semi-definite matrices of rankror lessPSSDFixedRank(n,r): Manifold ofn×nsymmetric positive semi-definite matrices of rankrProductManifold(M₁,...,Mₖ): Product of manifoldsM₁ × ... × MₖEvery space of dimension(n, k)can be applied to tensors of shape(*, n, k), so we also get efficient parallel implementations of product spaces such asObliqueManifold(n,k): Matrix with unit length columns,Sⁿ⁻¹ ×...ᵏ⁾× Sⁿ⁻¹Using GeoTorch in your CodeThe files inexamples/copying_problem.pyandexamples/sequential_mnist.pyserve as tutorials to see how to handle the initialization and usage of GeoTorch in some real code. They also show how to implement Riemannian Gradient Descent and some other tricks. For an introduction to how the library is actually implemented, see the Jupyter Notebookexamples/parametrisations.ipynb.You may try GeoTorch installing it withpipinstallgit+https://github.com/Lezcano/geotorch/GeoTorch is tested in Linux, Mac, and Windows environments for Python >= 3.6 and supports PyTorch >= 1.9Sharing Weights, Parametrizations, and Normalizing FlowsIf one wants to use a parametrized tensor in different places in their model, or uses one parametrized layer many times, for example in an RNN, it is recommended to wrap the forward pass as follows to avoid each parametrization to be computed many times:withgeotorch.parametrize.cached():logits=model(input_)Of course, thiswithstatement may be used simply inside the forward function where the parametrized layer is used several times.These ideas fall in the context of parametrized optimization, where one wraps a tensorXwith a functionf, and rather than usingX, usesf(X). Particular examples of this idea are pruning, weight normalization, and spectral normalization among others. This repository implements a framework to approach this kind of problems. This framework was accepted to core PyTorch 1.8. It can be found undertorch.nn.utils.parametrizeandtorch.nn.utils.parametrizations. When using PyTorch 1.10 or higher, these functions are used, and the user can interact with the parametrizations in GeoTorch using the functions in PyTorch.As every space in GeoTorch is, at its core, a map from a flat space into a manifold, the tools implemented here also serve as a building block in normalizing flows. Using a factorized space such asLowRank(n,k,r)it is direct to compute the determinant of the transformation it defines, as we have direct access to the singular values of the layer.BibliographyPlease cite the following work if you found GeoTorch useful. This paper exposes a simplified mathematical explanation of part of the inner-workings of GeoTorch.@inproceedings{lezcano2019trivializations,title={Trivializations for gradient-based optimization on manifolds},author={Lezcano-Casado, Mario},booktitle={Advances in Neural Information Processing Systems, NeurIPS},pages={9154--9164},year={2019},}