Engram Context-Aware Gating

Problem

Implement Engram context-aware gating, a mechanism inspired by memory engrams in neuroscience. The gate modulates a hidden state based on contextual input, selectively amplifying or suppressing features depending on the context signal.

Solution

import numpy as np

def sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

def engram_gating(
    hidden: np.ndarray,
    context: np.ndarray,
    W_gate: np.ndarray,
    b_gate: np.ndarray,
    W_transform: np.ndarray,
    b_transform: np.ndarray
) -> np.ndarray:
    # Concatenate hidden and context
    combined = np.concatenate([hidden, context], axis=-1)

    # Compute gate values (0 to 1) via sigmoid
    gate = sigmoid(combined @ W_gate + b_gate)

    # Compute transformed context
    transformed = np.tanh(context @ W_transform + b_transform)

    # Gated output: blend hidden state with transformed context
    output = gate * hidden + (1 - gate) * transformed

    return output

Explanation

1. Concatenate the hidden state and context vector to form a combined input.
2. Pass the combined input through a linear layer followed by sigmoid to produce gate values in [0, 1].
3. Transform the context through a separate linear layer with tanh activation.
4. The output is a gated blend: gate hidden + (1 - gate) transformed_context. When the gate is near 1, the hidden state passes through unchanged; when near 0, the context-derived signal dominates.

Complexity

• Time: O(d^2) where d is the hidden dimension (matrix multiplications)
• Space: O(d) for intermediate gate and transform vectors