https://github.com/helmholtz-ai-energy/oialr
https://github.com/helmholtz-ai-energy/oialr
Last synced: about 1 month ago
JSON representation
- Host: GitHub
- URL: https://github.com/helmholtz-ai-energy/oialr
- Owner: Helmholtz-AI-Energy
- License: bsd-3-clause
- Created: 2024-02-14T10:18:47.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2024-05-19T19:44:33.000Z (over 1 year ago)
- Last Synced: 2025-05-30T12:12:14.033Z (5 months ago)
- Language: Python
- Size: 1.11 MB
- Stars: 0
- Watchers: 3
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Orthogonality-Informed Adaptive Low-Rank Training
# Description
The orthogonal bases inherent to the the weights of a neural networks stabalizes during training.
This framework aims to make use of this to fact to train low-rank neural networks.
# Orthogonal Stabilization
For more details see our paper [here](https://arxiv.org/abs/2401.08505)
The short version, is that a network's weights learn to the basis first, while learning the mixing throughout training.
This is like learning a language, where one learns the basics, before learning how to use them.
To show this, we tracked the basis `U @ V.T` with a metric we call *Stability*, and we tracked the linear mixing via the `R` matrix from QR decomposition. For the `R` tracking we used the enclidean distance.
### *Stability* of the basis:
ImageNet + ResNet-RS 101 | ImageNet + Vision Transformer B/16
:-------------------------:|:-------------------------:
 | 
### Enclidean Similarity of the `R` mixing matrix:
ImageNet + ResNet-RS 101 | ImageNet + Vision Transformer B/16
:-------------------------:|:-------------------------:
 | 
We use this to train low rank networks which stop training the `U` and `V` matrices during training. To do this we wrap the network in a class and replace the layers.
This is found in the `src/madonna/models/svd/model.py` in the `OIALRModel` class.
# Getting started
## Create the pipeline environment and install the madonna package
We recommend using an NVIDIA container with torch preinstalled.
Before using this make sure to install the package with
```
pip install -e .
```
This package use Hydra for config/experiment management, although most things are managed though OmegaConf when needed.
If this is used, it will load the sub-configs for the run.
## Starting example
The best place to start is with a minified ViT model training on CIFAR10.
The config for this experiment is in `configs/experiment/cifar10-minivit.yaml`.
We trained this using a slurm launcher and `enroot`, you can look at how to do this with either the `scripts/launch_run.sbatch`.
You can also use the `singularity` for launching the runs with `scripts/launch_singularity.sbatch`.
These scripts all build off the `configs/ortho_train.yaml` config. This is where the configs for all the options are inhereted from. However, they are all over-ridden by the `experiment` config, if it is secified.
# Project Organization
```
├── configs <- Hydra configuration files
│ ├── callbacks <- Callbacks configs
│ ├── data <- Datamodule configs
│ ├── debug <- Debugging configs
│ ├── experiment <- Experiment configs
│ ├── local <- Local configs
│ ├── log_dir <- Logging directory configs
│ ├── logger <- Logger configs
│ ├── model <- Model configs
│ ├── trainer <- Trainer configs
│ │
│ ├── test.yaml <- Main config for testing
│ └── train.yaml <- Main config for training
│
├── docs <- Directory for Sphinx documentation in rst or md.
├── models <- Trained and serialized models, model predictions
├── notebooks <- Jupyter notebooks.
├── reports <- Generated analysis as HTML, PDF, LaTeX, etc.
│ └── figures <- Generated plots and figures for reports.
├── scripts <- Scripts used in project
│ ├── job_submission.sbatch <- Submit training job to slurm
│ ├── job_submission_interactive.sbatch <- Submit training job to slurm (interactive node)
│ ├── test.py <- Run testing
│ └── train.py <- Run training
│
├── src/madonna <- Source code
│ ├── datamodules <- Lightning datamodules
│ ├── models <- Lightning models
│ └── utils <- Utility scripts
│
├── .gitignore <- List of files/folders ignored by git
├── .pre-commit-config.yaml <- Configuration of pre-commit hooks for code formatting
├── requirements.txt <- File for installing python dependencies
├── setup.cfg <- Configuration of linters and pytest
├── LICENSE.txt <- License as chosen on the command-line.
├── pyproject.toml <- Build configuration. Don't change! Use `pip install -e .`
│ to install for development or to build `tox -e build`.
├── setup.cfg <- Declarative configuration of your project.
├── setup.py <- [DEPRECATED] Use `python setup.py develop` to install for
│ development or `python setup.py bdist_wheel` to build.
└── README.md
```