Model card: inference and training examples (both executed), five-seed result; fused connectome-kernels path in modeling code

README.md

@@ -81,7 +81,7 @@ h_i ← (1 − leak_i) h_i + leak_i · proposal_i        (2 microsteps per chara
 512 output neuron states → linear → 65 logits
 ```
 
-## Usage
+## Inference
 
 ```python
 import torch
@@ -93,9 +93,83 @@ model = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", trust_remote_code
 ids = tok("ROMEO:", return_tensors="pt").input_ids
 out = model.generate(ids, max_new_tokens=300, do_sample=True, temperature=0.8)
 print(tok.decode(out[0]))
+
+# The degree-preserving scrambled control (same neurons, same degrees, shuffled wiring), for comparison:
+scrambled = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", subfolder="scrambled", trust_remote_code=True, dtype=torch.float32)
+print(tok.decode(scrambled.generate(ids, max_new_tokens=300, do_sample=True, temperature=0.8)[0]))
+
+# Neuron activity, for visualization: [1, T, 5000] states after each character, plus MaleCNS body ids
+with torch.no_grad():
+    states = model(ids).state                                   # [B, N] after the last character
+body_ids = model.graph.node_id                                  # index -> MaleCNS body id, for lookup in QuixiAI/MaleCNS
+```
+
+The tokenizer is strict: only the 65 characters of Tiny Shakespeare are encodable. `generate()` carries the neuron
+state between characters instead of a KV cache.
+
+## Training
+
+The recurrent core has one trainable weight per real synaptic connection. With the
+[connectome-kernels](https://github.com/QuixiAI/connectome-kernels) package installed, the model's forward pass
+runs on fused CUDA kernels (about 13× faster than `torch.sparse`, identical gradients); without it, it falls back
+to `torch.sparse` automatically.
+
+```python
+# Fine-tune / continue training FlyGPT on Tiny Shakespeare (character-level).
+# pip install transformers safetensors
+# pip install --no-build-isolation git+https://github.com/QuixiAI/connectome-kernels   # fused CUDA path, ~13x faster
+import requests, torch, torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tok = AutoTokenizer.from_pretrained("QuixiAI/FlyGPT")
+model = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", trust_remote_code=True, dtype=torch.float32).cuda()
+# start from the untrained initialization instead:  subfolder="init"
+
+text = requests.get("https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt").text
+data = torch.tensor(tok(text).input_ids)
+train, val = data[: int(0.9 * len(data))], data[int(0.9 * len(data)):]   # FlyGPT's fixed 90/10 split
+
+def batch(split, B=32, T=64):
+    i = torch.randint(0, len(split) - T - 1, (B,))
+    x = torch.stack([split[j : j + T] for j in i]); y = torch.stack([split[j + 1 : j + T + 1] for j in i])
+    return x.cuda(), y.cuda()
+
+recurrent = list(model.recurrent.parameters())                         # one weight per real synapse, bias, leak
+adapters = [p for n, p in model.named_parameters() if not n.startswith("recurrent.")]
+opt = torch.optim.AdamW([{"params": adapters, "lr": 1e-3}, {"params": recurrent, "lr": 3e-4}], weight_decay=0.01)
+
+for step in range(1, 501):
+    x, y = batch(train)
+    logits = model(x).logits                                           # [B, T, 65]; state resets to zero per window
+    loss = F.cross_entropy(logits.reshape(-1, 65), y.reshape(-1))
+    opt.zero_grad(set_to_none=True); loss.backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
+    if step % 100 == 0:
+        with torch.no_grad():
+            vx, vy = batch(val); vl = F.cross_entropy(model(vx).logits.reshape(-1, 65), vy.reshape(-1))
+        print(f"step {step}  train {loss.item():.3f}  val {vl.item():.3f}")
+
+model.save_pretrained("flygpt-finetuned"); tok.save_pretrained("flygpt-finetuned")
 ```
 
-The tokenizer is strict: only the 65 characters of Tiny Shakespeare are encodable.
+## Result: does the wiring matter?
+
+Pre-registered rule (before any result was seen): "wiring matters" is claimed only if all 5 paired
+differences Δ = loss(scrambled) − loss(real) have the same sign **and** the mean Δ is at least 0.05 nats/char.
+
+| seed | real | degree-preserving scramble | Δ |
+|--:|--:|--:|--:|
+| 1 | 1.7275 | 1.7450 | +0.0175 |
+| 2 | 1.7285 | 1.7410 | +0.0125 |
+| 3 | 1.7490 | 1.7594 | +0.0104 |
+| 4 | 1.7516 | 1.7592 | +0.0076 |
+| 5 | 1.7305 | 1.7377 | +0.0073 |
+| **mean** | **1.7374** | **1.7485** | **+0.0111** |
+
+Validation loss in nats/char at the end of 20,000 steps, same data order, batches, adapter init and edge-value RNG
+stream per seed. Bigram reference on this split: 2.482. All five differences favour the real wiring, but the mean
+gap is below the pre-registered minimum effect, so the verdict is: **no detectable difference at this scale.** The fly connectome learns
+Shakespeare; whether its specific wiring helps, beyond its degree sequence, is not resolved at 5,000 neurons.
 
 ## Citation
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c1c0f9c6fc00d98d924fc273c9ea4b0ea0247701424240483ace1f344fbf59fe
+oid sha256:d8c8b47fc37d85fb591f215a3cc54aa3d045e0edfb69a9d6116c1b1f87bbd317
 size 7495610

modeling_flygpt.py

@@ -7,8 +7,10 @@ Dynamics (one scalar state per neuron, plan.md §7 of the FlyGPT spec):
     h_i_new    = (1 - leak_i) * h_i + leak_i * proposal_i
 
 The connectome is stored in `model.safetensors` as integer tensors (`graph.*`); only the learned
-per-edge values and the adapters are floating point (bf16 on disk). The sparse recurrent matmul is
-rebuilt in fp32 at runtime (rows = destination, columns = source).
+per-edge values and the adapters are floating point (bf16 on disk). The sparse recurrent matmul runs
+in fp32: through the fused kernels of the `connectome-kernels` package when it is installed and a
+CUDA device is used (training speed), else through torch.sparse COO (rows = destination). Both give
+the same logits and gradients.
 """
 from __future__ import annotations
 
@@ -133,17 +135,50 @@ class FlyGPTForCausalLM(FlyGPTPreTrainedModel, GenerationMixin):
     def logits_from_state(self, state: torch.Tensor) -> torch.Tensor:
         return self.lm_head(state[:, self.graph.output_nodes].to(self.lm_head.weight.dtype)).float()
 
+    # ---- fused CUDA path via the connectome-kernels package (optional, used for training) ------------
+    def _fused_graph(self):
+        from connectome_kernels import SparseGraph
+        dev = self.graph.edge_index.device
+        if getattr(self, "_fg", None) is None or self._fg_device != dev:
+            self._fg = SparseGraph(self.graph.edge_index[0].long(), self.graph.edge_index[1].long(), self.num_neurons,
+                                   self.graph.input_nodes)
+            self._fg_device = dev
+        return self._fg
+
+    def _fused_available(self, device) -> bool:
+        if device.type != "cuda":
+            return False
+        if not hasattr(self, "_fused_ok"):
+            try:
+                from connectome_kernels import available
+                self._fused_ok = available()
+            except Exception:
+                self._fused_ok = False
+        return self._fused_ok
+
+    def _forward_fused(self, input_ids, state):
+        from connectome_kernels import sparse_recurrence
+        drives = self.input_proj(self.embed(input_ids)).float().permute(1, 2, 0).contiguous()       # [T, n_in, B]
+        vals = self.recurrent.edge_values.float() * self.edge_scale()
+        out = sparse_recurrence(vals, self.leak, self.recurrent.bias.float(), drives, state,
+                                self._fused_graph(), self.config.microsteps)                         # [T, B, N]
+        logits = self.lm_head(out[:, :, self.graph.output_nodes].to(self.lm_head.weight.dtype)).float().permute(1, 0, 2)
+        return logits, out[-1]
+
     def forward(self, input_ids: torch.LongTensor, state: Optional[torch.Tensor] = None,
                 labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None,
                 return_dict: Optional[bool] = None, **kwargs) -> FlyGPTOutput:
         B, T = input_ids.shape
         state = self.init_state(B, input_ids.device) if state is None else state
-        W = self.sparse_weight()
-        outs = []
-        for t in range(T):
-            state = self.step(state, input_ids[:, t], W)
-            outs.append(self.logits_from_state(state))
-        logits = torch.stack(outs, 1)
+        if self._fused_available(input_ids.device):
+            logits, state = self._forward_fused(input_ids, state)
+        else:
+            W = self.sparse_weight()
+            outs = []
+            for t in range(T):
+                state = self.step(state, input_ids[:, t], W)
+                outs.append(self.logits_from_state(state))
+            logits = torch.stack(outs, 1)
         loss = None
         if labels is not None:
             loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]), labels[:, 1:].reshape(-1))

tokenizer_config.json

@@ -1,6 +1,6 @@
 {
   "tokenizer_class": "PreTrainedTokenizerFast",
-  "model_max_length": 1000000,
+  "model_max_length": 10000000,
   "clean_up_tokenization_spaces": false,
   "add_prefix_space": false
 }