Спайковые нейросети на Swift, часть II: Синапсы / Хабр

Всем привет! В первой части мы обсудили концепцию SNN в общих чертах, выбрали модель Ижикевича для нейрона и реализовали её на Swift. Напомню основные тезисы:

Спайковые нейронные сети (SNN) стремятся точнее повторять биологические.
Как и живые нейроны, искусственные спайковые нейроны взаимодействуют друг с другом короткими импульсами — спайками.
Спайковые нейроны способны интегрировать информацию и обучаться, но механика этих процессов отличается от «традиционных» нейросетей. Там, где важно подчеркнуть различие, будем называть их ANN.
Спайковые нейроны, как и биологические, могут быть разных типов — тормозящими, возбуждающими. Кроме того, они могут отличаться паттернами генерации спайков.
Степень возбуждения нейрона определяется значением мембранного потенциала — некоторого физического свойства нейрона.

Но от одного нейрона нет никакой практической пользы. Интересные свойства проявляются только на уровне больших сетей, когда нейроны взаимодействуют друг с другом. Поэтому в этой статье реализуем второй важный элемент — синапсы. Через них наши нейроны смогут общаться. При моделировании синапсов снова будем ориентироваться на биологические механизмы.

Входы и выходы нейронов

В биологическом мозге у нейронов нет глаз. Они живут в полной темноте, получая от соседей лишь малый кусочек информации о мире в виде серии спайков. Исключением, пожалуй, являются только клетки-рецепторы, которые отвечают за сенсорное восприятие и непосредственно реагируют на изменения в окружающей среде. Затем они передают информацию глубоким структурам для обработки. В этих структурах, в свою очередь, рождаются такие эффекты, как память и сознание.

То есть, всё самое интересное в нейронных сетях возникает в результате интенсивного взаимодействия нервных клеток. Но, прежде чем развивать эту тему, позаимствуем некоторые важные понятия из мира нейрофизиологии.

Клетка Пуркинье с разветвлённой сетью дендритов. Рамон-и-Кахаль

Дендриты. Отростки нервной клетки, на которых расположены входы нейронов. Могут принимать и интегрировать сигналы от тысяч других клеток.

Аксоны. Выходы нервной клетки. По ним распространяются спайки.

Терминаль. Окончание аксона. Из неё под воздействием спайка выделяются химические вещества — медиаторы.

Синапс. Устойчивый контакт между выходом и входом. То есть, между терминалью одного нейрона и дендритом другого. В более редких случаях соединяет аксон с телом, с другим дендритом и т.п.

Пресинаптический нейрон. Не��рон, который производит спайки и выделяет медиаторы.

Постсинаптический нейрон. Приёмник, который поглощает медиаторы. Может сам быть пресинаптическим для других нейронов сети.

Медиаторы — химические вещества, которыми обмениваются пресинаптический и постсинаптический нейроны.

Вес синапса. В модели — чувствительность синапса к медиаторам.

Таким образом, синапс — это связующее звено между нейронами, через которое они поддерживают общение. Как сочетание вилки и розетки. При этом синапсы бывают двух типов: химические и электрические. Далее рассмотрим основной тип — химический.

Обмен сигналами между нейронами через химические синапсы происходит следующим образом:

Пресинаптический нейрон производит спайк.
Спайк распространяется по аксону, в конце концов попадая на терминаль — окончание аксона.
В терминали начинаются процессы, которые приводят к выбросу медиаторов в окружающее пространство — синаптическую щель.
Партнёр по синапсу — постсинаптический нейрон — впитывает эти медиаторы.
Захваченные медиаторы служат чем-то вроде ключей к замку. Они увеличивают проводимость мембраны нейрона. Через неё начинает течь ток.
Ток изменяет мембранный потенциал: он повышается (возбуждение) или снижается (торможение).
Если мембранный потенциал постсинаптического нейрона достигает нужного уровня, в нём также возникает спайк.

Спайк попадает на терминаль и вызывает выброс медиатора. Мембранный потенциал соседнего нейрона увеличивается

Видно, что обмен является опосредованным — спайк пресинаптического нейрона лишь подталкивают соседа сделать собственный.

Так, если синапс слабый, нейрон будет реагировать на подталкивания неохотно. Но, по мере стимуляции, чувствительность повышается. Это часть процесса обучения — синапсы постепенно «раскачиваются» или угасают. Фактически, вся информация в мозге хранится в виде синаптических весов.

Будет синапс тормозящим или возбуждающим, определяется действием медиаторов. Тормозящие медиаторы заставляют ток вытекать из нейрона, разряжая его и снижая мембранный потенциал. Возбуждающие — наоборот.

