Neurophox

Neurophox is an open source machine learning and photonic simulation framework based on unitary mesh networks presented in arxiv/1808.00458 and arxiv/1903.04579.

Note: neurophox is under active development and is not yet stable, so expect to see the documentation change frequently.

neurophox neurophox

Motivation

Orthogonal and unitary neural networks have interesting properties and have been studied for synthetic natural language processing tasks (see unitary mesh-based RNN, unitary evolution RNN, and orthogonal evolution RNN). Furthermore, new energy-efficient photonic technologies are being built to realize unitary mesh-based neural networks using light as the computing medium as opposed to conventional analog electronics.

Introduction

Neurophox provides a general framework for mesh network layers in orthogonal and unitary neural networks. We use an efficient definition for any feedforward mesh architecture, neurophox.meshmodel.MeshModel, to develop mesh layer architectures in Numpy (neurophox.numpy.layers), Tensorflow 2.0 (neurophox.tensorflow.layers), and (soon) PyTorch.

Scattering matrix models used in unitary mesh networks for photonics simulations are provided in neurophox.components. The models for all layers are fully defined in neurophox.meshmodel, which provides a general framework for efficient implementation of any unitary mesh network.

Dependencies and requirements

Some important requirements for neurophox are:

  1. Python >=3.6

  2. Tensorflow >=2.0

  3. PyTorch >=1.3

The dependencies for neurophox (specified in requirements.txt) are:

numpy>=1.16
scipy
matplotlib
tensorflow>=2.0.0a

The user may also optionally install torch>=1.3 to run the neurophox.torch module.

Getting started

Installation

There are currently two options to install neurophox:

  1. Installation via pip:

    pip install neurophox
    
  2. Installation from source via pip:

    git clone https://github.com/solgaardlab/neurophox
    pip install -e neurophox
    pip install -r requirements.txt
    

If installing other dependencies manually, ensure you install PyTorch (since PyTorch mesh layers are currently in development) and Tensorflow 2.0.

Using the GPU

If using a GPU, we recommend using a conda environement to install GPU dependencies using CUDA 10.0 with the following commands:

conda install pytorch torchvision cudatoolkit=10.0 -c pytorch
pip install tensorflow-gpu==2.0.0-alpha0

Imports

import numpy as np
from neurophox.numpy import RMNumpy
from neurophox.tensorflow import RM

N = 16

tf_layer = RM(N)
np_layer = RMNumpy(N, phases=tf_layer.phases)

np.allclose(tf_layer.matrix, np_layer.matrix)  # True
np.allclose(tf_layer(np.eye(N)), np_layer.matrix)  # True

Inspection

We can inspect the parameters for each layer using neurophox.control.MeshPhases which can be accessed via tf_layer.phases and np_layer.phases.

We can inspect the matrix elements implemented by each layer as follows via tf_layer.matrix and np_layer.matrix.

Bottleneck

The bottleneck in time for a neurophox mesh layer is not dominated by the number of inputs/outputs (referred to as units or \(N\)), but rather the number of “vertical layers” in the mesh layer (referred to as num_layers or \(L\)). This is to say that fairly large unitary matrices can be implemented by neurophox as long as the num_layers parameter is kept small (e.g., empirically, \(\log N\) layers seems to be enough for butterfly mesh architectures used in unitary mesh-based RNNs).

Visualizations

Phase shift settings visualization

The phase shift patterns used to generate the above propagation patterns can also be visualized by plotting np_layer.phases:

Rectangular mesh:

neurophox

Triangular mesh:

neurophox

Light propagation visualization

For the phase shift settings above, we can visualize the propagation of light (field magnitude), as the data “flows” through the mesh.

Rectangular mesh:

neurophox

Triangular mesh:

neurophox

The code to generate these visualization examples are provided in neurophox-notebooks.

Small machine learning example

It is simple to compose neurophox Tensorflow layers into unitary neural networks using Sequential to solve machine learning problems.

import tensorflow as tf
from neurophox.tensorflow import RM
from neurophox.ml.nonlinearities import cnorm, cnormsq

ring_model = tf.keras.Sequential([
    RM(3, activation=tf.keras.layers.Activation(cnorm)),
    RM(3, activation=tf.keras.layers.Activation(cnorm)),
    RM(3, activation=tf.keras.layers.Activation(cnorm)),
    RM(3, activation=tf.keras.layers.Activation(cnorm)),
    tf.keras.layers.Activation(cnormsq),
    tf.keras.layers.Lambda(lambda x: tf.math.real(x[:, :2])), # get first 2 output ports,
    tf.keras.layers.Activation('softmax')
])

ring_model.compile(
    loss='categorical_crossentropy',
    optimizer=tf.keras.optimizers.Adam(lr=0.0025)
)

Below is a visualization for many planar classification problems (including the ring model defined above):

neurophox neurophox

The code to generate the above example is provided in neurophox-notebooks.

Authors and citing this repository

neurophox was written by Sunil Pai (email: sunilpai@stanford.edu).

If you find this repository useful, please cite at least one of the following papers depending on your application:

  1. Calibration of optical neural networks

    @article{pai2019parallel,
      title={Parallel fault-tolerant programming of an arbitrary feedforward photonic network},
      author={Pai, Sunil and Williamson, Ian AD and Hughes, Tyler W and Minkov, Momchil and Solgaard, Olav and Fan, Shanhui and Miller, David AB},
      journal={arXiv preprint arXiv:1909.06179},
      year={2019}
    }
    
  2. Optimization of unitary mesh networks:

    @article{pai_matrix_2019,
      author = {Pai, Sunil and Bartlett, Ben and Solgaard, Olav and Miller, David A. B.},
      doi = {10.1103/PhysRevApplied.11.064044},
      journal = {Physical Review Applied},
      month = jun,
      number = {6},
      pages = {064044},
      title = {Matrix Optimization on Universal Unitary Photonic Devices},
      volume = {11},
      year = {2019}
    }
    
  3. Optical neural network nonlinearities:

    @article{williamson_reprogrammable_2020,
      author = {Williamson, I. A. D. and Hughes, T. W. and Minkov, M. and Bartlett, B. and Pai, S. and Fan, S.},
      doi = {10.1109/JSTQE.2019.2930455},
      issn = {1077-260X},
      journal = {IEEE Journal of Selected Topics in Quantum Electronics},
      month = jan,
      number = {1},
      pages = {1-12},
      title = {Reprogrammable Electro-Optic Nonlinear Activation Functions for Optical Neural Networks},
      volume = {26},
      year = {2020}
    }
    

Contributions

neurophox is under active development and is not yet stable.

If you find a bug or would like to recommend a feature, please submit an issue on Github.

We welcome pull requests and contributions from the broader community. If you would like to contribute to the codebase, please submit a pull request and title your branch bug/bug-fix-title or feature/feature-title.

Documentation

neurophox is under active development and is not yet stable, so expect to see the documentation change frequently.

Summary

The layer APIs of neurophox are provided via the subpackages:

  • neurophox.tensorflow

  • neurophox.numpy

  • neurophox.torch

Some machine learning experiment starter code with neurophox can be found in neurophox-notebooksand/or neurophox.ml.

All other subpackages have code relevant to the inner workings of neurophox.