File size: 14,448 Bytes
b88c26d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 | #!/usr/bin/env python3
"""
Train Boolean actor network using MMEP (Min-Max Equilibrium Propagation).
License: BSL-1.1
Contact: jessica@collectivekitty.com
MMEP training loop:
1. Free phase: run network forward to equilibrium (energy minimization)
2. Nudge phase: clamp output toward target, re-equilibrate with nudge factor beta
3. Local update: compute weight updates from difference in activations between phases
4. Huntington regularization: penalize violations of Boolean algebra postulates
The Boolean perceptron actors enforce algebraic structure while MMEP provides
biologically plausible gradient-free learning.
"""
import argparse
import math
import sys
import time
from pathlib import Path
from typing import Dict, Optional
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import yaml
class BooleanPerceptronActor(nn.Module):
"""Single Boolean perceptron actor with Huntington-compatible operations."""
def __init__(self, input_dim: int, output_dim: int):
super().__init__()
self.input_dim = input_dim
self.output_dim = output_dim
self.weight = nn.Parameter(torch.randn(output_dim, input_dim) * 0.01)
self.bias = nn.Parameter(torch.zeros(output_dim))
# Boolean operation parameters (learnable)
self.or_weight = nn.Parameter(torch.randn(output_dim, 2 * output_dim) * 0.01)
self.and_weight = nn.Parameter(torch.randn(output_dim, 2 * output_dim) * 0.01)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass through sigmoid activation (Boolean-compatible)."""
return torch.sigmoid(F.linear(x, self.weight, self.bias))
def boolean_or(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""Approximate Boolean OR via learned gate."""
combined = torch.cat([a, b], dim=-1)
return torch.sigmoid(F.linear(combined, self.or_weight))
def boolean_and(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""Approximate Boolean AND via learned gate."""
combined = torch.cat([a, b], dim=-1)
return torch.sigmoid(F.linear(combined, self.and_weight))
class BooleanActorNetwork(nn.Module):
"""Network of Boolean perceptron actors trained with MMEP."""
def __init__(self, layer_dims: list, beta: float = 0.1):
super().__init__()
self.beta = beta
self.layers = nn.ModuleList()
for i in range(len(layer_dims) - 1):
self.layers.append(BooleanPerceptronActor(layer_dims[i], layer_dims[i + 1]))
self.layer_dims = layer_dims
def free_phase(self, x: torch.Tensor, num_steps: int = 10) -> list:
"""Run free phase to equilibrium. Returns activations at each layer."""
activations = [x]
h = x
for layer in self.layers:
h = layer(h)
activations.append(h)
# Iterate to equilibrium
for _ in range(num_steps - 1):
h = activations[0]
for i, layer in enumerate(self.layers):
h = layer(h)
activations[i + 1] = h
return activations
def nudge_phase(self, x: torch.Tensor, target: torch.Tensor, num_steps: int = 10) -> list:
"""Run nudge phase with target clamping. Returns activations at each layer."""
activations = [x]
h = x
for layer in self.layers:
h = layer(h)
activations.append(h)
# Nudge final layer toward target
for _ in range(num_steps - 1):
# Clamp output toward target
activations[-1] = (1 - self.beta) * activations[-1] + self.beta * target
h = activations[0]
for i, layer in enumerate(self.layers):
h = layer(h)
if i < len(self.layers) - 1:
activations[i + 1] = h
else:
activations[i + 1] = (1 - self.beta) * h + self.beta * target
return activations
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Standard forward pass for inference."""
h = x
for layer in self.layers:
h = layer(h)
return h
def huntington_regularization(network: BooleanActorNetwork, x: torch.Tensor) -> torch.Tensor:
"""Compute Huntington postulate violation as regularization loss.
Penalizes violations of:
- Commutativity: OR(a,b) = OR(b,a), AND(a,b) = AND(b,a)
- Idempotence: OR(a,a) = a, AND(a,a) = a
- Complement: OR(a, 1-a) ~ 1, AND(a, 1-a) ~ 0
"""
loss = torch.tensor(0.0, device=x.device)
count = 0
for layer in network.layers:
# Generate test activations
a = torch.sigmoid(torch.randn(x.shape[0], layer.output_dim, device=x.device))
b = torch.sigmoid(torch.randn(x.shape[0], layer.output_dim, device=x.device))
# Commutativity: OR(a,b) = OR(b,a)
or_ab = layer.boolean_or(a, b)
or_ba = layer.boolean_or(b, a)
loss = loss + F.mse_loss(or_ab, or_ba)
# Commutativity: AND(a,b) = AND(b,a)
and_ab = layer.boolean_and(a, b)
and_ba = layer.boolean_and(b, a)
loss = loss + F.mse_loss(and_ab, and_ba)
# Idempotence: OR(a,a) = a
or_aa = layer.boolean_or(a, a)
loss = loss + F.mse_loss(or_aa, a)
# Idempotence: AND(a,a) = a
and_aa = layer.boolean_and(a, a)
loss = loss + F.mse_loss(and_aa, a)
# Complement: OR(a, 1-a) ~ 1
comp_a = 1.0 - a
or_comp = layer.boolean_or(a, comp_a)
loss = loss + F.mse_loss(or_comp, torch.ones_like(or_comp))
# Complement: AND(a, 1-a) ~ 0
and_comp = layer.boolean_and(a, comp_a)
loss = loss + F.mse_loss(and_comp, torch.zeros_like(and_comp))
count += 6
return loss / max(count, 1)
def compute_mmep_update(
free_activations: list,
nudge_activations: list,
network: BooleanActorNetwork,
lr: float,
):
"""Compute and apply MMEP local weight updates.
The update rule is:
dW_i = (1/beta) * (s_i^nudge * s_{i-1}^nudge^T - s_i^free * s_{i-1}^free^T)
"""
beta = network.beta
for i, layer in enumerate(network.layers):
# Pre-synaptic and post-synaptic activations
pre_free = free_activations[i]
post_free = free_activations[i + 1]
pre_nudge = nudge_activations[i]
post_nudge = nudge_activations[i + 1]
# Compute correlation differences
# dW = (1/beta) * (post_nudge @ pre_nudge^T - post_free @ pre_free^T) / batch_size
batch_size = pre_free.shape[0]
corr_nudge = torch.mm(post_nudge.t(), pre_nudge) / batch_size
corr_free = torch.mm(post_free.t(), pre_free) / batch_size
dW = (1.0 / beta) * (corr_nudge - corr_free)
db = (1.0 / beta) * (post_nudge - post_free).mean(dim=0)
# Apply updates
with torch.no_grad():
layer.weight.add_(lr * dW)
layer.bias.add_(lr * db)
def load_config(config_path: str) -> dict:
"""Load configuration from YAML file."""
path = Path(config_path)
if not path.exists():
print(f"Warning: Config file {config_path} not found. Using defaults.")
return {}
with open(path, "r") as f:
return yaml.safe_load(f) or {}
def generate_boolean_data(batch_size: int, input_dim: int, device: torch.device):
"""Generate synthetic Boolean data for training."""
x = torch.randint(0, 2, (batch_size, input_dim), dtype=torch.float32, device=device)
# Target: XOR of first half and second half (non-trivial Boolean function)
half = input_dim // 2
target_bits = ((x[:, :half] + x[:, half:2*half]) % 2)
# Pad or truncate to match output dim
return x, target_bits
def main():
parser = argparse.ArgumentParser(
description="Train Boolean actor network with MMEP (Min-Max Equilibrium Propagation)."
)
parser.add_argument("--config", type=str, default="config/boolean_actor.yaml",
help="Path to config YAML (default: config/boolean_actor.yaml)")
parser.add_argument("--epochs", type=int, default=None, help="Number of epochs (overrides config)")
parser.add_argument("--device", type=str, default=None, help="Device (overrides config)")
args = parser.parse_args()
# Load config
config = load_config(args.config)
# Extract parameters with defaults
model_cfg = config.get("model", {})
train_cfg = config.get("training", {})
mmep_cfg = config.get("mmep", {})
input_dim = model_cfg.get("input_dim", 256)
hidden_dims = model_cfg.get("hidden_dims", [512, 256, 128])
output_dim = model_cfg.get("output_dim", 128)
layer_dims = [input_dim] + hidden_dims + [output_dim]
epochs = args.epochs or train_cfg.get("epochs", 100)
batch_size = train_cfg.get("batch_size", 256)
lr = train_cfg.get("learning_rate", 0.01)
huntington_weight = train_cfg.get("huntington_weight", 0.1)
log_interval = train_cfg.get("log_interval", 10)
beta = mmep_cfg.get("beta", 0.1)
free_steps = mmep_cfg.get("free_steps", 10)
nudge_steps = mmep_cfg.get("nudge_steps", 10)
device_str = args.device or train_cfg.get("device", "cpu")
device = torch.device(device_str)
print("BURT-IMMA Boolean MMEP Training")
print("=" * 60)
print(f" Config: {args.config}")
print(f" Layer dims: {layer_dims}")
print(f" Epochs: {epochs}")
print(f" Batch size: {batch_size}")
print(f" Learning rate: {lr}")
print(f" Beta (nudge): {beta}")
print(f" Free steps: {free_steps}")
print(f" Nudge steps: {nudge_steps}")
print(f" Huntington wt: {huntington_weight}")
print(f" Device: {device}")
print("=" * 60)
print()
# Create network
network = BooleanActorNetwork(layer_dims, beta=beta).to(device)
optimizer = torch.optim.Adam(network.parameters(), lr=lr)
# Training loop
metrics_history = []
start_time = time.time()
for epoch in range(1, epochs + 1):
epoch_loss = 0.0
epoch_accuracy = 0.0
epoch_huntington = 0.0
epoch_entropy = 0.0
num_batches = 0
# Generate training data
steps_per_epoch = max(1, 1000 // batch_size)
for step in range(steps_per_epoch):
x, target = generate_boolean_data(batch_size, input_dim, device)
# Ensure target matches output dim
if target.shape[1] != output_dim:
# Tile or truncate
if target.shape[1] < output_dim:
repeats = output_dim // target.shape[1] + 1
target = target.repeat(1, repeats)[:, :output_dim]
else:
target = target[:, :output_dim]
# === MMEP Training ===
# Phase 1: Free phase
free_acts = network.free_phase(x, num_steps=free_steps)
# Phase 2: Nudge phase
nudge_acts = network.nudge_phase(x, target, num_steps=nudge_steps)
# Phase 3: Local update (MMEP rule)
compute_mmep_update(free_acts, nudge_acts, network, lr)
# Compute task loss for logging
with torch.no_grad():
output = network(x)
task_loss = F.mse_loss(output, target)
# Phase 4: Huntington regularization (uses gradient)
optimizer.zero_grad()
h_loss = huntington_regularization(network, x)
reg_loss = huntington_weight * h_loss
reg_loss.backward()
optimizer.step()
# Compute metrics
with torch.no_grad():
output = network(x)
predictions = (output > 0.5).float()
accuracy = (predictions == target).float().mean().item()
# Binary entropy of outputs
p = torch.clamp(output, 1e-7, 1 - 1e-7)
entropy = -(p * torch.log(p) + (1 - p) * torch.log(1 - p)).mean().item()
epoch_loss += task_loss.item()
epoch_accuracy += accuracy
epoch_huntington += h_loss.item()
epoch_entropy += entropy
num_batches += 1
# Average metrics
avg_loss = epoch_loss / num_batches
avg_accuracy = epoch_accuracy / num_batches
avg_huntington = epoch_huntington / num_batches
avg_entropy = epoch_entropy / num_batches
metrics = {
"epoch": epoch,
"loss": avg_loss,
"accuracy": avg_accuracy,
"huntington_violation": avg_huntington,
"entropy": avg_entropy,
}
metrics_history.append(metrics)
if epoch % log_interval == 0 or epoch == 1:
elapsed = time.time() - start_time
print(
f"Epoch {epoch:4d}/{epochs} | "
f"loss: {avg_loss:.6f} | "
f"acc: {avg_accuracy:.4f} | "
f"hunt_viol: {avg_huntington:.6f} | "
f"entropy: {avg_entropy:.4f} | "
f"time: {elapsed:.1f}s"
)
# Final summary
elapsed = time.time() - start_time
print()
print("=" * 60)
print("Training Complete")
print("=" * 60)
print(f" Total time: {elapsed:.1f}s")
print(f" Final loss: {metrics_history[-1]['loss']:.6f}")
print(f" Final accuracy: {metrics_history[-1]['accuracy']:.4f}")
print(f" Final Huntington viol.: {metrics_history[-1]['huntington_violation']:.6f}")
print(f" Final entropy: {metrics_history[-1]['entropy']:.4f}")
# Save model
save_path = Path("checkpoints/boolean_mmep_latest.pt")
save_path.parent.mkdir(parents=True, exist_ok=True)
torch.save({
"model_state_dict": network.state_dict(),
"config": config,
"metrics": metrics_history,
"layer_dims": layer_dims,
}, save_path)
print(f" Model saved to: {save_path}")
if __name__ == "__main__":
main()
|