FixedPointFinder - A PyTorch / TensorFlow toolbox for finding fixed points and linearized dynamics in recurrent neural networks

Finds and analyzes the fixed points of recurrent neural networks that have been built using Tensorflow.

If you are using FixedPointFinder in research to be published, please cite our accompanying paper in your publication:

Golub and Sussillo (2018), "FixedPointFinder: A TensorFlow toolbox for identifying and characterizing fixed points in recurrent neural networks," Journal of Open Source Software, 3(31), 1003, https://doi.org/10.21105/joss.01003 .

Recommended Installation

1. Clone or download this repository.

2. Create a virtual environment for the required dependencies: To create a new virtual environment specific, enter at the command line:

   $ python3 -m venv --system-site-packages your-virtual-env-name

   where your-virtual-env-name is a path to the the virtual environment you would like to create (e.g.: /home/fpf). Then activate your new virtual environment:

   $ source your-virtual-env-name/bin/activate

   When you are finished working in your virtual environment (not now), enter:

   $ deactivate

3. Automatically assemble all dependencies using pip and the requirements*.txt files.

   For PyTorch, use:

   $ pip install -r requirements-torch.txt

   For TensorFlow, use:

   $ pip install -r requirements-tf.txt

Advanced Installation

Advanced Python users and those wishing to develop contributions may prefer a custom install. Such installs should adhere to the following general template:

1. Clone or download this repository.

2. Install compatible versions of the following prerequisites.

• NumPy, SciPy, Matplotlib (install SciPy stack, contains all of them).

• Scikit-learn (install).

• TensorFlow (recommended version: 2.8; requires at least version 1.14; versions beyond 2.8 are not currently supported) (install).

• RecurrentWhisperer (install).

3. Add the directories for FixedPointFinder and RecurrentWhisperer to your Python path:

   $ export PYTHONPATH=$PYTHONPATH:/path/to/your/directory/fixed-point-finder/
   $ export PYTHONPATH=$PYTHONPATH:/path/to/your/directory/recurrent-whisperer/

   where "/path/to/your/directory" is replaced with the path to the corresponding repository. This step must be performed each time you launch a new terminal to work with FixedPointFinder, and thus you may want to add the lines above to a startup script (e.g., the .bashrc / .bashprofile script in your home folder or an activate script in your virtual environment).

Example

FixedPointFinder includes an end-to-end example, implemented separately in PyTorch and TensorFlow, that trains an RNN to solve a task and then identifies and visualizes the fixed points of the trained RNN. To run the example, descend into the example directory: fixed-point-finder/examples/ and execute:

For PyTorch:

>>> python run_FlipFlop_torch.py

For TensorFlow:

>>> python run_FlipFlop_tf.py

The task is the "flip-flop" task previously described in Sussillo and Barak (2013). Briefly, the task is to implement a 3-bit binary memory, in which each of 3 input channels delivers signed transient pulses (-1 or +1) to a corresponding bit of the memory, and an input pulse flips the state of that memory bit (also -1 or +1) whenever a pulse's sign is opposite of the current state of the bit. The example trains a 16-unit LSTM RNN to solve this task (Fig. 1). Once the RNN is trained, the example uses FixedPointFinder to identify and characterize the trained RNN's fixed points. Finally, the example produces a visualization of these results (Fig. 2). In addition to demonstrating a working use of FixedPointFinder, this example provides a testbed for experimenting with different RNN architectures (e.g., numbers of recurrent units, LSTMs vs. GRUs vs. vanilla RNNs) and characterizing how these lower-level model design choices manifest in the higher-level dynamical implementation used to solve a task.

---

Figure 1. Inputs (gray), target outputs (cyan), and outputs of a trained LSTM RNN (purple) from an example trial of the flip-flop task. Signed input pulses (gray) flip the corresponding bit's state (green) whenever an input pulse has the opposite sign of the current bit state (e.g., if gray goes high when green is low). The RNN has been trained to nearly perfectly reproduce the target memory state (purple closely overlaps cyan).

---

