如何在tensorflow中实现stdp？

Posted 2021-04-08

我正在尝试在tensorflow中实现STDP（Spike-Timing Dependent Plasticity）。这有点复杂。任何想法（完全在张量流图中运行）？

它的工作方式如下：说我有2个输入神经元，它们通过这个矩阵连接到3个输出神经元：[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]]（输入神经元0连接到输出神经元0和1 ......）。

假设输入神经元有这些尖峰（2个神经元，7个时间步长）：

Input Spikes:
[[0, 0, 1, 1, 0, 1, 0],
 [1, 1, 0, 0, 0, 0, 1]]

这些输出神经元的峰值（3个神经元，7个时间步长）：

Output Spikes:
[[0, 0, 0, 1, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 1, 1, 1]]

现在，对于每个非零权重，我想计算一个dw。例如，对于连接到输出神经元0的输入神经元0：

输入神经元0的尖峰的时间戳是[2, 3, 5]，输出神经元0的时间戳是[3, 6]。现在，我计算所有增量时间：

Delta Times = [ 2-3, 2-6, 3-3, 3-6, 5-3, 5-6 ] = [ -1, -4, 0, -3, 2, -1 ]

然后，我计算一些函数（实际的STDP函数，这对于这个问题并不重要 - 一些指数的东西）

 dw = SUM [ F(-1), F(-4), F(0), F(-3), F(2), F(-1) ]

这是连接输入神经元0和输出神经元0的权重的dw。对所有非零权重重复。

所以我可以在numpy中完成所有这些，但我希望能够在单个张量流图中完全完成。特别是，我坚持计算delta时间。以及如何对所有非零权重并行执行所有这些操作。

这是实际的stdp函数，btw（常量可以是参数）：

def stdp_f(x):
    return tf.where(
        x == 0, np.zeros(x.shape), tf.where(
            x > 0, 1.0 * tf.exp(-1.0 * x / 10.0), -1.0 * 1.0 * tf.exp(x / 10.0)))

---

关于性能的说明：下面的@jdehesa给出的方法既正确又巧妙。但事实证明它也很慢。特别地，对于输入400个神经元的784个输入神经元的真实神经网络，超过500个时间步长，spike_match =步骤执行（784,1,500,1）和（1,400,1,500）张量的乘法。

答案

我不熟悉STDP，所以我希望我理解你的意思。我认为这符合你的描述：

import tensorflow as tf

def f(x):
    # STDP function
    return x * 1

def stdp(input_spikes, output_spikes):
    input_shape = tf.shape(input_spikes)
    t = input_shape[-1]
    # Compute STDP function for all possible time difference values
    stdp_values = f(tf.cast(tf.range(-t + 1, t), dtype=input_spikes.dtype))
    # Arrange in matrix such that position [i, j] contains f(i - j)
    matrix_idx = tf.expand_dims(tf.range(t - 1, 2 * t - 1), 1) + tf.range(0, -t, -1)
    stdp_matrix = tf.gather(stdp_values, matrix_idx)
    # Find spike matches
    spike_match = (input_spikes[:, tf.newaxis, :, tf.newaxis] *
                   output_spikes[tf.newaxis, :, tf.newaxis, :])
    # Sum values where there are spike matches
    return tf.reduce_sum(spike_match * stdp_matrix, axis=(2, 3))

# Test
input_spikes = [[0, 0, 1, 1, 0, 1, 0],
                [1, 1, 0, 0, 0, 0, 1]]
output_spikes = [[0, 0, 0, 1, 0, 0, 1],
                 [1, 0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 1, 1, 1]]
with tf.Graph().as_default(), tf.Session() as sess:
    ins = tf.placeholder(tf.float32, [None, None])
    outs = tf.placeholder(tf.float32, [None, None])
    res = stdp(ins, outs)
    res_val = sess.run(res, feed_dict={ins: input_spikes, outs: output_spikes})
    print(res_val)
    # [[ -7.  10. -15.]
    #  [-13.   7. -24.]]

在这里我假设f可能很昂贵（并且它的值对于每对神经元都是相同的），所以我只为每个可能的时间增量计算一次，然后在矩阵中重新分配计算值，所以我可以乘以输入和输出尖峰发生的坐标对。

我使用f的标识函数作为占位符，因此结果值实际上只是这种情况下时间差的总和。

编辑：仅供参考，用您包含的STDP函数替换f：

def f(x):
    return tf.where(x == 0,
                    tf.zeros_like(x),
                    tf.where(x > 0,
                             1.0 * tf.exp(-1.0 * x / 10.0),
                             -1.0 * 1.0 * tf.exp(x / 10.0)))

结果是：

[[-3.4020822   2.1660795  -5.694256  ]
 [-2.974073    0.45364904 -3.1197631 ]]