Перейдём к моделированию поведения синапса.

Моделирование синапса

Для начала вспомним модель Ижикевича, которую мы выбрали в предыдущей части для описания поведения нейрона:

$\begin{cases} C\dot{v} = k\,(v - v_r)\,(v - v_t) - u + I(t) \\ \dot{u} = a\big[b\,(v - v_r) - u\big] \end{cases}$

Ток , который изменяет мембранный потенциал нейрона, теперь станет совокупностью всех токов, которые протекают через синапсы. Токи эти возникают под воздействием входящих спайков.

В реальных нейронах интеграция может быть очень сложна — токи складываются или гасят друг друга. В нашем случае общий ток это просто $I(t)=\sum I_{syn}(t)$ , где

$I_{syn}(t) = E(t)-vG(t)$

То есть, мы рассматриваем нейрон как точечный объект и простой сумматор. Здесь $I_{syn}$ — ток через каждый синапс нейрона, а

$G(t)=\sum g_i(t),\qquad E(t)=\sum g_i(t)\,E_i(t),$

где это текущая суммарная проводимость синапса, — константы. Сейчас нам нужно знать про них только то, что они зависят от типа синапса.

С проводимостью всё сложнее. Это тоже динамическая характеристика, которая зависит от времени. Поэтому для описания нам потребуется ещё одно дифференциальное уравнение:

$\tau_s \dot{g} = -g + w_sS \sum x\delta (t - t^f)$

Проводимость постоянно угасает со скоростью $\tau_s$ . Переменная — вес синапса, — константа масштабирования, — фактор пластичности синапса в момент спайка.

Выражение с суммой равно 1 тогда, когда в момент зарегистрирован спайк, и 0 в других случаях. Соответственно — время регистрации спайка пресинаптического нейрона, $\delta$ — функция Дирака.

Как и говорили ранее, важен только факт спайка. Так, если спайки перестанут поступать, синапс со временем "закроется", ток прекратится: $G(t) \to 0, I_{syn} \to 0$ при $t \to \infty$ . А в момент спайка проводимости увеличатся на $w_s\sum{x}$ и значение мембранного потеницала снова придёт в движение.

В нашем подходе мы не будем передавать в модель Ижикевича напрямую, а вместо этого будем передавать суммарные значения и по всем синапсам.

Подводя итог, каждый новый спайк увеличивает проводимость синапса. Проводимость — величина, обратная сопротивлению из закона Ома. Её рост увеличит интенсивность токов через этот синапс. И если общий ток со всех синапсов окажется достаточно сильным, произойдёт спайк. Так происходит передача информации между нейронами.

Обучение синапсов и нейронов

Свойство синапсов и нейронов изменять свою чувствительность ("вес") и обучаться называют пластичностью. Пластичность бывает нескольких видов, в этой части мы рассмотрим один — Spike-Timing Dependent Plasticity или STDP.

Spike-Timing Dependent Plasticity

В этом случае вес синапса изменяется на достаточно длительный период, но медленно убывает со временем. Этот вид пластичности лежит в основе обучения Хебба, затронутое в первой части. Вес синапса изменяется по следующим правилам:

$\begin{aligned} \Delta w=A_{+}e^{(t_{pre}-t_{post})\tau_{+}}, t_{post}\ge t_{pre} \\ \Delta w=-A_{-}e^{-(t_{pre}-t_{post})\tau_{-}}, t_{post}<t_{pre} \end{aligned}$

Корректировка выполняется по факту каждого спайка с любой стороны. Идея следующая: когда пресинаптический нейрон спайкует непосредственно до постсинаптического, это увеличивает вес. А если после — уменьшает. И чем больше разница во времени, тем меньше эффект.

То есть, если один нейрон спайкует сразу за другим, связь между ними усиливаетс��. Это усиление зависит от того, насколько "сразу" был спайк. Если через длительный промежуток времени, то считается, что причинно-следственной связи между спайками нет и на вес они почти не влияют. За это отвечают экспоненты формулах.

В то же время, если первый нейрон постоянно спайкует, а второй на это сразу не реагирует, вес связи будет уменьшаться.

Параметры возьмём такие: $\tau_+=\tau_-=100,~A_+=1,~A_-=2$ . То есть, штрафы будут в два раза сильнее бонусов. Два последних параметра можно трактовать как "скорость обучения" (learning rate).

Со временем, с каждым новым спайком, значение будет бесконтрольно расти или станет отрицательным. Чтобы этого избежать, вес нужно ограничивать. Мы будем делать это так:

