814 lines
30 KiB
Rust
814 lines
30 KiB
Rust
//! Reward-modulated STDP, bit-exact with `model/plasticity.ts`.
|
|
//!
|
|
//! Pre-post spike timing writes an eligibility trace; a caller-supplied scalar reward turns that
|
|
//! trace into a change in a per-edge gain. Nothing else in the library learns.
|
|
|
|
use crate::bitset::{BitSet, RankedBitSet};
|
|
use crate::dataset::BrainDataset;
|
|
use crate::error::{bail, Result};
|
|
use crate::jsmath::{exp, js_max, js_min, tanh};
|
|
use crate::pool::{SharedSlice, WorkerPool};
|
|
use crate::version::{fnv1a32_step, version_for};
|
|
|
|
/// Version of the default configuration; kept verbatim so existing checkpoints stay loadable.
|
|
pub const PLASTICITY_VERSION: &str = "fly-kc-mbon-rstdp-v2";
|
|
|
|
/// Prefix of a derived version string. Deliberately not the full default string.
|
|
const PLASTICITY_VERSION_PREFIX: &str = "rstdp-v2:";
|
|
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct PlasticityConfig {
|
|
/// Anatomical role a plastic edge's source must belong to.
|
|
pub pre_role: String,
|
|
/// Anatomical role a plastic edge's target must belong to.
|
|
pub post_role: String,
|
|
/// Maximum number of plastic edges (strongest positive weights win).
|
|
pub budget: usize,
|
|
/// Eligibility-trace decay time constant in milliseconds.
|
|
pub trace_ms: f64,
|
|
/// Spike-pair exponential time constant in milliseconds.
|
|
pub pair_ms: f64,
|
|
/// Largest spike interval that still pairs, in milliseconds.
|
|
pub pair_window_ms: f64,
|
|
/// Eligibility added for a causal (pre before post) pair at dt = 0.
|
|
pub potentiation: f64,
|
|
/// Eligibility subtracted for an anti-causal (post before pre) pair at dt = 0.
|
|
pub depression: f64,
|
|
/// Gain step per unit of modulator times eligibility.
|
|
pub learning_rate: f64,
|
|
/// Restoring pull back towards a gain of 1, applied on reinforcement only.
|
|
pub restoring: f64,
|
|
/// Lower gain clamp.
|
|
pub min_gain: f64,
|
|
/// Upper gain clamp.
|
|
pub max_gain: f64,
|
|
}
|
|
|
|
/// Original constants: strongest 16,384 positive KC->MBON edges, gains clamped to [0.9, 1.1].
|
|
impl Default for PlasticityConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
pre_role: "kenyon".to_string(),
|
|
post_role: "mbon".to_string(),
|
|
budget: 16_384,
|
|
trace_ms: 5000.0,
|
|
pair_ms: 20.0,
|
|
pair_window_ms: 100.0,
|
|
potentiation: 0.1,
|
|
depression: 0.05,
|
|
learning_rate: 0.002,
|
|
restoring: 0.0001,
|
|
min_gain: 0.9,
|
|
max_gain: 1.1,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl PlasticityConfig {
|
|
/// Numeric parameters that define the learning rule, in version-hash order.
|
|
///
|
|
/// Role names and budget are deliberately absent: they change which edges are selected, which
|
|
/// the topology hash already covers.
|
|
fn version_params(&self) -> [f64; 9] {
|
|
[
|
|
self.trace_ms,
|
|
self.pair_ms,
|
|
self.pair_window_ms,
|
|
self.potentiation,
|
|
self.depression,
|
|
self.learning_rate,
|
|
self.restoring,
|
|
self.min_gain,
|
|
self.max_gain,
|
|
]
|
|
}
|
|
}
|
|
|
|
/// `PLASTICITY_VERSION` for the default rule constants, otherwise `rstdp-v2:<fnv1a32>`.
|
|
pub fn plasticity_version(config: &PlasticityConfig) -> String {
|
|
version_for(
|
|
PLASTICITY_VERSION,
|
|
PLASTICITY_VERSION_PREFIX,
|
|
&config.version_params(),
|
|
&PlasticityConfig::default().version_params(),
|
|
)
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct LearningStats {
|
|
pub version: String,
|
|
pub enabled: bool,
|
|
pub synapses: usize,
|
|
/// Count of selected edges (hash group 1). Historical field name.
|
|
pub mushroom: u64,
|
|
/// Always 0: no second edge group exists. Kept for checkpoint/UI compatibility.
|
|
pub output: u64,
|
|
pub updates: f64,
|
|
pub changed: u64,
|
|
pub mean_change: f64,
|
|
pub max_change: f64,
|
|
pub signal: f64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct PlasticityState {
|
|
pub version: String,
|
|
pub topology: u32,
|
|
pub enabled: bool,
|
|
pub updates: f64,
|
|
pub signal: f64,
|
|
pub gains: Vec<f32>,
|
|
pub traces: Vec<f32>,
|
|
pub touched: Vec<f64>,
|
|
}
|
|
|
|
/// Phenomenological three-factor plasticity; anatomical sites, not fitted dopamine compartments.
|
|
#[derive(Debug, Clone)]
|
|
pub struct RewardModulatedStdp {
|
|
pub config: PlasticityConfig,
|
|
/// Which base edges are plastic, and the slot of each.
|
|
///
|
|
/// This was a dense `Vec<i32>` of "slot, or -1" with one entry per base edge: 10.8 MB for
|
|
/// `fafb-v783`, read once per propagated edge, which is enough traffic on its own to evict the
|
|
/// membrane arrays from L2. The ranked bitset is 506 KB and answers "not plastic" — the
|
|
/// overwhelming majority — from a single word. A rank *is* the old slot number, because the
|
|
/// selected edges are numbered in ascending edge order either way.
|
|
slots: RankedBitSet,
|
|
pub edges: Vec<u32>,
|
|
pub gains: Vec<f32>,
|
|
pub traces: Vec<f32>,
|
|
pub touched: Vec<f64>,
|
|
/// Version string of this configuration; written into exported state.
|
|
pub version: String,
|
|
pub enabled: bool,
|
|
sources: Vec<u32>,
|
|
/// `data.targets[edges[slot]]`, cached so `observe` never touches the base CSR arrays.
|
|
slot_targets: Vec<u32>,
|
|
selected: Vec<u8>,
|
|
/// `incoming[neuron]` and `outgoing.get(neuron)` from the original, flattened into CSR over
|
|
/// neurons. Both preserve ascending edge order, which keeps pairing order identical to the
|
|
/// original base-edge traversal; an absent entry is an empty range, matching the original's
|
|
/// `if (!outgoing) continue`.
|
|
incoming: SlotIndex,
|
|
outgoing: SlotIndex,
|
|
/// `exp(-dt / pair_ms)` for every integer `dt` the pairing window admits, and
|
|
/// `exp(-gap / trace_ms)` for every short integer gap. `observe` calls `exp` about 7,800 times
|
|
/// per tick on `fafb-v783` -- two per pairing -- and every one of those arguments is an integer
|
|
/// number of milliseconds divided by a fixed constant, so almost all of them are one of a few
|
|
/// hundred values. See [`table_exp`]: this is a lookup that returns identical bits, not an
|
|
/// approximation, and the fallback path is the original call.
|
|
pair_decay: Vec<f64>,
|
|
trace_decay: Vec<f64>,
|
|
topology: u32,
|
|
updates: f64,
|
|
signal: f64,
|
|
}
|
|
|
|
/// Per-neuron slot lists in CSR form, plus a bit per neuron saying whether its list is non-empty.
|
|
///
|
|
/// The bitset is the whole point of the type. `observe` runs for every spiking neuron — about
|
|
/// 1,730 per tick on `fafb-v783` — and asking "does this neuron have plastic edges?" through
|
|
/// `offsets` costs two random loads from a 557 KB array, which is four L3 hits per spike between
|
|
/// the two indices. Almost every answer is no: only 4,730 of 139,255 neurons carry a plastic edge.
|
|
/// A 17 KB bitset gives the same answer out of L1.
|
|
#[derive(Debug, Clone, Default)]
|
|
struct SlotIndex {
|
|
offsets: Vec<u32>,
|
|
slots: Vec<u32>,
|
|
present: BitSet,
|
|
}
|
|
|
|
impl SlotIndex {
|
|
/// Build from `(neuron, slot)` pairs already in ascending slot order.
|
|
fn build(neurons: usize, pairs: impl Iterator<Item = (u32, u32)> + Clone) -> Self {
|
|
let mut counts = vec![0u32; neurons + 1];
|
|
let mut present = BitSet::new(neurons);
|
|
for (neuron, _) in pairs.clone() {
|
|
counts[neuron as usize + 1] += 1;
|
|
present.insert(neuron as usize);
|
|
}
|
|
for index in 1..counts.len() {
|
|
counts[index] += counts[index - 1];
|
|
}
|
|
let offsets = counts;
|
|
let mut cursor = offsets.clone();
|
|
let mut slots = vec![0u32; offsets[neurons] as usize];
|
|
for (neuron, slot) in pairs {
|
|
let at = &mut cursor[neuron as usize];
|
|
slots[*at as usize] = slot;
|
|
*at += 1;
|
|
}
|
|
Self {
|
|
offsets,
|
|
slots,
|
|
present,
|
|
}
|
|
}
|
|
|
|
#[inline]
|
|
fn get(&self, neuron: u32) -> &[u32] {
|
|
let from = self.offsets[neuron as usize] as usize;
|
|
let to = self.offsets[neuron as usize + 1] as usize;
|
|
&self.slots[from..to]
|
|
}
|
|
|
|
#[inline]
|
|
fn words(&self) -> &[u64] {
|
|
self.present.words()
|
|
}
|
|
}
|
|
|
|
/// Entries in the trace-decay table: `exp(-gap / trace_ms)` for gaps of 0..1023 ms.
|
|
///
|
|
/// 8 KB. Traces are touched about a quarter of a tick apart on average, so the gap is nearly always
|
|
/// a few milliseconds; a longer gap falls through to `exp` and costs what it always did.
|
|
const TRACE_TABLE: usize = 1024;
|
|
|
|
/// Upper bound on the pairing-window table, so a pathological `pair_window_ms` cannot allocate.
|
|
const MAX_PAIR_TABLE: usize = 4096;
|
|
|
|
/// `exp(-argument / scale)`, read from `table` when `argument` is an integer the table covers.
|
|
///
|
|
/// `table[k]` is built as `exp(-(k as f64) / scale)`, the same expression with the same operands,
|
|
/// so a hit returns the same bits the call would — this is a lookup, not an approximation. Both of
|
|
/// the kernel's `exp` arguments are integer millisecond differences, because `ms` is an integer
|
|
/// tick counter and `touched` and `lastSpikeMs` only ever hold an `ms`; the integer test is there
|
|
/// because `observe` and `reinforce` are public and a caller may pass anything.
|
|
#[inline]
|
|
fn table_exp(table: &[f64], argument: f64, scale: f64) -> f64 {
|
|
let index = argument as usize;
|
|
if index as f64 == argument {
|
|
if let Some(value) = table.get(index) {
|
|
return *value;
|
|
}
|
|
}
|
|
exp(-argument / scale)
|
|
}
|
|
|
|
/// The pairing-window constants `observe` reads, grouped so a shard takes one argument.
|
|
struct PairWindow {
|
|
pair_ms: f64,
|
|
pair_window_ms: f64,
|
|
potentiation: f64,
|
|
depression: f64,
|
|
trace_ms: f64,
|
|
}
|
|
|
|
/// The read-only slot indices and decay tables `observe` reads, likewise.
|
|
struct SlotArrays<'a> {
|
|
incoming: &'a SlotIndex,
|
|
outgoing: &'a SlotIndex,
|
|
sources: &'a [u32],
|
|
slot_targets: &'a [u32],
|
|
pair_decay: &'a [f64],
|
|
trace_decay: &'a [f64],
|
|
}
|
|
|
|
/// The slots in `list` that lie in `offset..limit`, as one contiguous window.
|
|
///
|
|
/// Every per-neuron list is built in ascending slot order ([`SlotIndex::build`] is fed pairs in
|
|
/// slot order), so the shard's slots are a contiguous run found by binary search rather than a
|
|
/// per-slot range test. `SHARDED` is false for the unsharded walk, where the window is the whole
|
|
/// list and the search folds away.
|
|
#[inline(always)]
|
|
fn slot_window<const SHARDED: bool>(list: &[u32], offset: usize, limit: usize) -> &[u32] {
|
|
if !SHARDED {
|
|
return list;
|
|
}
|
|
let from = list.partition_point(|slot| (*slot as usize) < offset);
|
|
let to = list.partition_point(|slot| (*slot as usize) < limit);
|
|
&list[from..to]
|
|
}
|
|
|
|
/// Observe this tick's spike pairs for the slots in `offset..limit`, whose traces and touch
|
|
/// stamps are `traces` and `touched`.
|
|
///
|
|
/// # Determinism
|
|
///
|
|
/// **For a fixed slot the sequence of trace writes is the sequential one, whatever the shard
|
|
/// bounds** — the same argument [`crate::lif`]'s `propagate_shard` rests on. The sequential walk
|
|
/// visits the spike list in order and, per spiking neuron, its arriving slots then its leaving
|
|
/// slots, each in ascending slot order. A shard walks the same spike list in the same order and
|
|
/// the same lists in the same order, and writes exactly the slots in `offset..limit` — a
|
|
/// subsequence of the sequential write sequence, which preserves the relative order of every pair
|
|
/// it keeps. Restricted to one slot it is therefore the sequential sequence for that slot, and a
|
|
/// trace write depends on nothing but that slot's own previous value and touch stamp: the lazy
|
|
/// decay, the clamp and the single rounding to f32 happen at the same points with the same
|
|
/// operands. Slots in different shards never interact and every slot is in exactly one shard.
|
|
///
|
|
/// `last_spike` is read, never written — the kernel stamps this tick's spikes after `observe`
|
|
/// returns — so the shards need no ordering between them at all.
|
|
#[allow(clippy::too_many_arguments)]
|
|
#[inline(always)]
|
|
fn observe_shard<const SHARDED: bool>(
|
|
traces: &mut [f32],
|
|
touched: &mut [f64],
|
|
offset: usize,
|
|
limit: usize,
|
|
spikes: &[u32],
|
|
last_spike: &[f64],
|
|
ms: f64,
|
|
index: &SlotArrays<'_>,
|
|
window: &PairWindow,
|
|
) {
|
|
if traces.is_empty() {
|
|
return;
|
|
}
|
|
let (incoming_words, outgoing_words) = (index.incoming.words(), index.outgoing.words());
|
|
for neuron in spikes.iter().copied() {
|
|
// Most spiking neurons carry no plastic edge in either direction, and both bitsets fit
|
|
// in L1, so the common case is two loads from the same pair of cache lines and no
|
|
// touch of `offsets`, `slots`, `sources` or `slot_targets` at all.
|
|
let word = neuron as usize >> 6;
|
|
let bit = 1u64 << (neuron & 63);
|
|
let (arriving, leaving) = (incoming_words[word], outgoing_words[word]);
|
|
if (arriving | leaving) & bit == 0 {
|
|
continue;
|
|
}
|
|
// Causal: each slot arriving at the spiking neuron.
|
|
if arriving & bit != 0 {
|
|
let list = slot_window::<SHARDED>(index.incoming.get(neuron), offset, limit);
|
|
for slot in list.iter().copied() {
|
|
let dt = ms - last_spike[index.sources[slot as usize] as usize];
|
|
if dt > 0.0 && dt <= window.pair_window_ms {
|
|
let pair =
|
|
window.potentiation * table_exp(index.pair_decay, dt, window.pair_ms);
|
|
write_trace(
|
|
traces,
|
|
touched,
|
|
index.trace_decay,
|
|
window.trace_ms,
|
|
slot as usize - offset,
|
|
ms,
|
|
pair,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
// Anti-causal: each slot leaving the spiking neuron.
|
|
if leaving & bit != 0 {
|
|
let list = slot_window::<SHARDED>(index.outgoing.get(neuron), offset, limit);
|
|
for slot in list.iter().copied() {
|
|
let dt = ms - last_spike[index.slot_targets[slot as usize] as usize];
|
|
if dt > 0.0 && dt <= window.pair_window_ms {
|
|
let pair = -window.depression * table_exp(index.pair_decay, dt, window.pair_ms);
|
|
write_trace(
|
|
traces,
|
|
touched,
|
|
index.trace_decay,
|
|
window.trace_ms,
|
|
slot as usize - offset,
|
|
ms,
|
|
pair,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// The one lazy trace update: advance the trace to `ms`, add `pair`, clamp, stamp.
|
|
///
|
|
/// A free function rather than a method so `observe` can hold the slot indices and the trace
|
|
/// arrays borrowed at the same time.
|
|
#[inline]
|
|
fn write_trace(
|
|
traces: &mut [f32],
|
|
touched: &mut [f64],
|
|
decay_table: &[f64],
|
|
trace_ms: f64,
|
|
slot: usize,
|
|
ms: f64,
|
|
pair: f64,
|
|
) {
|
|
let decay = table_exp(decay_table, js_max(0.0, ms - touched[slot]), trace_ms);
|
|
let value = f64::from(traces[slot]) * decay + pair;
|
|
traces[slot] = js_max(-1.0, js_min(1.0, value)) as f32;
|
|
touched[slot] = ms;
|
|
}
|
|
|
|
impl RewardModulatedStdp {
|
|
pub fn new(data: &BrainDataset, config: PlasticityConfig) -> Self {
|
|
let version = plasticity_version(&config);
|
|
let pre: std::collections::HashSet<u32> =
|
|
data.role(&config.pre_role).iter().copied().collect();
|
|
let post: std::collections::HashSet<u32> =
|
|
data.role(&config.post_role).iter().copied().collect();
|
|
|
|
// Candidates: positive, non-self edges from a preRole neuron to a postRole neuron.
|
|
let mut candidates: Vec<(u32, u32)> = Vec::new();
|
|
for source in 0..data.meta.neurons as u32 {
|
|
let from = data.indptr[source as usize] as usize;
|
|
let to = data.indptr[source as usize + 1] as usize;
|
|
for edge in from..to {
|
|
let target = data.targets[edge];
|
|
if data.weights[edge] <= 0 || source == target {
|
|
continue;
|
|
}
|
|
if pre.contains(&source) && post.contains(&target) {
|
|
candidates.push((edge as u32, source));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fixed deterministic budget: strongest anatomical pre->post connections.
|
|
//
|
|
// `(a, b) => weights[b] - weights[a] || a - b` is a total order (edge indices are
|
|
// unique), so a stable sort with the same comparator reproduces V8's result exactly.
|
|
candidates.sort_by(|a, b| {
|
|
let weight_a = f64::from(data.weights[a.0 as usize]);
|
|
let weight_b = f64::from(data.weights[b.0 as usize]);
|
|
match weight_b.partial_cmp(&weight_a) {
|
|
Some(std::cmp::Ordering::Equal) | None => a.0.cmp(&b.0),
|
|
Some(order) => order,
|
|
}
|
|
});
|
|
candidates.truncate(config.budget);
|
|
candidates.sort_by_key(|candidate| candidate.0);
|
|
|
|
let count = candidates.len();
|
|
// `candidates` is edge-sorted and edge indices are unique, so the k-th set bit belongs to
|
|
// the k-th slot: a rank in this set is exactly the slot the dense array used to store.
|
|
let slots = RankedBitSet::from_ascending(
|
|
data.meta.edges,
|
|
candidates.iter().map(|(edge, _)| *edge as usize),
|
|
);
|
|
let mut hash: u32 = 2_166_136_261;
|
|
for (edge, source) in candidates.iter().copied() {
|
|
// A negative Int16 weight sign-extends into the int32 XOR, as in JavaScript.
|
|
for value in [
|
|
edge as i32,
|
|
source as i32,
|
|
data.targets[edge as usize] as i32,
|
|
i32::from(data.weights[edge as usize]),
|
|
1,
|
|
] {
|
|
hash = fnv1a32_step(hash, value);
|
|
}
|
|
}
|
|
|
|
let pair_span = if config.pair_window_ms >= 0.0 {
|
|
(config.pair_window_ms as usize).min(MAX_PAIR_TABLE)
|
|
} else {
|
|
0
|
|
};
|
|
let pair_decay = (0..=pair_span)
|
|
.map(|dt| exp(-(dt as f64) / config.pair_ms))
|
|
.collect();
|
|
let trace_decay = (0..TRACE_TABLE)
|
|
.map(|gap| exp(-(gap as f64) / config.trace_ms))
|
|
.collect();
|
|
|
|
let neurons = data.meta.neurons;
|
|
// `candidates` is edge-sorted, so building in slot order keeps each per-neuron list in
|
|
// ascending edge order.
|
|
let incoming = SlotIndex::build(
|
|
neurons,
|
|
candidates
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(slot, (edge, _))| (data.targets[*edge as usize], slot as u32)),
|
|
);
|
|
let outgoing = SlotIndex::build(
|
|
neurons,
|
|
candidates
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(slot, (_, source))| (*source, slot as u32)),
|
|
);
|
|
|
|
Self {
|
|
slots,
|
|
edges: candidates.iter().map(|candidate| candidate.0).collect(),
|
|
sources: candidates.iter().map(|candidate| candidate.1).collect(),
|
|
slot_targets: candidates
|
|
.iter()
|
|
.map(|candidate| data.targets[candidate.0 as usize])
|
|
.collect(),
|
|
// Every candidate belongs to hash group 1; there is no second edge group.
|
|
selected: vec![1; count],
|
|
gains: vec![1.0; count],
|
|
traces: vec![0.0; count],
|
|
touched: vec![0.0; count],
|
|
incoming,
|
|
outgoing,
|
|
pair_decay,
|
|
trace_decay,
|
|
topology: hash,
|
|
updates: 0.0,
|
|
signal: 0.0,
|
|
enabled: true,
|
|
version,
|
|
config,
|
|
}
|
|
}
|
|
|
|
/// Effective multiplier for a base edge; 1 for every non-plastic edge.
|
|
#[inline]
|
|
pub fn gain(&self, edge: usize) -> f64 {
|
|
match self.slots.rank_of(edge) {
|
|
None => 1.0,
|
|
Some(slot) => f64::from(self.gains[slot as usize]),
|
|
}
|
|
}
|
|
|
|
/// Plastic slot of a base edge, or `None` for the immutable majority.
|
|
#[inline]
|
|
pub fn slot_of(&self, edge: usize) -> Option<u32> {
|
|
self.slots.rank_of(edge)
|
|
}
|
|
|
|
/// The plastic-edge presence bitset, so a loop over a run of consecutive edge indices — a CSR
|
|
/// row — loads one word and asks [`RewardModulatedStdp::gain_in_word`] per edge.
|
|
#[inline]
|
|
pub fn slot_words(&self) -> &[u64] {
|
|
self.slots.words()
|
|
}
|
|
|
|
/// [`RewardModulatedStdp::gain`] for `edge`, given the already-loaded bitset word `edge >> 6`.
|
|
#[inline]
|
|
pub fn gain_in_word(&self, edge: usize, word: u64) -> f64 {
|
|
if word & (1u64 << (edge & 63)) == 0 {
|
|
return 1.0;
|
|
}
|
|
f64::from(self.gains[self.slots.rank_within(edge, word) as usize])
|
|
}
|
|
|
|
/// Bytes the plastic-edge index occupies; quoted by the crate README.
|
|
pub fn slot_index_bytes(&self) -> usize {
|
|
self.slots.heap_bytes()
|
|
}
|
|
|
|
pub fn topology(&self) -> u32 {
|
|
self.topology
|
|
}
|
|
|
|
pub fn updates(&self) -> f64 {
|
|
self.updates
|
|
}
|
|
|
|
pub fn signal(&self) -> f64 {
|
|
self.signal
|
|
}
|
|
|
|
pub fn clear_eligibility(&mut self, ms: f64) {
|
|
self.traces.fill(0.0);
|
|
self.touched.fill(ms);
|
|
self.signal = 0.0;
|
|
}
|
|
|
|
/// Called once per tick by the kernel, *before* this tick's spikes are written into
|
|
/// `last_spike`, so a pair needs a strictly positive `dt`.
|
|
pub fn observe(&mut self, spikes: &[u32], count: usize, last_spike: &[f64], ms: f64) {
|
|
self.observe_sharded(spikes, count, last_spike, ms, None);
|
|
}
|
|
|
|
/// Slot-range boundaries for `workers` workers, length `workers + 1`.
|
|
///
|
|
/// An even split of the slot range. Slots are numbered in ascending base-edge order, which is
|
|
/// source-major, so a range is a contiguous run of pre-role neurons' edges; the long per-neuron
|
|
/// lists belong to the post-role neurons and are spread across the whole range, which is where
|
|
/// nearly all of the work is. Balance is not a correctness property either way — see
|
|
/// [`observe_shard`].
|
|
pub fn slot_shards(&self, workers: usize) -> Vec<u32> {
|
|
let slots = self.gains.len();
|
|
let workers = workers.max(1);
|
|
let chunk = slots.div_ceil(workers).max(1);
|
|
(0..=workers)
|
|
.map(|worker| (chunk * worker).min(slots) as u32)
|
|
.collect()
|
|
}
|
|
|
|
/// [`RewardModulatedStdp::observe`], optionally sharded by slot range over a worker pool.
|
|
///
|
|
/// `shards` is a pool and the boundaries [`RewardModulatedStdp::slot_shards`] returned for it.
|
|
/// The result is identical to the sequential run for any worker count; the argument is in
|
|
/// [`observe_shard`].
|
|
pub fn observe_sharded(
|
|
&mut self,
|
|
spikes: &[u32],
|
|
count: usize,
|
|
last_spike: &[f64],
|
|
ms: f64,
|
|
shards: Option<(&WorkerPool, &[u32])>,
|
|
) {
|
|
if !self.enabled {
|
|
return;
|
|
}
|
|
let Self {
|
|
config,
|
|
traces,
|
|
touched,
|
|
sources,
|
|
slot_targets,
|
|
incoming,
|
|
outgoing,
|
|
pair_decay,
|
|
trace_decay,
|
|
..
|
|
} = self;
|
|
let window = PairWindow {
|
|
pair_ms: config.pair_ms,
|
|
pair_window_ms: config.pair_window_ms,
|
|
potentiation: config.potentiation,
|
|
depression: config.depression,
|
|
trace_ms: config.trace_ms,
|
|
};
|
|
let spikes = &spikes[..count];
|
|
let slots = traces.len();
|
|
let index = SlotArrays {
|
|
incoming,
|
|
outgoing,
|
|
sources,
|
|
slot_targets,
|
|
pair_decay,
|
|
trace_decay,
|
|
};
|
|
|
|
let Some((pool, bounds)) =
|
|
shards.filter(|(pool, bounds)| bounds.len() == pool.workers() + 1)
|
|
else {
|
|
observe_shard::<false>(
|
|
traces, touched, 0, slots, spikes, last_spike, ms, &index, &window,
|
|
);
|
|
return;
|
|
};
|
|
let traces = SharedSlice::new(traces);
|
|
let touched = SharedSlice::new(touched);
|
|
pool.broadcast(&|worker| {
|
|
let (from, to) = (bounds[worker] as usize, bounds[worker + 1] as usize);
|
|
// SAFETY: `bounds` is ascending and ends at `slots`, so worker `w` owns
|
|
// `bounds[w]..bounds[w + 1]` and no two workers' ranges overlap. `broadcast` runs each
|
|
// index exactly once and returns only after every borrow here has been dropped.
|
|
let (traces, touched) = unsafe { (traces.range(from, to), touched.range(from, to)) };
|
|
observe_shard::<true>(
|
|
traces,
|
|
touched,
|
|
from,
|
|
to,
|
|
spikes,
|
|
last_spike,
|
|
ms,
|
|
&index,
|
|
&window,
|
|
);
|
|
});
|
|
}
|
|
|
|
pub fn reinforce(&mut self, reward: f64, ms: f64) {
|
|
if !self.enabled || !reward.is_finite() || reward == 0.0 {
|
|
return;
|
|
}
|
|
let learning_rate = self.config.learning_rate;
|
|
let restoring = self.config.restoring;
|
|
let min_gain = self.config.min_gain;
|
|
let max_gain = self.config.max_gain;
|
|
let trace_ms = self.config.trace_ms;
|
|
self.signal = tanh(reward);
|
|
let signal = self.signal;
|
|
let Self {
|
|
traces,
|
|
touched,
|
|
gains,
|
|
trace_decay,
|
|
..
|
|
} = self;
|
|
let mut changed = false;
|
|
for slot in 0..gains.len() {
|
|
write_trace(traces, touched, trace_decay, trace_ms, slot, ms, 0.0);
|
|
let old = gains[slot];
|
|
// Small restoring term limits long-run drift; excitation never changes sign.
|
|
let next = f64::from(old) + learning_rate * signal * f64::from(traces[slot])
|
|
- restoring * (f64::from(old) - 1.0);
|
|
gains[slot] = js_max(min_gain, js_min(max_gain, next)) as f32;
|
|
changed |= gains[slot] != old;
|
|
}
|
|
if changed {
|
|
self.updates += 1.0;
|
|
}
|
|
}
|
|
|
|
pub fn statistics(&self) -> LearningStats {
|
|
let mut changed = 0u64;
|
|
let mut sum = 0.0f64;
|
|
let mut max = 0.0f64;
|
|
let mut mushroom = 0u64;
|
|
for index in 0..self.gains.len() {
|
|
let delta = (f64::from(self.gains[index]) - 1.0).abs();
|
|
if delta > 0.000001 {
|
|
changed += 1;
|
|
}
|
|
sum += delta;
|
|
max = js_max(max, delta);
|
|
mushroom += u64::from(self.selected[index]);
|
|
}
|
|
let synapses = self.edges.len();
|
|
LearningStats {
|
|
version: self.version.clone(),
|
|
enabled: self.enabled,
|
|
synapses,
|
|
mushroom,
|
|
output: synapses as u64 - mushroom,
|
|
updates: self.updates,
|
|
changed,
|
|
mean_change: sum / if synapses == 0 { 1.0 } else { synapses as f64 },
|
|
max_change: max,
|
|
signal: self.signal,
|
|
}
|
|
}
|
|
|
|
pub fn export_state(&self) -> PlasticityState {
|
|
PlasticityState {
|
|
version: self.version.clone(),
|
|
topology: self.topology,
|
|
enabled: self.enabled,
|
|
updates: self.updates,
|
|
signal: self.signal,
|
|
gains: self.gains.clone(),
|
|
traces: self.traces.clone(),
|
|
touched: self.touched.clone(),
|
|
}
|
|
}
|
|
|
|
/// Validate everything before mutating anything.
|
|
pub fn import_state(&mut self, state: Option<&PlasticityState>) -> Result<()> {
|
|
let Some(state) = state else {
|
|
self.gains.fill(1.0);
|
|
self.traces.fill(0.0);
|
|
self.touched.fill(0.0);
|
|
self.updates = 0.0;
|
|
self.signal = 0.0;
|
|
return Ok(());
|
|
};
|
|
if state.version != self.version || state.topology != self.topology {
|
|
bail!("Incompatible plasticity topology/version");
|
|
}
|
|
if state.gains.len() != self.gains.len()
|
|
|| state.traces.len() != self.traces.len()
|
|
|| state.touched.len() != self.touched.len()
|
|
{
|
|
bail!("Invalid plasticity dimensions");
|
|
}
|
|
if state.updates.fract() != 0.0
|
|
|| !state.updates.is_finite()
|
|
|| state.updates < 0.0
|
|
|| !state.signal.is_finite()
|
|
|| state.signal.abs() > 1.0
|
|
{
|
|
bail!("Invalid plasticity metadata");
|
|
}
|
|
// 0.100001 for the default clamps: the widest legal displacement plus a float32 tolerance.
|
|
let radius = js_max(1.0 - self.config.min_gain, self.config.max_gain - 1.0) + 0.000001;
|
|
for index in 0..self.gains.len() {
|
|
let gain = f64::from(state.gains[index]);
|
|
let trace = f64::from(state.traces[index]);
|
|
let touched = state.touched[index];
|
|
if !gain.is_finite()
|
|
|| (gain - 1.0).abs() > radius
|
|
|| !trace.is_finite()
|
|
|| trace.abs() > 1.0
|
|
|| !touched.is_finite()
|
|
|| touched < 0.0
|
|
{
|
|
bail!("Invalid plasticity values");
|
|
}
|
|
}
|
|
self.gains.copy_from_slice(&state.gains);
|
|
self.traces.copy_from_slice(&state.traces);
|
|
self.touched.copy_from_slice(&state.touched);
|
|
self.enabled = state.enabled;
|
|
self.updates = state.updates;
|
|
self.signal = state.signal;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
/// A table hit has to be the same bits as the call it replaces, and a miss has to fall through.
|
|
#[test]
|
|
fn the_exp_tables_are_a_lookup_not_an_approximation() {
|
|
for scale in [20.0f64, 5000.0, 1.0 / 3.0] {
|
|
let table: Vec<f64> = (0..256).map(|k| exp(-(k as f64) / scale)).collect();
|
|
for k in 0..256 {
|
|
assert_eq!(
|
|
table_exp(&table, k as f64, scale).to_bits(),
|
|
exp(-(k as f64) / scale).to_bits(),
|
|
"integer {k} at scale {scale}"
|
|
);
|
|
}
|
|
// Past the end, and non-integers inside it, go to `exp`.
|
|
for argument in [256.0, 1e9, 0.5, 12.25, -3.0, f64::NAN, f64::INFINITY] {
|
|
let want = exp(-argument / scale);
|
|
let got = table_exp(&table, argument, scale);
|
|
assert_eq!(
|
|
got.to_bits(),
|
|
want.to_bits(),
|
|
"non-hit {argument} at scale {scale}"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|