<div align="center">
  <img src="docs/src/lfp_logo.png" width="400"/>
  <p>Gradient-free Neural Network Training based on Layer-wise Relevance Propagation (LRP)</p>
</div>

## :rocket: Getting Started


### <a name="installation"></a> :fire: Installation

#### <a name="poetryinstall"></a> Using [Poetry](https://python-poetry.org/)

You can install the necessary dependencies to run the code and reproduce experiments via Poetry:

```shell
cd layerwise-feedback-propagation
poetry install
```

This requires ```poetry-core>=2.0.0,<3.0.0```.

### Overview

Our implementation of LFP is based the LRP-implementation of the [zennit](https://github.com/chr5tphr/zennit) and [LXT](https://github.com/rachtibat/LRP-eXplains-Transformers) libraries. Both of these libraries are based on [PyTorch](https://pytorch.org/) and modify the backward pass to return relevances instead of gradients.

```lfprop``` extends these libraries to return relevances not only w.r.t. activations, but also w.r.t. parameters. Similar to LXT and zennit, this requires registering a composite to modify the backward pass.

*LXT Backend*
```python
from lfprop.propagation import propagator_lxt
propagation_composite = propagator.LFPEpsilonComposite()
```
*Zennit Backend*
```python
from lfprop.propagation import propagator_zennit
propagation_composite = propagator.LFPEpsilonComposite()
```

*SNNs (Also Zennit Backend)*
```python
from lfprop.propagation import propagator_snn
propagation_composite = propagator.LRPRewardPropagator()
```

Instead of an initial relevance, LFP requires an initial reward at the output, to be decomposed throughout the model. We implement several reward functions, with a similar signature to ```torch.nn.Loss``` functions.

```python
from lfprop.rewards import reward_functions as rewards
reward_func = rewards.SoftmaxLossReward(device)
```

To apply the modified backward pass, the composite simply needs to be registered.

After the backward pass is finished, the computed LFP-feedback can then be accessed via the (newly added) ```.feedback``` attribute of each parameter.

The model can simply be optimized using any ```torch.nn.Optimizer```, by first overwriting the ```.grad``` attribute by the corresponding (negative) feedback.

This results in the following training step:

```python
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
optimizer.zero_grad()

with propagation_composite.context(model) as modified:
    inputs = inputs.detach().requires_grad_(True)
    outputs = modified(inputs)

    # Calculate reward
    reward = torch.from_numpy(
        reward_func(outputs, labels).detach().cpu().numpy()
    ).to(device)

    # Calculate LFP
    input_reward = torch.autograd.grad(
        (outputs,), (inputs,), grad_outputs=(reward,), retain_graph=False
    )[0]

    # Write LFP Values into .grad attributes.
    for name, param in model.named_parameters():
        param.grad = -param.feedback

    # Optimization step
    optimizer.step()
```

### Examples

A simple, full example of how to train a LeNet model on MNIST can be found under ```minimal_example.ipynb```. An example using SNNs can be found under ```minimal_example_spiking_nets.ipynb```

### :mag: Reproducing Experiments

Most toy data experiments can then be reproduced by simply running the corresponding notebooks under ```nbs/```. You can find the used hyperparameters for the notebooks within the first two cells.

For reproducing the experiments that require training on more complex data and models (LFP for Non-ReLU and Pruning experiments), the training script is implemented in ```run_experiment.py```.
Hyperparameters for these experiments can be generated using the scripts under ```configs/```.

For reproducing the Non-SNN experiments first run
```bash
# 1. generate the config files
python configs/<experimentname>/config_generator<somesuffix>.py
# 2. run training script
python run_experiment.py --config_file=configs/<experimentname>/cluster/<selected-config-name>
```

For the pruning experiments, you can then run the ```nbs/*eval-clusterresults-pruning.ipynb``` notebooks using the obtained models.

For reproducing the SNN experiments run
```bash
# 1. generate the config files
python configs/spiking_neural_networks/config_generator_mnist_training.py
# 2. run training script
python run_snn_experiment.py --config_file=configs/spiking_neural_networks/cluster/<selected-config-name>
```
