292 lines
13 KiB
TypeScript
292 lines
13 KiB
TypeScript
import type { BrainDataset } from '../dataset/format';
|
|
import { RewardModulatedStdp, type PlasticityConfig, type PlasticityState } from './plasticity';
|
|
import { DEFAULT_RETINA_CONFIG, projectFrame, type RetinaConfig } from './retina';
|
|
import { Xorshift32 } from './rng';
|
|
import { versionFor } from './version';
|
|
|
|
/** Version of the default configuration; kept verbatim so existing checkpoints stay loadable. */
|
|
export const NEURAL_KERNEL_VERSION = 'lif-1ms-f64-v2';
|
|
|
|
/** Rate roles tracked by default, on top of every `command_*` role the dataset declares. */
|
|
const DEFAULT_EXTRA_RATE_ROLES = ['steer_left', 'steer_right', 'forward', 'backward', 'proboscis', 'reward_pam'];
|
|
|
|
/**
|
|
* Prefix of the macro-type populations, tracked by default like the `command_*` buttons.
|
|
*
|
|
* `docs/design/macros.md` section 12: a macro is a button, pressed by its own population, so its
|
|
* rate is read exactly as a button's is. A dataset that predates the roles simply has none, and a
|
|
* checkpoint that predates them restores them at zero (`importState`).
|
|
*/
|
|
const MACRO_RATE_ROLE_PREFIX = 'macro_';
|
|
|
|
/**
|
|
* Rate roles are packed into a per-neuron bitmask of {@link ROLE_MASK_WORDS} 32-bit words, one
|
|
* bit per role.
|
|
*
|
|
* 64 rather than 32 since `docs/design/macros.md` section 11: the eight `command_*` buttons and
|
|
* the six historical motor/reward roles are fourteen, and the twenty-two `macro_<type>`
|
|
* populations bring the default set to thirty-six. The rates themselves are unchanged — a role's
|
|
* bit index only decides which counter a spike lands in — so a wider mask is capacity, not
|
|
* semantics.
|
|
*/
|
|
export const MAX_RATE_ROLES = 64;
|
|
|
|
/** 32-bit words per neuron in the role bitmask. */
|
|
export const ROLE_MASK_WORDS = MAX_RATE_ROLES / 32;
|
|
|
|
export interface LifConfig {
|
|
/** Membrane decay time constant in milliseconds; decay per tick is fround(exp(-1/decayMs)). */
|
|
decayMs: number;
|
|
/** Spike threshold in membrane units. */
|
|
threshold: number;
|
|
/** Ticks a neuron stays refractory after spiking. */
|
|
refractoryMs: number;
|
|
/** Multiplier applied to dataset weights when a spike propagates. */
|
|
synapseScale: number;
|
|
/** Upper bound of the per-neuron random baseline drive. */
|
|
baselineMax: number;
|
|
/** Random membrane kicks per tick. */
|
|
noiseKicks: number;
|
|
/** Membrane increment per noise kick. */
|
|
noiseAmount: number;
|
|
/** Exponential-moving-average coefficient for rate estimates. */
|
|
rateAlpha: number;
|
|
/** Lower clamp on the membrane after inhibitory input. */
|
|
membraneFloor: number;
|
|
/** Seed of the deterministic noise generator. */
|
|
seed: number;
|
|
/** Role driven by `stimulate()` (the reward pulse) and the drive it receives per tick. */
|
|
stimulation: { role: string; drive: number };
|
|
/**
|
|
* Roles whose population rate is tracked. Defaults to every `command_*` role plus the
|
|
* historical motor/reward roles, in dataset role order. Names absent from the dataset are
|
|
* dropped; at most {@link MAX_RATE_ROLES} roles may be tracked.
|
|
*/
|
|
rateRoles?: string[];
|
|
/** Retina gain and the frame size assumed by `setVisualFrame` when none is given. */
|
|
retina: RetinaConfig;
|
|
}
|
|
|
|
/** Original constants of the FAFB kernel. */
|
|
export const DEFAULT_LIF_CONFIG: LifConfig = {
|
|
decayMs: 20, threshold: 1, refractoryMs: 2, synapseScale: 0.005,
|
|
baselineMax: 0.06, noiseKicks: 300, noiseAmount: 0.42, rateAlpha: 1 / 25,
|
|
membraneFloor: -2, seed: 22_222,
|
|
stimulation: { role: 'reward_pam', drive: 0.20 },
|
|
retina: { ...DEFAULT_RETINA_CONFIG },
|
|
};
|
|
|
|
/** Numeric parameters of the kernel, in version-hash order. Role names never enter the hash. */
|
|
function versionParams(config: LifConfig): number[] {
|
|
return [
|
|
config.decayMs, config.threshold, config.refractoryMs, config.synapseScale,
|
|
config.baselineMax, config.noiseKicks, config.noiseAmount, config.rateAlpha,
|
|
config.membraneFloor, config.seed, config.stimulation.drive,
|
|
config.retina.gain, config.retina.width, config.retina.height,
|
|
];
|
|
}
|
|
|
|
/** `NEURAL_KERNEL_VERSION` for the default constants, otherwise `lif-1ms-f64-v2:<fnv1a32>`. */
|
|
export function kernelVersion(config: Partial<LifConfig> = {}): string {
|
|
const merged = { ...DEFAULT_LIF_CONFIG, ...config };
|
|
return versionFor(NEURAL_KERNEL_VERSION, `${NEURAL_KERNEL_VERSION}:`, versionParams(merged), versionParams(DEFAULT_LIF_CONFIG));
|
|
}
|
|
|
|
export interface LifState {
|
|
membrane: Float32Array;
|
|
refractory: Uint8Array;
|
|
lastSpikeMs: Float64Array;
|
|
visualDrive: Float32Array;
|
|
rng: number;
|
|
rewardRemaining: number;
|
|
ms: number;
|
|
populationRate: number;
|
|
rates: Record<string, number>;
|
|
plasticity: PlasticityState;
|
|
}
|
|
|
|
/**
|
|
* Leaky integrate-and-fire network over a connectome dataset, stepped in 1-ms ticks.
|
|
*
|
|
* Membrane state is Float32; spike and eligibility timestamps are Float64 milliseconds so that
|
|
* 1-ms differences survive long sessions. Noise is a deterministic xorshift stream, so a
|
|
* checkpoint plus the dataset reproduces a run exactly.
|
|
*/
|
|
export class LifNetwork {
|
|
readonly config: LifConfig;
|
|
/** Version string of this configuration; pin it in checkpoints. */
|
|
readonly version: string;
|
|
readonly plasticity: RewardModulatedStdp;
|
|
readonly membrane: Float32Array;
|
|
readonly refractory: Uint8Array;
|
|
readonly baseline: Float32Array;
|
|
readonly lastSpikeMs: Float64Array;
|
|
readonly rates: Record<string, number> = {};
|
|
readonly roleNames: string[];
|
|
/** Role membership per neuron, {@link ROLE_MASK_WORDS} words each, low word first. */
|
|
readonly roleMasks: Uint32Array;
|
|
private readonly decay: number;
|
|
private readonly stimulationTargets: number[];
|
|
private readonly spikes: Uint32Array;
|
|
private readonly roleCounts: Uint32Array;
|
|
private readonly rng: Xorshift32;
|
|
private visualDrive: Float32Array;
|
|
private rewardRemaining = 0;
|
|
ms = 0;
|
|
populationRate = 0;
|
|
|
|
constructor(readonly data: BrainDataset, config: Partial<LifConfig> & { plasticity?: Partial<PlasticityConfig> } = {}) {
|
|
const { plasticity, ...kernel } = config;
|
|
this.config = { ...DEFAULT_LIF_CONFIG, ...kernel };
|
|
this.version = kernelVersion(this.config);
|
|
this.decay = Math.fround(Math.exp(-1 / this.config.decayMs));
|
|
this.plasticity = new RewardModulatedStdp(data, plasticity);
|
|
const n = data.meta.neurons;
|
|
this.membrane = new Float32Array(n);
|
|
this.refractory = new Uint8Array(n);
|
|
this.baseline = new Float32Array(n);
|
|
this.visualDrive = new Float32Array(data.visualIndices.length);
|
|
this.lastSpikeMs = new Float64Array(n).fill(-1_000_000);
|
|
this.spikes = new Uint32Array(n);
|
|
this.stimulationTargets = data.meta.roles[this.config.stimulation.role] ?? [];
|
|
const requested = this.config.rateRoles;
|
|
this.roleNames = requested
|
|
? requested.filter((name) => name in data.meta.roles)
|
|
: Object.keys(data.meta.roles).filter(
|
|
(name) =>
|
|
name.startsWith('command_')
|
|
|| name.startsWith(MACRO_RATE_ROLE_PREFIX)
|
|
|| DEFAULT_EXTRA_RATE_ROLES.includes(name),
|
|
);
|
|
if (this.roleNames.length > MAX_RATE_ROLES) {
|
|
throw new Error(`Too many tracked rate roles (${this.roleNames.length}); at most ${MAX_RATE_ROLES} fit the role bitmask`);
|
|
}
|
|
this.roleMasks = new Uint32Array(n * ROLE_MASK_WORDS);
|
|
this.roleCounts = new Uint32Array(this.roleNames.length);
|
|
for (let role = 0; role < this.roleNames.length; role++) {
|
|
this.rates[this.roleNames[role]] = 0;
|
|
const word = role >>> 5;
|
|
const bit = 1 << (role & 31);
|
|
for (const neuron of data.meta.roles[this.roleNames[role]]) this.roleMasks[neuron * ROLE_MASK_WORDS + word] |= bit;
|
|
}
|
|
this.rng = new Xorshift32(this.config.seed);
|
|
for (let i = 0; i < n; i++) this.baseline[i] = this.rng.next() * this.config.baselineMax;
|
|
}
|
|
|
|
/** Project one RGBA frame onto the retina columns; the size defaults to `config.retina`. */
|
|
setVisualFrame(rgba: Uint8Array, width = this.config.retina.width, height = this.config.retina.height): void {
|
|
const count = this.data.visualIndices.length;
|
|
if (this.visualDrive.length !== count) this.visualDrive = new Float32Array(count);
|
|
projectFrame(
|
|
rgba, width, height,
|
|
{ xy: this.data.visualXY, hemisphere: this.data.visualHemisphere, count },
|
|
this.config.retina.gain,
|
|
this.visualDrive,
|
|
);
|
|
}
|
|
|
|
/** Drive the stimulation role for `durationMs` ticks; overlapping pulses take the maximum. */
|
|
stimulate(durationMs = 120): void {
|
|
this.rewardRemaining = Math.max(this.rewardRemaining, durationMs);
|
|
}
|
|
|
|
/** Historical name of {@link LifNetwork.stimulate}. */
|
|
reward(durationMs = 120): void {
|
|
this.stimulate(durationMs);
|
|
}
|
|
|
|
step(milliseconds: number): number {
|
|
let totalSpikes = 0;
|
|
for (let tick = 0; tick < milliseconds; tick++) totalSpikes += this.stepOne();
|
|
return totalSpikes;
|
|
}
|
|
|
|
private stepOne(): number {
|
|
const { indptr, targets, weights, visualIndices, meta } = this.data;
|
|
const { threshold, refractoryMs, synapseScale, noiseKicks, noiseAmount, rateAlpha, membraneFloor } = this.config;
|
|
const decay = this.decay;
|
|
const n = meta.neurons;
|
|
for (let i = 0; i < noiseKicks; i++) this.membrane[this.rng.nextUint() % n] += noiseAmount;
|
|
for (let i = 0; i < visualIndices.length; i++) this.membrane[visualIndices[i]] += this.visualDrive[i];
|
|
if (this.rewardRemaining > 0) {
|
|
for (const neuron of this.stimulationTargets) this.membrane[neuron] += this.config.stimulation.drive;
|
|
this.rewardRemaining--;
|
|
}
|
|
|
|
let spikeCount = 0;
|
|
for (let neuron = 0; neuron < n; neuron++) {
|
|
if (this.refractory[neuron] > 0) {
|
|
this.refractory[neuron]--;
|
|
this.membrane[neuron] *= decay;
|
|
continue;
|
|
}
|
|
const voltage = this.membrane[neuron] * decay + this.baseline[neuron];
|
|
if (voltage >= threshold) {
|
|
this.membrane[neuron] = 0;
|
|
this.refractory[neuron] = refractoryMs;
|
|
this.spikes[spikeCount++] = neuron;
|
|
} else {
|
|
this.membrane[neuron] = voltage;
|
|
}
|
|
}
|
|
|
|
this.roleCounts.fill(0);
|
|
this.plasticity.observe(this.spikes, spikeCount, this.lastSpikeMs, this.ms);
|
|
for (let spike = 0; spike < spikeCount; spike++) {
|
|
const source = this.spikes[spike];
|
|
this.lastSpikeMs[source] = this.ms;
|
|
for (let word = 0; word < ROLE_MASK_WORDS; word++) {
|
|
let mask = this.roleMasks[source * ROLE_MASK_WORDS + word];
|
|
for (let role = word * 32; mask; role++, mask >>>= 1) if (mask & 1) this.roleCounts[role]++;
|
|
}
|
|
for (let edge = indptr[source]; edge < indptr[source + 1]; edge++) {
|
|
const target = targets[edge];
|
|
this.membrane[target] = Math.max(membraneFloor, this.membrane[target] + weights[edge] * this.plasticity.gain(edge) * synapseScale);
|
|
}
|
|
}
|
|
for (let role = 0; role < this.roleNames.length; role++) {
|
|
const name = this.roleNames[role];
|
|
const size = meta.roles[name].length;
|
|
const instantaneous = size ? this.roleCounts[role] * 1000 / size : 0;
|
|
this.rates[name] += (instantaneous - this.rates[name]) * rateAlpha;
|
|
}
|
|
this.populationRate += (spikeCount * 1000 / n - this.populationRate) * rateAlpha;
|
|
this.ms++;
|
|
return spikeCount;
|
|
}
|
|
|
|
exportState(): LifState {
|
|
return {
|
|
membrane: this.membrane.slice(),
|
|
refractory: this.refractory.slice(),
|
|
lastSpikeMs: this.lastSpikeMs.slice(),
|
|
visualDrive: this.visualDrive.slice(),
|
|
rng: this.rng.state,
|
|
rewardRemaining: this.rewardRemaining,
|
|
ms: this.ms,
|
|
populationRate: this.populationRate,
|
|
rates: { ...this.rates },
|
|
plasticity: this.plasticity.exportState(),
|
|
};
|
|
}
|
|
|
|
importState(state: Omit<LifState, 'plasticity'> & { plasticity?: PlasticityState }): void {
|
|
if (state.membrane.length !== this.membrane.length || state.refractory.length !== this.refractory.length || state.lastSpikeMs.length !== this.lastSpikeMs.length) {
|
|
throw new Error('Brain checkpoint dimensions do not match the loaded dataset');
|
|
}
|
|
if (state.visualDrive.length !== this.data.visualIndices.length || !Number.isSafeInteger(state.ms) || state.ms < 0 || !Number.isInteger(state.rng) || !Number.isFinite(state.populationRate) || !Number.isFinite(state.rewardRemaining) || state.rewardRemaining < 0 || Object.values(state.rates).some(value => !Number.isFinite(value)) || [state.membrane, state.lastSpikeMs, state.visualDrive].some(values => values.some(value => !Number.isFinite(value)))) throw new Error('Invalid neural checkpoint values');
|
|
this.plasticity.importState(state.plasticity);
|
|
this.membrane.set(state.membrane);
|
|
this.refractory.set(state.refractory);
|
|
this.lastSpikeMs.set(state.lastSpikeMs);
|
|
this.visualDrive = state.visualDrive.slice();
|
|
this.rng.state = state.rng;
|
|
this.rewardRemaining = state.rewardRemaining;
|
|
this.ms = state.ms;
|
|
this.populationRate = state.populationRate;
|
|
for (const name of this.roleNames) this.rates[name] = state.rates[name] ?? 0;
|
|
}
|
|
}
|
|
|
|
/** Historical name of {@link LifNetwork}. */
|
|
export { LifNetwork as FlyBrain };
|