$\hat{w}=\frac{hw_{max}}{1+(\frac{w}{\theta(1-w})^{-\gamma}}$

Эта функция обладает высоким контрастом в середине, что важно для эффективного обучения. В то же время, она ограничивает вес слева и справа. Здесь $h=1,~\theta=1.25,~\gamma=6$ , а $w_{max}$ — ограничение сверху, своё для каждого типа пресинаптического нейрона.

Теперь пришло время реализовать эту модель. Приступим.

Реализация синапса

Моделирование синапсов также будем выполнять на Swift. Сначала определим базовые структуры:

struct SynapseParams {
	let presynType: NeuronType // Тип синапса (возбуждающий, тормозящий)
	let weightLimit: Double // ограничение на вес "сверху"
}

struct SynapseOutput: Hashable, Identifiable {
  	var w: Double  // вес синапса
	var G: Double = 0 // суммарная проводимость
	var giEi: Double = 0
	var id: ObjectIdentifier // идентификатор
	var key: NeuronKey // пресинаптический ключ
	var stsp: Double // зарезервировано
	var E: Double // зарезервировано

	static func == (lhs: SynapseOutput, rhs: SynapseOutput) -> Bool {
		return lhs.id == rhs.id
	}

	func hash(into hasher: inout Hasher) {
		hasher.combine(id)
	}
}

// Идентификатор нейрона. Потребуется в будущем
struct NeuronKey: Hashable {
	var x: Int = 0
	var y: Int = 0
	var T: Int = 0
	var L: Int = 0
	var C: Int = 0
}

// представление в виде строки для удобства
extension NeuronKey: CustomStringConvertible {
	var description: String {
		"[(\(x), \(y)), T:\(T), L:\(L), C:\(C)]"
	}
}

Теперь структуру, которая будет отвечать за проводимости:

// разные типы проводимости gi
enum ConductanceType {
	case ampa, nmda, gabaa, gabab
}

struct Conductance {
	let tau: Double
	var g: Double = 0.0
	let type: ConductanceType
	let ratio: Double // удельный вес этой проводимости
	let E_L: Double // обратный потенциал

    // глутаматные синапсы NMDA зависит от потенциала,
    // учитываем это здесь
	private func nmda(_ v: Double) -> Double {
		let y = pow((v + 80)/60, 2)
		return y/(1 + y)
	}

    // расчёт проводимости
	func adjustedG(_ v: Double) -> Double {
		let s = switch type {
		case .nmda: nmda(v)*ratio
		default: ratio
		}
		
		return g * s
	}

    // флаг равен true, когда проводимость около нуля
	var isStable: Bool { g < EPSILON }

    // угасание проводимости со временем
	mutating func decay() {
		g -= dt * g / tau
	}

    // задаём начальные параметры. Значения E_L
    // и tau стандартные для своих типов
	init(_ type: ConductanceType, ratio: Double = 1.0)
	{
		switch type {
		case .ampa: E_L = 0; tau = 5
		case .gabaa: E_L = -70; tau = 6
		case .gabab: E_L = -90; tau = 150
		case .nmda: E_L = 0; tau = 150
		}
		self.ratio = ratio
		self.type = type
	}
}

Теперь базовый класс для синапса:

actor Synapse: Identifiable {
	private var w: Double // вес
	private let tau_s: Double = 20 // скорость угасания
	private let scale: Double = 0.01 // переменная масштабирования
	private	var c: [Conductance] = [] // вектор проводимостей
	nonisolated let presynKey: NeuronKey
	nonisolated let postsynKey: NeuronKey
	nonisolated let params: SynapseParams // параметры
	private var postSpikeT: Double = -Double.infinity // время постспайка
	private var presSpikeT: Double = -Double.infinity // время преспайка
	private var postsynV: Double = 0 // постсинаптичнский потенциал
	private let wGABA: Double = 4.0 // постоянный вес для ГАМК
	private let Am: Double = 2.0 // LR для штрафа к весу
	private let Ap: Double = 1.0 // LR для бонуса
  	private var r: Double = 0.5 // зарезервировано
    private var stspX: Double = 1.0 // зарезервировано
	private var E: Double = 0 // зарезервировано

    // Синапс в стабильном состоянии, когда все проводимости
    // околонулевые
	var isStable: Bool {
		c.allSatisfy { $0.isStable }
	}
			
	init(
		postsynKey: NeuronKey, presynKey: NeuronKey,
		params: SynapseParams)
	{
		self.params = params
		self.postsynKey = postsynKey
		self.presynKey = presynKey

        // тормозящие синапсы -- ГАМКовые
        // возбуждающие синапсы -- глутаматные (AMPA, NMDA)
		c = switch params.presynType {
		case .inhibitory:
			[Conductance(.gabaa), Conductance(.gabab)]
		case .excitatory:
			[Conductance(.ampa), Conductance(.nmda)]
		}

        // инициализируем начальные веса
		self.w = switch params.presynType {
		case .excitatory:
			Double.random(in: 0...1)
		case .inhibitory:
			wGABA
		}
	}
}

Для каждого типа синапса (возбуждающий, тормозящий) будет свой набор проводимостей
Веса w инициализируем случайно из диапазона [0, 1] для каждого синапса с возбуждающими нейронами
Синапсы с тормозящими нейронами не обучаются, их вес фиксирован

Добавим основную функциональность:

extension Synapse {	
    // режем проводимости по экспоненциальному
    // закону. Если был спайк -- обновляем веса. Дальше
    // рассчитываем общие проводимости по всем типам
    // и возвращаем
  	func update(_ wasSpike: Bool, v: Double) -> SynapseOutput {
		if wasSpike {
			w = updateWeights(.postsynSpike)
		}
		
		let output = c.reduce(
			into: SynapseOutput(w: w, id: id, key: presynKey, 
                                stsp: stspX, E: E)
        ) { x, y in
			x.giEi += y.adjustedG(v) * y.E_L
			x.G += y.adjustedG(v)
		}

		c.indices.forEach { j in c[j].decay() }
		
		return output
	}

    // если получили спайк от пресинаптического нейрона,
    // корректируем веса соответственно и увеличиваем
    // проводимости
	func presynSpike(_ stsp: Double = 1.0) {
		w = updateWeights(.presynSpike)
		for j in c.indices { c[j].g += w * scale * x }
	}
}

И, наконец, реализуем обучение синапса через STDP. Сначала служебные структуры:

fileprivate enum SpikeType {
	case postsynSpike, presynSpike
}

// расширяем Double, добавляем наш ограничитель весов
fileprivate extension Double {
	func smooth(_ scale: Self) -> Self {
		let (theta, gamma, h) = (1.25, 6.0, 1.0)
		return h*scale/(1 + pow(self/(theta*(1-self)), -gamma))
	}
}

Теперь основную логику:

extension Synapse {
	private func updateWeights(_ type: SpikeType) -> Double {
		var dW: Double = 0
        guard params.presynType == .excitatory
		else { return w }

		let globalTimer = Double(
			atomicGt.load(ordering: .sequentiallyConsistent)
		)
		
		switch type {
		case .presynSpike:
			let exp0 = -(globalTimer - postSpikeT)/tau_s
			presSpikeT = globalTimer
			dW = -Am * exp(exp0)
		case .postsynSpike:
			let exp0 = (presSpikeT - globalTimer)/tau_s
			postSpikeT = globalTimer
			dW = +Ap * exp(exp0)
		}

        let updatedW = w + dW
      
		return updatedW.smooth(params.weightLimit)
	}
}

Здесь globalTimer — счётчик времени, полученный из глобального атомика atomicGt.

В прошлой части мы возбуждали наш нейрон, передавая ему значение тока напрямую, как константу. Теперь это можно делать через синапс:

var wasPostsynSpike: Bool = false
var injectPresynSpike: Bool = false
let dummyKey = NeuronKey(x: 0, y: 0)
let isRunning = false

// let params = IzhikevichParams(...)
// let neuron = Izhikevich(params)

let sParams = SynapseParams(presynType: .excitatory, weightLimit: 10.0)
let s = Synapse(postsynKey: dummyKey, presynKey: dummyKey, params: sParams)

while isRunning {
  // 1. Вручную подаём входящий спайк
  if injectPresynSpike {
    await s.presynSpike()
    injectPresynSpike = false
  }

  // 2. и смотрим что получилось
  let synOut = await synapse.update(wasPostsynSpike, v: neuron.V)

  // 3. Передаём проводимости на вход нейрону
  wasPostsynSpike = neuron.updateState(synOut, constI: 0) 

  if wasPostsynSpike {
    //реагируем на выходной спайк, если он был
  }
}

В этом примере мы больше не передаём машине состояний нейрона значение тока напрямую, а используем проводимости synOut, которые получили с синапса
Подача входящего спайка в синапс делается через presynSpike()
injectPresynSpike контролирует, когда нужно подавать спайк

Контроллер Playground из прошлой части теперь можно модифицировать так, чтобы он работал с синапсом, а не напрямую с нейроном.

А пока на этом всё. В следующей части соберём все ингредиенты вместе.

Спайковые нейросети на Swift, часть I: Принципы и модель

Спасибо за прочтение!