RMLR: Extending Multinomial Logistic Regression into General Geometries

This is the official code for our NeurIPS 2024 publication: RMLR: Extending Multinomial Logistic Regression into General Geometries.

If you find this project helpful, please consider citing us as follows:

@inproceedings{chen2024rmlr,
    title={{RMLR}: Extending Multinomial Logistic Regression into General Geometries},
    author={Ziheng Chen and Yue Song and Rui Wang and Xiaojun Wu and Nicu Sebe},
    booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
    year={2024},
}

And our previous work on the flat SPD MLR, ALEM, and product Cholesky metrics (PCM and BWCM):

@inproceedings{chen2024spdmlr,
    title={Riemannian Multinomial Logistics Regression for {SPD} Neural Networks},
    author={Ziheng Chen and Yue Song and Gaowen Liu and Ramana Rao Kompella and Xiaojun Wu and Nicu Sebe},
    booktitle={Conference on Computer Vision and Pattern Recognition 2024},
    year={2024}
}

@article{chen2024adaptive,
  title={Adaptive Log-Euclidean metrics for {SPD} matrix learning},
  author={Chen, Ziheng and Song, Yue and Xu, Tianyang and Huang, Zhiwu and Wu, Xiao-Jun and Sebe, Nicu},
  journal={IEEE Transactions on Image Processing},
  year={2024}
}

@article{chen2024product,
  title={Product geometries on {Cholesky} manifolds with applications to {SPD} manifolds},
  author={Chen, Ziheng and Song, Yue and Wu, Xiao-Jun and Sebe, Nicu},
  journal={arXiv preprint arXiv:2407.02607},
  year={2024}
}

If you have any problem, do not hesitate to contact me ziheng_ch@163.com.

Geometries

This source code contains eight SPD MLRs:

"LEM": SPDLogEuclideanMetric,
"ALEM": SPDAdaptiveLogEuclideanMetric,
"LCM": SPDLogCholeskyMetric,
"AIM": SPDAffineInvariantMetric,
"BWM": SPDBuresWassersteinMetric,
"PEM": SPDEuclideanMetric,
"PCM": SPDPowerCholeskyMetric,
"BWCM": SPDBuresWassersteinCholeskyMetric,

where ALEM comes from our TIP24, and PCM & BWCM come from our Arxiv24.

We also release the torch implementation for SO(3) in Geometry.rotation.

Requirements

Install necessary dependencies by conda:

conda env create --file environment.yaml

Note that the hydra package is used to manage configuration files.

Demos

Demos of typical use can be found in demo_spd.py:

import torch as th
from RieNets.spdnets.SPDMLR import SPDRMLR

# Set parameters
batch_size = 8  # Batch size
n = 5  # Dimension of SPD matrices
c = 3  # Number of classes

# Generate random SPD matrices of shape (batch_size, 1, n, n)
X = th.randn(batch_size, 1, n, n)
X = X @ X.transpose(-1, -2)  # Ensure positive definiteness
X += th.eye(n) * 1e-3  # Add a small perturbation to guarantee strict positive definiteness

# Initialize the model
model = SPDRMLR(n=n, c=c, metric='LEM')

# Forward computation
output = model(X)

Experiments

The implementation of SPD MLR is based on our previous work:

Riemannian Multinomial Logistics Regression for SPD Neural Networks [CVPR 2024] [code].

Dataset

The preprocessed SPD from the HDM05 dataset can be found [here].

Please download the datasets and put them in your folder. If necessary, change the path in conf/RiemNets/dataset/HDM05_SPD.yaml,

Note: Other datasets for SPD networks can be found in our CVPR24 paper.

Running experiments

To run experiments on the SPDNet (Tabs. 4.), run this command:

bash exp_spdnets.sh

Note: You also can change the hdm05_path in exp_spdnets.sh, which will override the hydra config.

Visualization

The pyrhon code for visualization can be found in spd_vis.py and rot_vis.py, corresponding to Figs. 2-3.

GitWR / RMLR