diff options
Diffstat (limited to 'sdil/bci.py')
| -rw-r--r-- | sdil/bci.py | 302 |
1 files changed, 302 insertions, 0 deletions
diff --git a/sdil/bci.py b/sdil/bci.py new file mode 100644 index 0000000..9847ffe --- /dev/null +++ b/sdil/bci.py @@ -0,0 +1,302 @@ +"""Continuous BCI dynamics for the preregistered Harnett-signature study. + +Learning is manual and local. Autograd is never enabled. The environment's +scalar loss is observed only through antithetic neural perturbations; the +resulting causal targets calibrate a cell-specific apical vectorizer. +""" +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class BCIConfig: + n_plus: int = 5 + n_minus: int = 5 + n_background: int = 30 + context_dim: int = 16 + steps_per_episode: int = 28 + episodes_per_day: int = 64 + days: int = 14 + target: float = 0.8 + inertia: float = 0.65 + process_noise: float = 0.12 + context_ar: float = 0.8 + coupling_scale: float = 1.0 + predictor_eta: float = 0.2 + vectorizer_eta: float = 0.03 + forward_eta: float = 0.01 + perturb_sigma: float = 0.03 + perturb_every: int = 4 + kappa: float = 0.0 + feedback: str = "error" + + @property + def n_neurons(self): + return self.n_plus + self.n_minus + self.n_background + + @property + def feedback_dim(self): + return 1 if self.feedback == "error" else 2 + + def validate(self): + if self.feedback not in ("error", "error_velocity"): + raise ValueError(f"unknown BCI feedback: {self.feedback}") + if self.n_plus < 1 or self.n_minus < 1 or self.n_background < 1: + raise ValueError("all BCI populations must be nonempty") + if self.context_dim < 2: + raise ValueError("context_dim includes a bias and must be at least two") + if self.steps_per_episode < 2 or self.episodes_per_day < 2 or self.days < 2: + raise ValueError("BCI trajectories require multiple steps, episodes, and days") + if not 0 <= self.context_ar < 1: + raise ValueError("context_ar must lie in [0,1)") + if self.perturb_every < 1 or self.perturb_sigma <= 0: + raise ValueError("perturbation schedule must be positive") + + +@dataclass +class BCITrajectories: + context: torch.Tensor + process_noise: torch.Tensor + perturbations: torch.Tensor + task_seed: int + + @property + def shape(self): + return self.context.shape[:3] + + +def causal_roles(cfg, *, device="cpu", dtype=torch.float32): + """Experimenter-defined P+/P-/P0 contribution to the scalar cursor.""" + role = torch.zeros(cfg.n_neurons, device=device, dtype=dtype) + role[:cfg.n_plus] = 1.0 / cfg.n_plus + role[cfg.n_plus:cfg.n_plus + cfg.n_minus] = -1.0 / cfg.n_minus + return role + + +def generate_trajectories(cfg, task_seed, *, days=None, episodes=None, + device="cpu", dtype=torch.float32): + """Generate paired exogenous context, process noise, and causal probes. + + Model variants receive this same object, which makes all environmental and + perturbation randomness exactly paired. The first context component is a + constant bias; the others follow independent AR(1) dynamics within each + episode and reset at episode boundaries. + """ + cfg.validate() + days = cfg.days if days is None else days + episodes = cfg.episodes_per_day if episodes is None else episodes + generator = torch.Generator(device="cpu").manual_seed(task_seed) + shape = (days, episodes, cfg.steps_per_episode) + innovations = torch.randn( + *shape, cfg.context_dim - 1, generator=generator, dtype=dtype) + context = torch.empty(*shape, cfg.context_dim, dtype=dtype) + context[..., 0] = 1.0 + state = torch.zeros(days, episodes, cfg.context_dim - 1, dtype=dtype) + innovation_scale = (1.0 - cfg.context_ar ** 2) ** 0.5 + for step in range(cfg.steps_per_episode): + state = cfg.context_ar * state + innovation_scale * innovations[:, :, step] + context[:, :, step, 1:] = state + noise = cfg.process_noise * torch.randn( + *shape, cfg.n_neurons, generator=generator, dtype=dtype) + perturbations = torch.empty( + *shape, cfg.n_neurons, dtype=dtype).bernoulli_(0.5, generator=generator) + perturbations.mul_(2).sub_(1) + return BCITrajectories( + context=context.to(device), process_noise=noise.to(device), + perturbations=perturbations.to(device), task_seed=task_seed) + + +class BCISDIL: + """Manual local learner for the continuous synthetic BCI.""" + + def __init__(self, cfg, model_seed=0, *, device="cpu", dtype=torch.float32): + cfg.validate() + self.cfg = cfg + self.device = device + self.dtype = dtype + self.role = causal_roles(cfg, device=device, dtype=dtype) + generator = torch.Generator(device="cpu").manual_seed(model_seed) + self.W = (0.15 * torch.randn( + cfg.n_neurons, cfg.context_dim, generator=generator, dtype=dtype) + / cfg.context_dim ** 0.5).to(device) + self.A = (0.05 * torch.randn( + cfg.n_neurons, cfg.feedback_dim, generator=generator, dtype=dtype) + / cfg.feedback_dim ** 0.5).to(device) + self.coupling = (cfg.coupling_scale * torch.exp( + 0.2 * torch.randn(cfg.n_neurons, generator=generator, dtype=dtype))).to(device) + self.P = torch.zeros(cfg.n_neurons, device=device, dtype=dtype) + self.P_bias = torch.zeros(cfg.n_neurons, device=device, dtype=dtype) + self.model_seed = model_seed + + def clone(self): + other = BCISDIL( + self.cfg, self.model_seed, device=self.device, dtype=self.dtype) + for name in ("W", "A", "coupling", "P", "P_bias"): + setattr(other, name, getattr(self, name).clone()) + return other + + def feedback_features(self, error, previous_abs_error): + if self.cfg.feedback == "error": + return error.unsqueeze(1) + improvement_velocity = (torch.zeros_like(error) if previous_abs_error is None + else previous_abs_error - error.abs()) + return torch.stack((error, improvement_velocity), dim=1) + + def apical_components(self, soma, feedback): + ordinary = self.coupling * soma + instructional = feedback @ self.A.t() + raw = ordinary + instructional + baseline = self.P * soma + self.P_bias + return raw, raw - baseline, ordinary, instructional + + def causal_targets(self, soma, target, perturbations): + """One simultaneous antithetic forward-only causal target per cell.""" + sigma = self.cfg.perturb_sigma + plus = soma + sigma * perturbations + minus = soma - sigma * perturbations + error_plus = target - plus @ self.role + error_minus = target - minus @ self.role + loss_plus = 0.5 * error_plus.square() + loss_minus = 0.5 * error_minus.square() + directional_derivative = (loss_plus - loss_minus) / (2.0 * sigma) + return -directional_derivative.unsqueeze(1) * perturbations + + def exact_causal_direction(self, soma, target): + """Analytic diagnostic only; never used by a learning update.""" + error = target - soma @ self.role + return error.unsqueeze(1) * self.role.unsqueeze(0) + + def update_predictor(self, soma, ordinary, active): + """Per-cell centered normalized-LMS fit of neutral apical traffic.""" + if not bool(active.any()): + return + selected_soma = soma[active] + selected_target = ordinary[active] + residual = selected_target - ( + self.P * selected_soma + self.P_bias) + soma_centered = selected_soma - selected_soma.mean(0) + residual_centered = residual - residual.mean(0) + variance = soma_centered.square().mean(0) + self.P += self.cfg.predictor_eta * ( + (residual_centered * soma_centered).mean(0) / (variance + 1e-6)) + self.P_bias += self.cfg.predictor_eta * residual.mean(0) + + def update_vectorizer(self, feedback, innovation, causal_target, active): + if not bool(active.any()): + return + c = feedback[active] + calibration_error = causal_target[active] - innovation[active] + self.A += self.cfg.vectorizer_eta * ( + calibration_error.t() @ c / c.shape[0]) + + def update_forward(self, context, soma, innovation, active, plasticity_gain): + if plasticity_gain == 0 or not bool(active.any()): + return + eligibility_gain = 1.0 - soma.square() + local_delta = plasticity_gain * innovation * eligibility_gain + self.W += self.cfg.forward_eta * ( + local_delta[active].t() @ context[active] / active.sum()) + + def step(self, context, process_noise, perturbations, previous_soma, + previous_abs_error, active, global_step, *, control_gain=1.0, + plasticity_gain=1.0, learn_vectorizer=True, + learn_predictor=True): + """Advance one contiguous BCI step and optionally apply local updates. + + ``control_gain`` affects state but not the plasticity signal. + ``plasticity_gain`` affects W but not online state. This separation is + the preregistered phase lesion. All tensors are minibatches of episodes. + """ + with torch.no_grad(): + u = context @ self.W.t() + self.cfg.inertia * previous_soma + process_noise + base_soma = torch.tanh(u) + base_cursor = base_soma @ self.role + base_error = self.cfg.target - base_cursor + feedback = self.feedback_features(base_error, previous_abs_error) + raw, innovation, ordinary, instructional = self.apical_components( + base_soma, feedback) + innovation = innovation * active.unsqueeze(1) + controlled = torch.clamp( + base_soma + control_gain * self.cfg.kappa * innovation, -1.0, 1.0) + soma = torch.where(active.unsqueeze(1), controlled, previous_soma) + cursor = soma @ self.role + error = self.cfg.target - cursor + loss = 0.5 * error.square() + + did_perturb = global_step % self.cfg.perturb_every == 0 + causal_target = None + if did_perturb: + causal_target = self.causal_targets( + base_soma, self.cfg.target, perturbations) + + # All updates see the same pre-update state. The order below only + # mutates parameters after every required local quantity is stored. + self.update_forward( + context, base_soma, innovation, active, plasticity_gain) + if did_perturb and learn_vectorizer: + self.update_vectorizer( + feedback, innovation, causal_target, active) + if learn_predictor: + self.update_predictor(base_soma, ordinary, active) + + return { + "soma": soma, + "base_soma": base_soma, + "cursor": cursor, + "error": error, + "abs_error": error.abs(), + "loss": loss, + "feedback": feedback, + "raw_apical": raw, + "innovation": innovation, + "ordinary_apical": ordinary, + "instructional_apical": instructional, + "causal_target": causal_target, + "did_perturb": did_perturb, + } + + +def run_day(model, trajectories, day, *, control_gain=1.0, + plasticity_gain=1.0, learn_vectorizer=True, + learn_predictor=True, collect=False, global_step=0): + """Run one parallel batch of episodes and return success plus event data.""" + cfg = model.cfg + episodes = trajectories.context.shape[1] + soma = torch.zeros( + episodes, cfg.n_neurons, device=model.device, dtype=model.dtype) + previous_abs_error = None + active = torch.ones(episodes, device=model.device, dtype=torch.bool) + success = torch.zeros_like(active) + events = [] + for step_index in range(cfg.steps_per_episode): + record = model.step( + trajectories.context[day, :, step_index], + trajectories.process_noise[day, :, step_index], + trajectories.perturbations[day, :, step_index], + soma, previous_abs_error, active, global_step, + control_gain=control_gain, plasticity_gain=plasticity_gain, + learn_vectorizer=learn_vectorizer, + learn_predictor=learn_predictor) + newly_successful = active & (record["cursor"] >= cfg.target) + success |= newly_successful + if collect: + events.append({ + key: value.detach().cpu().clone() + for key, value in record.items() + if isinstance(value, torch.Tensor) + } | { + "active": active.detach().cpu().clone(), + "success": success.detach().cpu().clone(), + "step": step_index, + }) + soma = record["soma"] + previous_abs_error = record["abs_error"] + active = active & ~newly_successful + global_step += 1 + return { + "success": success.detach().cpu(), + "success_rate": success.float().mean().item(), + "events": events, + "global_step": global_step, + } |