""" corvidae.py -- the base architecture. This is the foundation every other file (species_memory_bank.py, corvid_extensions.py, corvid_aviary.py) imports and extends. It was referenced throughout those files but never actually written until now; this fills that gap and, per the new batch of papers, upgrades three of its core components so they reflect specific published findings rather than being generic placeholders: CorvidaeMemory (generic DNC-style store) A simplified differentiable neural computer: content-addressed read/write over a persistent memory matrix, plus a `consolidate()` step that blends working memory into a slower long-term buffer (a crude analogue of sleep-dependent memory consolidation) -- called periodically by the trainer via `model.consolidate_memory()`. MetatoolPlanningBuffer Gruber et al. (2019, both Current Biology papers) showed New Caledonian crows keep mental representations of a functional sub-goal, a DISTRACTOR sub-goal, and a final goal simultaneously in mind across multi-stage, out-of-sight tool problems, and suppress attention to the distractor rather than just failing to notice it. This buffer keeps three role-tagged slots (functional sub-goal / distractor sub-goal / final goal) alive across the sequence and explicitly down-weights the distractor slot's contribution to the planning readout, rather than only encoding "the tool. CausalRelationModule Taylor et al. (2009, New Caledonian crows) found crows that solved a trap-tube transferred to a trap-TABLE that shared no visual features at all (different shape, colour, material) -- evidence they'd abstracted a causal relation (object-hole interaction), not a perceptually-bound rule. This module encodes a "causal" embedding via a bottleneck trained to predict task outcome while a gradient-reversal adversary actively strips out predictability of *which surface/apparatus context* produced it -- forcing the representation to be invariant to surface appearance, which is exactly the property needed for trap-tube -> trap-table-style analogical transfer. InhibitionModule reflects a recurring theme across the rook and NC crow trap-tube papers: success correlated with willingness to switch sides / stop an already-started pull (Bird & Emery 2009; Taylor et al. 2009 NC crows) -- i.e., inhibitory control over a prepotent response, not just correct classification. It's implemented as a learned partial gate on the output logits rather than a hard mask, since inhibition in these studies was graded (some individuals showed partial hesitation), not all-or-nothing. """ from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F # ====================== GENERIC DNC-STYLE MEMORY ====================== class CorvidaeMemory(nn.Module): """ A simplified differentiable neural computer. Content-addressed read/write over a persistent (B, memory_size, word_size) matrix, with `num_read_heads` independent read heads. State is lazily created on first forward() call for a given batch size and cleared by `reset()` -- mirroring the lazy-init/explicit-reset convention used by the stateful specialists in corvid_extensions.py. `consolidate(strength)` blends the working memory into a slower long-term buffer and partially blends that consolidated trace back into working memory -- a crude analogue of sleep-dependent replay consolidation. Call this periodically (e.g. every `consolidate_every` training steps), not every step. """ def __init__(self, input_dim: int, memory_size: int = 32, word_size: int = 16, num_read_heads: int = 2): super().__init__() self.memory_size = memory_size self.word_size = word_size self.num_read_heads = num_read_heads interface_size = (num_read_heads * word_size) + num_read_heads + word_size + 1 + word_size + word_size self.interface_proj = nn.Linear(input_dim, interface_size) self._memory: Optional[torch.Tensor] = None # (B, N, W) working memory self._long_term: Optional[torch.Tensor] = None # (B, N, W) consolidated trace def _ensure_state(self, batch_size: int, device): if self._memory is None or self._memory.size(0) != batch_size: self._memory = torch.zeros(batch_size, self.memory_size, self.word_size, device=device) if self._long_term is None or self._long_term.size(0) != batch_size: self._long_term = torch.zeros(batch_size, self.memory_size, self.word_size, device=device) def reset(self): """Clear memory. Lazily re-created on the next forward() call, whatever batch size that call uses -- so this needs no batch_size argument.""" self._memory = None self._long_term = None @torch.no_grad() def consolidate(self, strength: float = 0.16): if self._memory is None or self._long_term is None: return self._long_term = (1 - strength) * self._long_term + strength * self._memory.detach() self._memory = (1 - strength * 0.5) * self._memory + (strength * 0.5) * self._long_term def forward(self, control: torch.Tensor) -> torch.Tensor: """ control: (batch, seq_len, input_dim) Returns: (batch, seq_len, num_read_heads * word_size) concatenated read vectors. """ b, t, _ = control.shape self._ensure_state(b, control.device) interface = self.interface_proj(control) i = 0 read_keys = interface[..., i: i + self.num_read_heads * self.word_size] i += self.num_read_heads * self.word_size read_keys = read_keys.view(b, t, self.num_read_heads, self.word_size) read_strengths = F.softplus(interface[..., i: i + self.num_read_heads]); i += self.num_read_heads write_key = interface[..., i: i + self.word_size]; i += self.word_size write_strength = torch.sigmoid(interface[..., i: i + 1]); i += 1 erase = torch.sigmoid(interface[..., i: i + self.word_size]); i += self.word_size add = torch.tanh(interface[..., i: i + self.word_size]); i += self.word_size reads = [] for step in range(t): wk = write_key[:, step, :] ws = write_strength[:, step, :] er = erase[:, step, :] ad = add[:, step, :] mem_n = F.normalize(self._memory, dim=-1) wk_n = F.normalize(wk, dim=-1).unsqueeze(1) write_weight = F.softmax(torch.einsum('bnd,bod->bn', mem_n, wk_n), dim=-1) * ws self._memory = self._memory * (1 - write_weight.unsqueeze(-1) * er.unsqueeze(1)) \ + write_weight.unsqueeze(-1) * ad.unsqueeze(1) step_reads = [] for h in range(self.num_read_heads): rk = read_keys[:, step, h, :] rs = read_strengths[:, step, h:h + 1] mem_n2 = F.normalize(self._memory, dim=-1) rk_n = F.normalize(rk, dim=-1).unsqueeze(1) read_weight = F.softmax(rs * torch.einsum('bnd,bod->bn', mem_n2, rk_n), dim=-1) read_vec = torch.einsum('bn,bnd->bd', read_weight, self._memory) step_reads.append(read_vec) reads.append(torch.cat(step_reads, dim=-1)) return torch.stack(reads, dim=1) # ====================== METATOOL PLANNING BUFFER ====================== class MetatoolPlanningBuffer(nn.Module): """ Gruber et al. (2019) x2: NC crows solving multi-stage metatool problems kept mental representations of a functional sub-goal, a distractor sub-goal, and the final goal active while stages were out of sight, and specifically avoided distractor apparatuses/tools rather than just failing to encode them -- an active suppression, not an absence of representation. Three role-tagged slots (index 0 = functional sub-goal, 1 = distractor sub-goal, 2 = final goal), initialized from learned role embeddings and slowly updated from context. A learned distractor gate actively suppresses slot 1's contribution to the read-out (rather than that slot simply never being written), and a stage-attention head decides which slot is most relevant to output at each step -- a crude analogue of shifting attention across sub-goal/goal representations as a multi-stage plan unfolds. State persists across a trial (call `clear()` at episode start, matching the lazy-reinit convention of CorvidaeMemory) since the whole point is remembering sub-goals while the relevant apparatus is out of sight. """ def __init__(self, embedding_dim: int, num_slots: int = 3): super().__init__() self.embedding_dim = embedding_dim self.num_slots = num_slots # convention: 0=functional subgoal, 1=distractor, 2=final goal self.role_embeddings = nn.Parameter(torch.randn(num_slots, embedding_dim) * 0.02) self.slot_write = nn.Linear(embedding_dim, embedding_dim) self.distractor_gate = nn.Sequential( nn.Linear(embedding_dim, embedding_dim), nn.ReLU(), nn.Linear(embedding_dim, 1) ) self.stage_attn = nn.Linear(embedding_dim, num_slots) self.output_proj = nn.Linear(embedding_dim, embedding_dim) self._slot_content: Optional[torch.Tensor] = None # (B, num_slots, D) def clear(self): self._slot_content = None def _ensure_state(self, batch_size: int, device): if self._slot_content is None or self._slot_content.size(0) != batch_size: self._slot_content = self.role_embeddings.unsqueeze(0).expand(batch_size, -1, -1).clone().to(device) def forward(self, x: torch.Tensor) -> torch.Tensor: """ x: (batch, seq_len, embedding_dim) Returns: (batch, seq_len, embedding_dim) plan read-out, distractor-suppressed. """ b, t, d = x.shape self._ensure_state(b, x.device) outs = [] for step in range(t): xt = x[:, step, :] update = self.slot_write(xt) new_slot0 = 0.9 * self._slot_content[:, 0, :] + 0.1 * update self._slot_content = torch.cat([ new_slot0.unsqueeze(1), self._slot_content[:, 1:, :] ], dim=1) stage_logits = self.stage_attn(xt) # (B, num_slots) distractor_suppress = torch.sigmoid(self.distractor_gate(self._slot_content[:, 1, :])) # (B,1) suppressed = stage_logits.clone() suppressed[:, 1] = suppressed[:, 1] - 5.0 * distractor_suppress.squeeze(-1) weights = F.softmax(suppressed, dim=-1) read = torch.einsum('bk,bkd->bd', weights, self._slot_content) outs.append(read) out = torch.stack(outs, dim=1) return self.output_proj(out) # ====================== CAUSAL RELATION MODULE (DOMAIN-INVARIANT) ====================== class _GradientReversal(torch.autograd.Function): @staticmethod def forward(ctx, x, lambd): ctx.lambd = lambd return x.view_as(x) @staticmethod def backward(ctx, grad_output): return -ctx.lambd * grad_output, None def grad_reverse(x: torch.Tensor, lambd: float = 1.0) -> torch.Tensor: return _GradientReversal.apply(x, lambd) class CausalRelationModule(nn.Module): """ Taylor et al. (2009): NC crows that solved a trap-tube transferred immediately to a trap-TABLE sharing no visual features (different shape, colour, material) -- the strongest available evidence for a causal/analogical representation rather than a perceptually-bound rule. The design choice that would produce this kind of transfer is a representation that predicts task-relevant outcome while being actively prevented from encoding which surface/apparatus context it came from. Implemented via a small domain-adversarial bottleneck: a causal encoder maps context to a `causal_dim` embedding; an outcome head predicts the task-relevant approach/avoid signal from it; a gradient-reversal adversary tries to predict which of `num_surface_contexts` surface contexts produced the embedding, and its gradient is negated before reaching the causal encoder -- so the encoder is pushed to become WORSE at leaking surface identity even as it stays predictive of outcome. That invariance is what should let a decision rule learned in one surface context transfer to a perceptually distinct one, as in the crows' trap-table transfer. `forward` returns only the causal content embedding (for drop-in compatibility with the rest of the fusion pipeline); the outcome/surface logits needed for the adversarial training objective are stashed on `self.last_outcome_logit` / `self.last_surface_logits` after each call. """ def __init__(self, embedding_dim: int, causal_dim: Optional[int] = None, num_surface_contexts: int = 8, grl_lambda: float = 1.0): super().__init__() causal_dim = causal_dim or embedding_dim self.causal_encoder = nn.Sequential( nn.Linear(embedding_dim, causal_dim), nn.ReLU(), nn.Linear(causal_dim, causal_dim) ) self.outcome_head = nn.Sequential( nn.Linear(causal_dim, causal_dim), nn.ReLU(), nn.Linear(causal_dim, 1) ) self.surface_adversary = nn.Sequential( nn.Linear(causal_dim, causal_dim), nn.ReLU(), nn.Linear(causal_dim, num_surface_contexts) ) self.output_proj = nn.Linear(causal_dim, embedding_dim) self.grl_lambda = grl_lambda self.last_outcome_logit: Optional[torch.Tensor] = None self.last_surface_logits: Optional[torch.Tensor] = None def forward(self, x: torch.Tensor) -> torch.Tensor: causal_repr = self.causal_encoder(x) self.last_outcome_logit = self.outcome_head(causal_repr).squeeze(-1) reversed_repr = grad_reverse(causal_repr, self.grl_lambda) self.last_surface_logits = self.surface_adversary(reversed_repr) return self.output_proj(causal_repr) @staticmethod def outcome_loss(outcome_logit: torch.Tensor, outcome_labels: torch.Tensor) -> torch.Tensor: return F.binary_cross_entropy_with_logits(outcome_logit, outcome_labels) @staticmethod def surface_adversary_loss(surface_logits: torch.Tensor, surface_labels: torch.Tensor) -> torch.Tensor: """Trained normally (not reversed again) -- the reversal already happened on the forward pass via grad_reverse; this loss just gives the adversary a real objective to (try to) solve, which is what makes the reversed gradient meaningful signal for the causal encoder.""" return F.cross_entropy(surface_logits.reshape(-1, surface_logits.size(-1)), surface_labels.reshape(-1)) # ====================== INHIBITORY CONTROL ====================== class InhibitionModule(nn.Module): """ Bird & Emery (2009, rooks) and Taylor et al. (2009, NC crows): success on the trap-tube correlated with willingness to switch/inhibit an already-started response (side-switching), and failure sometimes looked like partial hesitation rather than an outright inability to represent the task (their crow Espanol repeatedly "momentarily halted pulling actions... before eventually pulling it into the trap"). That's graded inhibition, not a hard veto -- so this is a soft, learned partial gate on the output logits rather than a mask that can zero them out entirely. """ def __init__(self, embedding_dim: int, num_classes: int): super().__init__() self.gate_proj = nn.Linear(embedding_dim, num_classes) def forward(self, inhibition_context: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: gate = torch.sigmoid(self.gate_proj(inhibition_context)).unsqueeze(1) # (B,1,num_classes) return logits * (0.5 + 0.5 * gate) # ====================== FULL BASE MODEL ====================== class Corvidae(nn.Module): """ The base architecture. Sensory encoding -> two lateralized "hemisphere" streams (ncl_left / ncl_right, named for the nidopallium caudolaterale, a key executive- function region in corvid brains per Marzluff et al. 2012 and the Swift/Marzluff/ Cross dead-conspecific imaging work) blended into a shared executive stream -> generic DNC memory + metatool planning buffer + causal relation module, fused with lateralized threat processing and a social-context head -> final encoder -> output logits, optionally passed through inhibitory gating. Subclasses (CorvidaeMultiSpecies, CorvidaeAviary) replace the plain `mem_proj` contribution to fusion with a routed bank of species-specific specialists, but this base class is fully runnable standalone. """ def __init__(self, num_embeddings: int, embedding_dim: int, max_seq_len: int, nhead: int = 4, dim_feedforward: int = 128, memory_size: int = 32, memory_word_size: int = 16, num_read_heads: int = 2, num_classes: Optional[int] = None, num_sensory_layers: int = 1, num_ncl_layers: int = 1, num_final_layers: int = 1, dropout: float = 0.1, num_tool_classes: int = 8, num_surface_contexts: int = 8): super().__init__() num_classes = num_classes or num_embeddings self.embedding_dim = embedding_dim self.max_seq_len = max_seq_len self.token_emb = nn.Embedding(num_embeddings, embedding_dim) self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) self.emb_dropout = nn.Dropout(dropout) def make_encoder(num_layers): layer = nn.TransformerEncoderLayer( embedding_dim, nhead, dim_feedforward, dropout=dropout, batch_first=True ) return nn.TransformerEncoder(layer, num_layers=num_layers) self.sensory = make_encoder(num_sensory_layers) self.ncl_left = make_encoder(num_ncl_layers) self.ncl_right = make_encoder(num_ncl_layers) self.final = make_encoder(num_final_layers) self.lateral_strength = nn.Parameter(torch.tensor(0.5)) self.controller = nn.Sequential(nn.Linear(embedding_dim, embedding_dim), nn.ReLU()) self.memory = CorvidaeMemory(embedding_dim, memory_size, memory_word_size, num_read_heads) self.mem_proj = nn.Linear(num_read_heads * memory_word_size, embedding_dim) self.plan_buffer = MetatoolPlanningBuffer(embedding_dim) self.plan_proj = nn.Linear(embedding_dim, embedding_dim) self.causal_relation = CausalRelationModule(embedding_dim, num_surface_contexts=num_surface_contexts) self.threat_left = nn.Linear(embedding_dim, embedding_dim) self.threat_right = nn.Linear(embedding_dim, embedding_dim) self.social_head = nn.Linear(embedding_dim, embedding_dim) fusion_in = embedding_dim * 6 # executive, species/mem, plan, causal, threat, social self.fusion_gate = nn.Linear(fusion_in, embedding_dim) self.fusion_dropout = nn.Dropout(dropout) self.output_proj = nn.Linear(embedding_dim, num_classes) self.inhibition = InhibitionModule(embedding_dim, num_classes) self.planning_head = nn.Linear(embedding_dim, embedding_dim) self.tool_predictor = nn.Linear(embedding_dim, num_tool_classes) def reset_memory(self): """No batch_size needed -- both the DNC memory and the plan buffer lazily re-create their state on the next forward() call, using whatever batch size that call provides.""" self.memory.reset() self.plan_buffer.clear() def consolidate_memory(self, strength: float = 0.16): self.memory.consolidate(strength) def forward(self, x: torch.Tensor, return_planning: bool = False, use_inhibition: bool = True): seq_len = x.size(1) if seq_len > self.max_seq_len: raise ValueError(f"Sequence too long: {seq_len}") causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=x.device) emb = self.token_emb(x) + self.pos_emb(torch.arange(seq_len, device=x.device)) emb = self.emb_dropout(emb) sensory = self.sensory(emb, mask=causal_mask, is_causal=True) left = self.ncl_left(sensory, mask=causal_mask, is_causal=True) right = self.ncl_right(sensory, mask=causal_mask, is_causal=True) lateral = torch.sigmoid(self.lateral_strength) executive = (1 - lateral) * left + lateral * right control = self.controller(executive) mem_read = self.memory(control) mem_proj = self.mem_proj(mem_read) plan_read = self.plan_buffer(executive) plan_proj = self.plan_proj(plan_read) causal_ctx = self.causal_relation(executive) threat_l = self.threat_left(executive.mean(dim=1)) threat_r = self.threat_right(executive.mean(dim=1)) threat_emb = (threat_l + 1.3 * threat_r) * 0.5 # right-hemisphere threat bias, Marzluff et al. 2012 fused_cat = torch.cat([ executive, mem_proj, plan_proj, causal_ctx, threat_emb.unsqueeze(1).expand(-1, seq_len, -1), self.social_head(executive.mean(dim=1)).unsqueeze(1).expand(-1, seq_len, -1), ], dim=-1) gate = torch.sigmoid(self.fusion_gate(fused_cat)) gate = self.fusion_dropout(gate) fused = executive + gate * (mem_proj + plan_proj + causal_ctx) final = self.final(fused, mask=causal_mask, is_causal=True) logits = self.output_proj(final) if use_inhibition: inhibition_context = (causal_ctx.mean(dim=1) + plan_proj.mean(dim=1)) * 0.5 logits = self.inhibition(inhibition_context, logits) if return_planning: plan = self.planning_head(final.mean(dim=1)) tool_logits = self.tool_predictor(plan) return logits, tool_logits return logits # ====================== BASE TRAINER ====================== class CorvidaeTrainer: """ Minimal training loop for the base Corvidae model. Subclassed by SpeciesAwareTrainer / AviaryTrainer to add specialist-specific auxiliary losses. """ def __init__(self, model: "Corvidae", lr: float = 3e-4, replay_weight: float = 0.35, consolidate_every: int = 40): self.model = model self.optimizer = torch.optim.Adam(model.parameters(), lr=lr) self.replay_weight = replay_weight self.consolidate_every = consolidate_every self.global_step = 0 def train_step(self, input_ids: torch.Tensor, targets: torch.Tensor, replay_targets: Optional[torch.Tensor] = None): self.model.train() self.optimizer.zero_grad() logits = self.model(input_ids) ce_loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) replay_loss = torch.tensor(0.0, device=logits.device) if replay_targets is not None: replay_loss = F.mse_loss(logits.mean(dim=1), replay_targets) total_loss = ce_loss + self.replay_weight * replay_loss total_loss.backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) self.optimizer.step() self.global_step += 1 if self.global_step % self.consolidate_every == 0: self.model.consolidate_memory(strength=0.16) return {"total_loss": total_loss.item(), "ce_loss": ce_loss.item(), "replay_loss": replay_loss.item()} if __name__ == "__main__": torch.manual_seed(0) vocab = 200 model = Corvidae(num_embeddings=vocab, embedding_dim=64, max_seq_len=32, nhead=4, dim_feedforward=128, memory_size=32, memory_word_size=16, num_read_heads=2, num_classes=vocab) trainer = CorvidaeTrainer(model, consolidate_every=2) x = torch.randint(0, vocab, (2, 16)) y = torch.randint(0, vocab, (2, 16)) for step in range(5): stats = trainer.train_step(x, y) print(step, stats) model.reset_memory() logits, tool_logits = model(x, return_planning=True) print("logits shape:", logits.shape, "tool_logits shape:", tool_logits.shape) print("smoke test passed")