change heads_dim to attn_dim, and modify the mlp to be as similar as possible to the attentive probing pattern

Author: rxng8

ngclearn/utils/analysis/attentive_probe.py

@@ -102,11 +102,15 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float =
     attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout)
     features = cross_attention(attn_params, learnable_query, encodings, mask, n_heads, dropout)
     features = features[:, 0]  # (B, 1, dim) => (B, dim)
-    hids = jnp.matmul((features + learnable_query[:, 0]), Whid) + bhid
-    hids = gelu(hids)
+    # MLP
+    residual = features
     if use_LN: ## normalize hidden layer output of probe predictor
-        hids = layer_normalize(hids, Wln_mu, Wln_scale)
-    outs = jnp.matmul(hids, Wy) + by
+        features = layer_normalize(features, Wln_mu, Wln_scale)
+    features = jnp.matmul((features), Whid) + bhid
+    features = gelu(features)
+    features = residual + features
+
+    outs = jnp.matmul(features, Wy) + by
     if use_softmax: ## apply softmax output nonlinearity
         outs = softmax(outs)
     return outs, features
@@ -178,10 +182,11 @@ class AttentiveProbe(Probe):
 
     """
     def __init__(
-            self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, head_dim=64,
+            self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, attn_dim=64,
             target_seq_length=1, learnable_query_dim=31, batch_size=1, hid_dim=32, use_LN=True, use_softmax=True, **kwargs
     ):
         super().__init__(dkey, batch_size, **kwargs)
+        assert attn_dim % num_heads == 0, f"`attn_dim` must be divisible by `num_heads`. Got {attn_dim} and {num_heads}."
         self.dkey, *subkeys = random.split(self.dkey, 12)
         self.num_heads = num_heads
         self.source_seq_length = source_seq_length
@@ -192,24 +197,24 @@ def __init__(
 
         sigma = 0.05
         ## cross-attention parameters
-        Wq = random.normal(subkeys[0], (learnable_query_dim, head_dim)) * sigma
-        bq = random.normal(subkeys[1], (1, head_dim)) * sigma
-        Wk = random.normal(subkeys[2], (input_dim, head_dim)) * sigma
-        bk = random.normal(subkeys[3], (1, head_dim)) * sigma
-        Wv = random.normal(subkeys[4], (input_dim, head_dim)) * sigma
-        bv = random.normal(subkeys[5], (1, head_dim)) * sigma
-        Wout = random.normal(subkeys[6], (head_dim, learnable_query_dim)) * sigma
+        Wq = random.normal(subkeys[0], (learnable_query_dim, attn_dim)) * sigma
+        bq = random.normal(subkeys[1], (1, attn_dim)) * sigma
+        Wk = random.normal(subkeys[2], (input_dim, attn_dim)) * sigma
+        bk = random.normal(subkeys[3], (1, attn_dim)) * sigma
+        Wv = random.normal(subkeys[4], (input_dim, attn_dim)) * sigma
+        bv = random.normal(subkeys[5], (1, attn_dim)) * sigma
+        Wout = random.normal(subkeys[6], (attn_dim, learnable_query_dim)) * sigma
         bout = random.normal(subkeys[7], (1, learnable_query_dim)) * sigma
         #params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout)
         learnable_query = jnp.zeros((batch_size, 1, learnable_query_dim))  # (B, T, D)
         #self.all_params = (learnable_query, *params)
         self.mask = np.zeros((batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor
         ## MLP parameters
-        Whid = random.normal(subkeys[8], (learnable_query_dim, hid_dim)) * sigma
-        bhid = random.normal(subkeys[9], (1, hid_dim)) * sigma
-        Wln_mu = jnp.zeros((1, hid_dim))
-        Wln_scale = jnp.ones((1, hid_dim))
-        Wy = random.normal(subkeys[8], (hid_dim, out_dim)) * sigma
+        Whid = random.normal(subkeys[8], (learnable_query_dim, learnable_query_dim)) * sigma
+        bhid = random.normal(subkeys[9], (1, learnable_query_dim)) * sigma
+        Wln_mu = jnp.zeros((1, learnable_query_dim))
+        Wln_scale = jnp.ones((1, learnable_query_dim))
+        Wy = random.normal(subkeys[8], (learnable_query_dim, out_dim)) * sigma
         by = random.normal(subkeys[9], (1, out_dim)) * sigma
         #mlp_params = (Whid, bhid, Wln_mu, Wln_scale, Wy, by)
         self.probe_params = (learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout, Whid, bhid, Wln_mu, Wln_scale, Wy, by)