Spectral Inference Networks (SpIN)

This package provides an implementation of Spectral Inference Networks, as in Pfau, Petersen, Agarwal, Barrett and Stachenfeld (2019).

This is not an officially supported Google product.

Prerequisites

SpIN requires a working installation of Python and TensorFlow. We recommend running it on GPU for faster convergence.

If you want to make use of the GUI (on by default) you will also need Tcl/Tk installed on your system.

Installation

After cloning the repo, run pip to install the package and its Python dependencies:

bash cd spectral_inference_networks pip install .

Usage

Training a spectral inference network is similar to most other deep learning pipelines: you must construct a data source, network architecture and optimizer. What makes spectral inference networks unique is that instead of a loss you provide a linear operator to diagonalize. The code expects an object of the LinearOperator class, which can be constructed from a similarity kernel or by other means. LinearOperator objects can be added together or multiplied by a scalar.

Below is a minimal example of training spectral inference networks:

```python import tensorflow as tf import spectralinferencenetworks as spin

batchsize = 1024 inputdim = 10 num_eigenvalues = 5 iterations = 1000 # number of training iterations

Create variables for simple MLP

w1 = tf.Variable(tf.random.normal([inputdim, 64])) w2 = tf.Variable(tf.random.normal([64, numeigenvalues]))

b1 = tf.Variable(tf.random.normal([64])) b2 = tf.Variable(tf.random.normal([num_eigenvalues]))

Create function to construct simple MLP

def network(x): h1 = tf.nn.relu(tf.matmul(x, w1) + b1) return tf.matmul(h1, w2) + b2

data = tf.random.normal([batchsize, inputdim]) # replace with actual data

Squared exponential kernel.

kernel = lambda x, y: tf.exp(-(tf.norm(x-y, axis=1, keepdims=True)**2)) linop = spin.KernelOperator(kernel) optim = tf.train.AdamOptimizer()

Constructs the internal training ops for spectral inference networks.

spectral_net = spin.SpectralNetwork( linop, network, data, [w1, w2, b1, b2])

Trivial defaults for logging and stats hooks.

loggingconfig = { 'config': {}, 'logimageevery': iterations, 'saveparamsevery': iterations, 'saverpath': '/tmp', 'saver_name': 'example', }

statshooks = { 'create': spin.util.createdefaultstats, 'update': spin.util.updatedefault_stats, }

Executes the training of spectral inference networks.

stats = spectralnet.train( optim, iterations, loggingconfig, stats_hooks) ```

We provide two examples in the examples folder, which you can run as follows:

bash python spectral_inference_networks/examples/hydrogen.py

and:

bash python spectral_inference_networks/examples/atari.py These correspond to experiments in section 5.1 and C.3 of the paper. Each example comes with a range of supported command line arguments. Please take a look in the source code for each example for further information and have a play with the many options.

Giving Credit

If you use this code in your work, we ask that you cite the paper:

David Pfau, Stig Petersen, Ashish Agarwal, David Barrett, Kim Stachenfeld. "Spectral Inference Networks: Unifying Deep and Spectral Learning." The 7th International Conference on Learning Representations (ICLR) (2019).

Acknowledgements

Special thanks to James Spencer for help with the open-source implementation of the code.