15_07_expt/scripts/ctm_components.py

"""Generic CTM building blocks: NeuronModel, SynapseNet, compute_synchronization."""

from __future__ import annotations

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


class NeuronModel(nn.Module):
    """Per-neuron linear model: maps each neuron's rolling history -> next activation."""

    def __init__(self, memory_length: int, d_model: int, dropout: float = 0.0):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(memory_length)

        bound = 1.0 / math.sqrt(memory_length + 2)
        self.weight = nn.Parameter(torch.empty(memory_length, 2, d_model).uniform_(-bound, bound))
        self.bias = nn.Parameter(torch.zeros(1, d_model, 2))
        self.temperature = nn.Parameter(torch.ones(1))

    def forward(self, input_trace: torch.Tensor) -> torch.Tensor:
        x = self.dropout(input_trace)
        x = self.layer_norm(x)
        x = torch.einsum("BNM,MON->BNO", x, self.weight) + self.bias
        return F.glu(x, dim=-1).squeeze(-1) / self.temperature


class SynapseNet(nn.Module):
    """Maps cat(attn_out, activated_state) -> new pre-NLM internal state."""

    def __init__(self, d_input: int, d_model: int, dropout: float = 0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(d_input + d_model, d_model * 2),
            nn.GLU(dim=-1),
            nn.LayerNorm(d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def compute_synchronization(
    activated_state: torch.Tensor,
    ema_numer: Optional[torch.Tensor],
    ema_denom: Optional[torch.Tensor],
    decay_rate: torch.Tensor,
    n_synch: int,
    idx_left: torch.Tensor,
    idx_right: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    select_left = activated_state[:, idx_left]
    select_right = activated_state[:, idx_right]
    outer_prod = select_left.unsqueeze(2) * select_right.unsqueeze(1)
    row_idx, col_idx = torch.triu_indices(n_synch, n_synch, device=activated_state.device)
    pairwise = outer_prod[:, row_idx, col_idx]

    if ema_numer is None or ema_denom is None:
        ema_numer = pairwise
        ema_denom = torch.ones_like(pairwise)
    else:
        ema_numer = decay_rate * ema_numer + pairwise
        ema_denom = decay_rate * ema_denom + 1.0

    return ema_numer / ema_denom.sqrt(), ema_numer, ema_denom