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

Motivation

Integrated optical neural networks or photonic neural networks are an ASIC technology that uses light as a computing medium as opposed to conventional analog electronics. Such devices are a new and exciting option for low-energy and practical machine learning accelerator technologies that can be deployed in data centers where optical fibers and photonic technologies are already used to transmit and process data.

Optical neural networks are composed of "optical matrix multipliers" (defined by two-port components arranged in a "unitary mesh" architecture as discussed in arxiv/1808.00458) and optical nonlinearities such as electro-optic activations.

The interesting property of the unitary mesh matrix multiplier are that they behave differently from conventional matrix multipliers:

1. They act as unitary operators rather than general linear operators, preserving the norm of the data flowing through the unitary mesh network.
2. The matrix elements are not directly trained during backpropagation. Instead, "phase shifts" or Euler angle parameters are trained.

This puts optical mesh networks in an interesting and relatively unexplored regime of machine learning. For example, orthogonal and unitary neural networks might circumvent vanishing and exploding gradient problems due to norm-preserving property. Such networks have been studied for synthetic long-term memory natural language processing tasks (see unitary mesh-based RNN, unitary evolution RNN, and orthogonal evolution RNN).

Introduction

Neurophox provides a robust and 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 (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 feedforward unitary mesh network.

Dependencies and requirements

Some important requirements for Neurophox are:

1. Python >=3.6
2. Tensorflow >=2.0.0a
3. PyTorch

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

numpy
scipy
matplotlib
tensorflow>=2.0.0a
torch>=1.1

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.

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: Triangular mesh:

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: Triangular mesh:

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

Small machine learning example

It is possible to compose Neurophox Tensorflow layers into unitary neural networks using tf.keras.Sequential to solve machine learning problems. Here we use absolute value nonlinearities and categorical cross entropy.

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 (we already know it is real from the activation),
    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:

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

Authors and citing this repository

Neurophox was written by Sunil Pai (email: [email protected]).

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

1. Optimization of unitary mesh networks:

   @article{pai2018matrix,
     title={Matrix optimization on universal unitary photonic devices},
     author={Pai, Sunil and Bartlett, Ben and Solgaard, Olav and Miller, David AB},
     journal={arXiv preprint arXiv:1808.00458},
     year={2018}
   }
   

2. Optical neural network nonlinearities:

   @article{williamson2019reprogrammable,
     title={Reprogrammable Electro-Optic Nonlinear Activation Functions for Optical Neural Networks},
     author={Williamson, Ian A.D. and Hughes, Tyler W. and Minkov, Momchil and Bartlett, Ben and Pai, Sunil and Fan, Shanhui},
     journal={arXiv preprint arXiv:1903.04579},
     year={2019}
   }
   

Future Work and Contributions

Neurophox is under development and is not yet stable.

If you find a bug, have a question, 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, please submit a pull request and title your branch bug/bug-fix-title or feature/feature-title.

Some future feature ideas include:

1. Implement multi-wavelength and multi-mode operators for photonics simulation using batch and broadcasting operations.
2. Add nonlinearities and a unitary mesh-based RNN cell for use in deep learning architectures.