Synaptic parameter sharing

Synaptic parameter sharing#

In this tutorial, you will learn how to:

  • flexibly share parameters of synapses

Here is a code snippet which you will learn to understand in this tutorial:

net = ...  # See tutorial on Basics of Jaxley.

# The same parameter for all synapses
net.make_trainable("Ionotropic_gS")

# An individual parameter for every synapse.
net.select(edges="all").make_trainable("Ionotropic_gS")

# Share synaptic conductances emerging from the same neurons.
net.copy_node_property_to_edges("cell_index")
sub_net = net.select(edges=[0, 1, 2])
sub_net.edges["controlled_by_param"] = sub_net.edges["pre_global_cell_index"]
sub_net.make_trainable("Ionotropic_gS")

In a previous tutorial about training networks, we briefly touched on parameter sharing. In this tutorial, we will show you how you can flexibly share parameters within a network.

import jaxley as jx
from jaxley.channels import Na, K, Leak
from jaxley.connect import fully_connect
from jaxley.synapses import IonotropicSynapse

Preface: Building the network#

We first build a network consisting of six neurons, in the same way as we showed in the previous tutorials:

dt = 0.025
t_max = 10.0

comp = jx.Compartment()
branch = jx.Branch(comp, ncomp=2)
cell = jx.Cell(branch, parents=[-1, 0])
net = jx.Network([cell for _ in range(6)])
fully_connect(net.cell([0, 1, 2]), net.cell([3, 4, 5]), IonotropicSynapse())

Sharing parameters by modifying controlled_by_param#

net.copy_node_property_to_edges("global_cell_index")

df = net.edges
df = df.query("pre_global_cell_index in [0, 1, 2]")
subnetwork = net.select(edges=df.index)

df = subnetwork.edges
df["controlled_by_param"] = df["pre_global_cell_index"]
subnetwork.make_trainable("IonotropicSynapse_gS")

Number of newly added trainable parameters: 3. Total number of trainable parameters: 3

Let’s look at this line by line. First, we exactly follow the previous tutorial in selecting the synapses which we are interested in training (i.e., the ones whose presynaptic neuron has index 0, 1, 2):

df = net.edges
df = df.query("pre_global_cell_index in [0, 1, 2]")
subnetwork = net.select(edges=df.index)

As second step, we enable parameter sharing. This is done by setting the controlled_by_param. Synapses that have the same value in controlled_by_param will be shared. Let’s inspect controlled_by_param before we modify it:

subnetwork.edges[["pre_global_cell_index", "controlled_by_param"]]
pre_global_cell_index controlled_by_param
0 0 0
1 0 1
2 0 2
3 1 3
4 1 4
5 1 5
6 2 6
7 2 7
8 2 8

Every synapse has a different value. Because of this, no synaptic parameters will be shared. To enable parameter sharing we override the controlled_by_param column with the presynaptic cell index:

df = subnetwork.edges
df["controlled_by_param"] = df["pre_global_cell_index"]

df[["pre_global_cell_index", "controlled_by_param"]]
pre_global_cell_index controlled_by_param
0 0 0
1 0 0
2 0 0
3 1 1
4 1 1
5 1 1
6 2 2
7 2 2
8 2 2

Now, all we have to do is to make these synaptic parameters trainable with the make_trainable() method:

subnetwork.make_trainable("IonotropicSynapse_gS")

Number of newly added trainable parameters: 3. Total number of trainable parameters: 6

It correctly says that we added three parameters (because we have three cells, and we share individual synaptic parameters). We now have 6 trainable parameters in total (because we already added 3 trainable parameters above).

A more involved example: sharing by pre- and post-synaptic cell type#

As an example, consider the following: We have a fully connected network of six cells. Each cell falls into one of three cell types:

from typing import Union, List

net = jx.Network([cell for _ in range(6)])
fully_connect(net.cell("all"), net.cell("all"), IonotropicSynapse())

net.cell([0, 1]).add_to_group("exc")
net.cell([2, 3]).add_to_group("inh")
net.cell([4, 5]).add_to_group("unknown")

We want to make all synapses that start from excitatory or inhibitory neurons trainable. In addition, we want to use the same parameter for synapses if they have the same pre- and post-synaptic cell type.

To achieve this, we will first want a column in net.nodes which indicates the cell type.

net.nodes["cell_type"] = net.nodes[["exc", "inh", "unknown"]].idxmax(axis=1)

net.nodes["cell_type"]

0         exc
1         exc
2         exc
3         exc
4         exc
5         exc
6         exc
7         exc
8         inh
9         inh
10        inh
11        inh
12        inh
13        inh
14        inh
15        inh
16    unknown
17    unknown
18    unknown
19    unknown
20    unknown
21    unknown
22    unknown
23    unknown
Name: cell_type, dtype: object

The cell_type is now part of the net.nodes. However, we would like to do parameter sharing of synapses based on the pre- and post-synaptic node values. To do so, we import the cell_type column into net.edges. To do this, we use the .copy_node_property_to_edges() which the name of the property you are copying from nodes:

net.copy_node_property_to_edges("cell_type")

After this, you have columns in the .edges which indicate the pre- and post-synaptic cell type:

net.edges[["pre_cell_type", "post_cell_type"]]
pre_cell_type post_cell_type
0 exc exc
1 exc exc
2 exc inh
3 exc inh
4 exc unknown
5 exc unknown
6 exc exc
7 exc exc
8 exc inh
9 exc inh
10 exc unknown
11 exc unknown
12 inh exc
13 inh exc
14 inh inh
15 inh inh
16 inh unknown
17 inh unknown
18 inh exc
19 inh exc
20 inh inh
21 inh inh
22 inh unknown
23 inh unknown
24 unknown exc
25 unknown exc
26 unknown inh
27 unknown inh
28 unknown unknown
29 unknown unknown
30 unknown exc
31 unknown exc
32 unknown inh
33 unknown inh
34 unknown unknown
35 unknown unknown

Next, we specify which parts of the network we actually want to change (in this case, all synapses which have excitatory or inhibitory presynaptic neurons):

df = net.edges
df = df.query(f"pre_cell_type in ['exc', 'inh']")
print(f"There are {len(df)} synapses to be changed.")

subnetwork = net.select(edges=df.index)

There are 24 synapses to be changed.

As the last step, we again have to specify parameter sharing by setting controlled_by_param. In this case, we want to share parameters that have the same pre- and post-synaptic neuron. We achieve this by grouping the synpases by their pre- and post-synaptic cell type (see pd.DataFrame.groupby for details):

# Step 6: use groupby to specify parameter sharing and make the parameters trainable.
subnetwork.edges["controlled_by_param"] = subnetwork.edges.groupby(["pre_cell_type", "post_cell_type"]).ngroup()
subnetwork.make_trainable("IonotropicSynapse_gS")

Number of newly added trainable parameters: 6. Total number of trainable parameters: 6

This created six trainable parameters, which makes sense as we have two types of pre-synaptic neurons (excitatory and inhibitory) and each has three options for the postsynaptic neuron (pre, post, unknown).

Summary#

In this tutorial, you learned how you can flexibly share synaptic parameters. This works by first using select() to identify which synapses to make trainable, and by then modifying controlled_by_param to customize parameter sharing.

Navigation