| |
| |
| |
| class FuzzyHeuristicRetriever(nn.Module): |
| """ |
| mu_C(h_t, S_i) = exp(-gamma * ||h_t - S_i||^2) |
| Top-k retrieval: O(log|S| * d) |
| """ |
| def __init__(self, d_state: int, n_heuristics: int, gamma: float = 1.0, |
| top_k: int = 5, *, rngs: nn.Rngs): |
| super().__init__() |
| self.d_state = d_state |
| self.n_heuristics = n_heuristics |
| self.top_k = top_k |
| self.gamma = nn.Param(jnp.array(gamma)) |
| key = rngs.params() |
| self.S = nn.Param(jax.random.normal(key, (n_heuristics, d_state)) * 0.1) |
| self.value_proj = nn.Linear(d_state, d_state, rngs=rngs) |
|
|
| def gaussian_membership(self, h_t, S_i): |
| h_sq = jnp.sum(h_t**2, axis=-1, keepdims=True) |
| s_sq = jnp.sum(S_i**2, axis=-1) |
| h_dot_s = h_t @ S_i.T |
| dist_sq = jnp.maximum(h_sq - 2*h_dot_s + s_sq, 0.0) |
| gamma = jax.nn.softplus(self.gamma.value) |
| return jnp.exp(-gamma * dist_sq) |
|
|
| def __call__(self, h_t): |
| if h_t.ndim == 3: |
| batch, seq_len, d = h_t.shape |
| h_t_flat = h_t.reshape(-1, d) |
| else: |
| h_t_flat = h_t |
| batch = seq_len = None |
| S_i = self.S.value |
| mu = self.gaussian_membership(h_t_flat, S_i) |
| top_k = min(self.top_k, self.n_heuristics) |
| top_k_mu, top_k_idx = jax.lax.top_k(mu, top_k) |
| top_k_S = S_i[top_k_idx] |
| weights = top_k_mu / (jnp.sum(top_k_mu, axis=-1, keepdims=True) + 1e-8) |
| retrieved = jnp.sum(weights[..., None] * top_k_S, axis=-2) |
| retrieved = self.value_proj(retrieved) |
| w_f = jnp.max(mu, axis=-1, keepdims=True) |
| if batch is not None: |
| retrieved = retrieved.reshape(batch, seq_len, self.d_state) |
| w_f = w_f.reshape(batch, seq_len, 1) |
| mu = mu.reshape(batch, seq_len, self.n_heuristics) |
| return retrieved, w_f, mu |
|
|
| print("[3/6] FuzzyHeuristicRetriever defined") |
|
|
