| |
| |
| |
| class CogniThreadModel(nn.Module): |
| """ |
| Full CogniThread SSM architecture combining: |
| 1. CogniThreadSSM (Mamba-style SSM with ZOH + parallel scan) |
| 2. FuzzyHeuristicRetriever (Gaussian fuzzy membership) |
| 3. LatentDiffusionSynthesizer (conditioned reverse diffusion) |
| 4. AdaptiveGate (variance-constrained gating) |
| |
| Forward: |
| h_t = SSM(x) # SSM trajectory |
| retrieved, w_f, mu = Fuzzy(h_t) # fuzzy retrieval |
| z_0 = Diffusion(h_t) # diffusion generation |
| y_final, var = Gate(retrieved, z_0, w_f) # adaptive gating |
| """ |
| def __init__(self, d_model: int, d_state: int = 16, n_heuristics: int = 64, |
| d_z: int = 32, T_diff: int = 10, d_conv: int = 4, |
| expand: int = 2, gamma: float = 1.0, top_k: int = 5, *, |
| rngs: nn.Rngs): |
| super().__init__() |
| self.d_model = d_model |
| self.d_state = d_state |
| self.d_z = d_z |
| self.ssm = CogniThreadSSM(d_model, d_state, d_conv, expand, rngs=rngs) |
| |
| self.state_proj = nn.Linear(d_model, d_state, rngs=rngs) |
| self.fuzzy = FuzzyHeuristicRetriever(d_state, n_heuristics, gamma, top_k, rngs=rngs) |
| self.diffusion = LatentDiffusionSynthesizer(d_z, d_state, T_diff, rngs=rngs) |
| |
| self.diff_proj = nn.Linear(d_z, d_model, rngs=rngs) |
| |
| self.fuzzy_proj = nn.Linear(d_state, d_model, rngs=rngs) |
| self.gate = AdaptiveGate(d_model, rngs=rngs) |
|
|
| def __call__(self, x, key=None): |
| if key is None: |
| key = jax.random.key(0) |
| |
| h_ssm = self.ssm(x) |
| |
| h_t = self.state_proj(h_ssm) |
| |
| retrieved, w_f, mu = self.fuzzy(h_t) |
| retrieved = self.fuzzy_proj(retrieved) |
| |
| z_0 = self.diffusion(h_t, key=key) |
| generated = self.diff_proj(z_0) |
| |
| y_final, var_y = self.gate(retrieved, generated, w_f) |
| return y_final, var_y, w_f, mu |
|
|
| print("[6/6] CogniThreadModel defined") |
|
|