Figure 2. Fixed-point structure of an LSTM RNN trained to solve the flip-flop task. FixedPointFinder identified 8 stable fixed points (black points), each of which corresponds to a unique state of the 3-bit memory. FixedPointFinder also identified a number of unstable fixed points (red points) along with their unstable modes (red lines), which mediate the set of state transitions trained into the RNN's dynamics. Here, each unstable fixed point is a "saddle" in the RNN's dynamical flow field, and the corresponding unstable modes indicate the directions that nearby states are repelled from the fixed point. State trajectories from example trials (blue) traverse about these fixed points. All quantities are visualized in the 3-dimensional space determined by the top 3 principal components computed across 128 example trials.

General Usage (PyTorch & TensorFlow)

1. Start by building, and if desired, training an RNN. FixedPointFinder works with Pytorch RNN objects (e.g., torch.nn.RNN, torch.nn.GRU) and Tensorflow RNNCell objects.

   Advanced: More generally, FixedPointFinder will work on any Pytorch or TensorFlow function f that satisfies the following:

   • f must be auto-differentiatiable.

   • f must map inputs and previous states to updated states.

   • f must match the argument specifications: _, h_next = f(input, h_prev)

     input: a tensor with shape (n, n_inputs) containing n inputs of dimension n_inputs.
     h_prev: a tensor with shape (n, n_states) containing n previous states of dimension n_states.
     h_next: a tensor with shape (n, n_states) containing the n updated states.

     Internally, f should map inputs[i] and h_prev[i] to h_next[i].

2. Build a FixedPointFinder object:

   • PyTorch: fpf = FixedPointFinder(your_rnn, **kwargs)
   • Tensorflow: fpf = FixedPointFinder(your_rnn_cell, tf_session, **kwargs)
     • Here, your_rnn_cell is the RNNCell that specifies the single-timestep transitions in your RNN, and tf_session is the Tensorflow session in which your model has been instantiated.

3. Specify the initial_states from which you'd like to initialize the local optimizations implemented by FixedPointFinder. These data should conform to shape and type expected by your_rnn_cell. For Tensorflow's BasicRNNCell, this would mean an (n, n_states) numpy array, where n is the number of initializations and n_states is the dimensionality of the RNN state (i.e., the number of hidden units). For Tensorflow's LSTMCell, initial_states should be an LSTMStateTuple containing one (n, nstates) numpy array specifying the initializations of the hidden states and another (n, nstates) numpy array specifying the cell states.

4. Specify the inputs under which you'd like to study your RNN. Currently, To study the RNN given a set of static inputs, inputs should be a numpy array with shape (1, n_inputs) where n_inputs is an int specifying the depth of the inputs expected by your_rnn_cell. Alternatively, you can search for fixed points under different inputs by specifying a potentially different input for each initial states by making inputs a (n, n_inputs) numpy array.

5. Run the local optimizations that find the fixed points:

   >>> fps = fpf.find_fixed_points(initial_states, inputs)

   The fixed points identified, the Jacobian of your RNN state transition function at those points, and some metadata corresponding to the optimizations will be returned in the FixedPoints object.fps (see FixedPoints.py for more detail).

6. Finally, visualize the identified fixed points:

   >>> fps.plot()

   You can also visualize these fixed points amongst state trajectories from your RNN (see plot in FixedPoints.py and the example in run_FlipFlop_torch.py and run_FlipFlop_tf.py)

Testing the Package

Tests are not currently functional due to package upgrades in 2022-2023. That said, the rest of the codebase should be fully usable, including the 3-bit flip flop examples. Stay tuned.

Earlier versions of FixedPointFinder included a test suite for confirming successful installation, and for ensuring that contributions have not introduced bugs into the main control flow. The tests run FixedPointFinder over a set of RNNs where ground truth fixed points have been previously identified, numerically confirmed, and saved for comparison.

To run the tests, descend into the test directory: fixed-point-finder/test/ and execute:

>>> python run_test.py

Contribution Guidelines

Contributions are welcome. Please see the contribution guidelines.