minor tweak to attentive prob code comments

Author: Alexander Ororbia

ngclearn/utils/analysis/attentive_probe.py

@@ -232,24 +232,24 @@ def __init__(
         bvs = random.normal(subkeys[13], (1, learnable_query_dim)) * sigma
         Wouts = random.normal(subkeys[14], (learnable_query_dim, learnable_query_dim)) * sigma
         bouts = random.normal(subkeys[15], (1, learnable_query_dim)) * sigma
-        Wlnattn_mu = jnp.zeros((1, learnable_query_dim))
-        Wlnattn_scale = jnp.ones((1, learnable_query_dim))
+        Wlnattn_mu = jnp.zeros((1, learnable_query_dim)) ## LN parameter (applied to output of attention)
+        Wlnattn_scale = jnp.ones((1, learnable_query_dim)) ## LN parameter (applied to output of attention)
         self_attn_params = (Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts, Wlnattn_mu, Wlnattn_scale)
         learnable_query = jnp.zeros((batch_size, 1, learnable_query_dim))  # (B, T, D)
         self.mask = np.zeros((batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor
         ## MLP parameters
         Whid1 = random.normal(subkeys[16], (learnable_query_dim, learnable_query_dim)) * sigma
         bhid1 = random.normal(subkeys[17], (1, learnable_query_dim)) * sigma
-        Wln_mu1 = jnp.zeros((1, learnable_query_dim))
-        Wln_scale1 = jnp.ones((1, learnable_query_dim))
+        Wln_mu1 = jnp.zeros((1, learnable_query_dim)) ## LN parameter
+        Wln_scale1 = jnp.ones((1, learnable_query_dim)) ## LN parameter
         Whid2 = random.normal(subkeys[18], (learnable_query_dim, learnable_query_dim * 4)) * sigma
         bhid2 = random.normal(subkeys[19], (1, learnable_query_dim * 4)) * sigma
-        Wln_mu2 = jnp.zeros((1, learnable_query_dim))
-        Wln_scale2 = jnp.ones((1, learnable_query_dim))
+        Wln_mu2 = jnp.zeros((1, learnable_query_dim)) ## LN parameter
+        Wln_scale2 = jnp.ones((1, learnable_query_dim)) ## LN parameter
         Whid3 = random.normal(subkeys[20], (learnable_query_dim * 4, learnable_query_dim)) * sigma
         bhid3 = random.normal(subkeys[21], (1, learnable_query_dim)) * sigma
-        Wln_mu3 = jnp.zeros((1, learnable_query_dim * 4))
-        Wln_scale3 = jnp.ones((1, learnable_query_dim * 4))
+        Wln_mu3 = jnp.zeros((1, learnable_query_dim * 4)) ## LN parameter
+        Wln_scale3 = jnp.ones((1, learnable_query_dim * 4)) ## LN parameter
         Wy = random.normal(subkeys[22], (learnable_query_dim, out_dim)) * sigma
         by = random.normal(subkeys[23], (1, out_dim)) * sigma
         mlp_params = (Whid1, bhid1, Wln_mu1, Wln_scale1, Whid2, bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3, Wy, by)