Butterfly

A Sinkhorn network is a supervised method used to learn to reconstruct a scrambled object $\tilde{X}$ (input) given training pairs of $(X, \tilde{X})$.

By applying non-linear transformations, the network parameterizes a mapping between the scrambled input $\tilde{X}$ and a permutation $P$. When $P^{-1}$ is applied to $\tilde{X}$, it reconstructs the original object: $X_{rec} = P^{-1}\tilde{X}$. This implementation is based on the paper Learning Latent Permutations with Gumbel-Sinkhorn Networks.

To evaluate a pre-trained Sinkhorn model, use the my_sinkhorn_eval.py script. The script uses the trained_model by default.

You can adjust the evaluation granularity by modifying the samples_per_num variable. The original paper tested values of 5, 10, 15, 80, 100, 120.

The evaluation calculates the following metrics:

1. L1 Loss: Mean absolute difference between sorted input and reconstructed input using hard learned permutations.
2. L2 Loss: Mean squared difference between sorted input and reconstructed input using hard learned permutations.
3. Prop. wrong: Proportion of errors in sorting.
4. Prop. any wrong: Proportion of sequences containing at least one error.
5. Kendall's tau: A measure of similarity between the correct ordering and the ordering learned by the network.

python my_sinkhorn_eval.py

To perform single node training on CIFAR, execute the cifar_experiment.py script. You can specify the model, optimizer, and decay settings via command-line arguments.

Available arguments in the example include:

• model: The architecture to use (e.g., LeNet).
• optimizer: The optimization algorithm (e.g., SGD).
• lr_decay: Boolean to enable learning rate decay.
• weight_decay: Boolean to enable weight decay.

python cifar_experiment.py with model=LeNet optimizer=SGD lr_decay=True weight_decay=True

To retrain the Sinkhorn model, run the my_sorting_train.py script.

The training process uses the mean squared reconstruction error as the loss function. It operates by applying soft permutations to random input numbers drawn i.i.d. from a standard uniform distribution and comparing the result to the correct ordering.

python my_sorting_train.py

The legacy interface is being rewritten and should only be used if required for specific features not yet available in the new interface.

To use the Butterfly module in butterfly/butterfly.py as a drop-in replacement for torch.nn.Linear, you must install the C++/CUDA extensions located in the subdirectories.

Note: If you do not install these extensions, butterfly multiplication will still work using a pure PyTorch implementation, but it will be significantly slower. The behavior is controlled by the use_extension variable in butterfly/butterfly_multiply.py.

# Installing legacy C++/CUDA extensions
cd butterfly/factor_multiply
python setup.py install
cd butterfly/factor_multiply_fast
python setup.py install

To run speed comparisons and benchmarks, you must compile the Cython extensions. This is only necessary for the speed benchmark functionality.

python setup.py build_ext --inplace

ImageNet experiments use a modified version of the imagenet-fast code. Run the experiments using the multiproc module with the imagenet_experiment.py script.

Required arguments:

• --lr: Learning rate.
• --epochs: Number of training epochs.
• --small: A flag (likely for a smaller subset or specific configuration).
• [DIR FOR IMAGENET DATA]: The positional argument specifying the directory containing the ImageNet dataset.

python -m multiproc imagenet_experiment.py --lr 0.4 --epochs 45 --small [DIR FOR IMAGENET DATA]

To use the high-performance C++/CUDA implementation of butterfly multiplication, run the setup.py file located in the root directory.

It is recommended to use a Conda environment with compatible versions of Python, SciPy, and PyTorch (including cudatoolkit).

# Example setup using conda
conda create --name butterfly python=3.8 scipy pytorch=1.8.1 cudatoolkit=11.0 -c pytorch
conda activate butterfly
python setup.py install

The modern interface for the butterfly C++/CUDA code is located in the csrc and torch_butterfly modules.

For specific structured linear operations, refer to torch_butterfly/special.py. This module provides methods to construct butterfly matrices that perform:

• FFT (Fast Fourier Transform)
• Inverse FFT
• Circulant matrix multiplication
• Hadamard transform
• torch.nn.Conv1d with circular padding

Example usage and verification can be found in tests/test_butterfly.py and tests/test_special.py.