136 lines
6.7 KiB
TypeScript
136 lines
6.7 KiB
TypeScript
import type { BrainDataset } from '../../src/dataset/format';
|
|
|
|
export const PLASTICITY_VERSION = 'fly-kc-mbon-rstdp-v2';
|
|
const TRACE_MS = 5000;
|
|
const PAIR_MS = 20;
|
|
export interface LearningStats {
|
|
version: string; enabled: boolean; synapses: number; mushroom: number; output: number;
|
|
updates: number; changed: number; meanChange: number; maxChange: number; signal: number;
|
|
}
|
|
export interface PlasticityState {
|
|
version: string; topology: number; enabled: boolean; updates: number; signal: number;
|
|
gains: Float32Array; traces: Float32Array; touched: Float64Array;
|
|
}
|
|
|
|
/** Phenomenological three-factor plasticity; anatomical sites, not fitted dopamine compartments. */
|
|
export class Plasticity {
|
|
readonly slots: Int32Array;
|
|
readonly edges: Uint32Array;
|
|
readonly gains: Float32Array;
|
|
readonly traces: Float32Array;
|
|
readonly touched: Float64Array;
|
|
private readonly sources: Uint32Array;
|
|
private readonly mushroom: Uint8Array;
|
|
private readonly incoming: number[][];
|
|
private readonly outgoing: Map<number, number[]> = new Map();
|
|
private readonly topology: number;
|
|
private updates = 0;
|
|
private signal = 0;
|
|
enabled = true;
|
|
|
|
constructor(private readonly data: BrainDataset) {
|
|
const kc = new Set(data.meta.roles.kenyon ?? []);
|
|
const mbon = new Set(data.meta.roles.mbon ?? []);
|
|
const candidates: [number, number, number][] = [];
|
|
for (let source = 0; source < data.meta.neurons; source++) {
|
|
for (let edge = data.indptr[source]; edge < data.indptr[source + 1]; edge++) {
|
|
const target = data.targets[edge];
|
|
if (data.weights[edge] <= 0 || source === target) continue;
|
|
const mb = kc.has(source) && mbon.has(target);
|
|
if (mb) candidates.push([edge, source, 1]);
|
|
}
|
|
}
|
|
// Fixed deterministic budget: strongest anatomical KC→MBON connections.
|
|
const chosen = candidates.sort((a, b) => data.weights[b[0]] - data.weights[a[0]] || a[0] - b[0]).slice(0, 16384).sort((a, b) => a[0] - b[0]);
|
|
this.slots = new Int32Array(data.meta.edges).fill(-1);
|
|
this.edges = Uint32Array.from(chosen.map(row => row[0]));
|
|
this.sources = Uint32Array.from(chosen.map(row => row[1]));
|
|
this.mushroom = Uint8Array.from(chosen.map(row => row[2]));
|
|
this.gains = new Float32Array(chosen.length).fill(1);
|
|
this.traces = new Float32Array(chosen.length);
|
|
this.touched = new Float64Array(chosen.length);
|
|
this.incoming = Array.from({ length: data.meta.neurons }, () => []);
|
|
let hash = 2166136261;
|
|
chosen.forEach(([edge, source, group], slot) => {
|
|
this.slots[edge] = slot;
|
|
this.incoming[data.targets[edge]].push(slot);
|
|
// chosen is edge-sorted: preserve the former base-CSR traversal order.
|
|
const outgoing = this.outgoing.get(source) ?? [];
|
|
outgoing.push(slot);
|
|
this.outgoing.set(source, outgoing);
|
|
for (const value of [edge, source, data.targets[edge], data.weights[edge], group]) hash = Math.imul(hash ^ value, 16777619) >>> 0;
|
|
});
|
|
this.topology = hash;
|
|
}
|
|
|
|
gain(edge: number): number { const slot = this.slots[edge]; return slot < 0 ? 1 : this.gains[slot]; }
|
|
|
|
clearEligibility(ms: number): void { this.traces.fill(0); this.touched.fill(ms); this.signal = 0; }
|
|
|
|
private trace(slot: number, ms: number, pair: number): void {
|
|
const decay = Math.exp(-Math.max(0, ms - this.touched[slot]) / TRACE_MS);
|
|
this.traces[slot] = Math.max(-1, Math.min(1, this.traces[slot] * decay + pair));
|
|
this.touched[slot] = ms;
|
|
}
|
|
|
|
observe(spikes: Uint32Array, count: number, lastSpike: Float64Array, ms: number): void {
|
|
if (!this.enabled) return;
|
|
for (let i = 0; i < count; i++) {
|
|
const neuron = spikes[i];
|
|
for (const slot of this.incoming[neuron]) {
|
|
const dt = ms - lastSpike[this.sources[slot]];
|
|
if (dt > 0 && dt <= 100) this.trace(slot, ms, 0.1 * Math.exp(-dt / PAIR_MS));
|
|
}
|
|
const outgoing = this.outgoing.get(neuron);
|
|
if (!outgoing) continue;
|
|
for (const slot of outgoing) {
|
|
const dt = ms - lastSpike[this.data.targets[this.edges[slot]]];
|
|
if (dt > 0 && dt <= 100) this.trace(slot, ms, -0.05 * Math.exp(-dt / PAIR_MS));
|
|
}
|
|
}
|
|
}
|
|
|
|
reinforce(reward: number, ms: number): void {
|
|
if (!this.enabled || !Number.isFinite(reward) || reward === 0) return;
|
|
this.signal = Math.tanh(reward);
|
|
let changed = false;
|
|
for (let slot = 0; slot < this.edges.length; slot++) {
|
|
this.trace(slot, ms, 0);
|
|
const old = this.gains[slot];
|
|
// Small restoring term limits long-run drift; excitation never changes sign.
|
|
const next = old + 0.002 * this.signal * this.traces[slot] - 0.0001 * (old - 1);
|
|
this.gains[slot] = Math.max(0.9, Math.min(1.1, next));
|
|
changed ||= this.gains[slot] !== old;
|
|
}
|
|
if (changed) this.updates++;
|
|
}
|
|
|
|
statistics(): LearningStats {
|
|
let changed = 0, sum = 0, max = 0, mushroom = 0;
|
|
for (let i = 0; i < this.gains.length; i++) {
|
|
const delta = Math.abs(this.gains[i] - 1);
|
|
if (delta > 0.000001) changed++;
|
|
sum += delta; max = Math.max(max, delta); mushroom += this.mushroom[i];
|
|
}
|
|
return { version: PLASTICITY_VERSION, enabled: this.enabled, synapses: this.edges.length, mushroom, output: this.edges.length - mushroom,
|
|
updates: this.updates, changed, meanChange: sum / (this.edges.length || 1), maxChange: max, signal: this.signal };
|
|
}
|
|
|
|
exportState(): PlasticityState {
|
|
return { version: PLASTICITY_VERSION, topology: this.topology, enabled: this.enabled, updates: this.updates, signal: this.signal,
|
|
gains: this.gains.slice(), traces: this.traces.slice(), touched: this.touched.slice() };
|
|
}
|
|
|
|
importState(state?: PlasticityState): void {
|
|
if (!state) { this.gains.fill(1); this.traces.fill(0); this.touched.fill(0); this.updates = 0; this.signal = 0; return; }
|
|
if (state.version !== PLASTICITY_VERSION || state.topology !== this.topology) throw new Error('Incompatible plasticity topology/version');
|
|
if (state.gains.length !== this.gains.length || state.traces.length !== this.traces.length || state.touched.length !== this.touched.length) throw new Error('Invalid plasticity dimensions');
|
|
if (typeof state.enabled !== 'boolean' || !Number.isInteger(state.updates) || state.updates < 0 || !Number.isFinite(state.signal) || Math.abs(state.signal) > 1) throw new Error('Invalid plasticity metadata');
|
|
for (let i = 0; i < this.gains.length; i++) {
|
|
const radius = 0.100001;
|
|
if (!Number.isFinite(state.gains[i]) || Math.abs(state.gains[i] - 1) > radius || !Number.isFinite(state.traces[i]) || Math.abs(state.traces[i]) > 1 || !Number.isFinite(state.touched[i]) || state.touched[i] < 0) throw new Error('Invalid plasticity values');
|
|
}
|
|
this.gains.set(state.gains); this.traces.set(state.traces); this.touched.set(state.touched);
|
|
this.enabled = state.enabled; this.updates = state.updates; this.signal = state.signal;
|
|
}
|
|
}
|