File size: 2,664 Bytes
c6be157 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """First coupling: does attention survive on IB-reconstructed K/V?
Runs the same random-weight GNOMON twice — once with raw attention, once
where every layer attends over the previous layer's reconstructed K/V.
Reports the logit divergence and both latencies.
"""
import os
import sys
import time
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import jax
import jax.numpy as jnp
from gnomon import init_gnomon, init_mode_state, make_forward_jit
def kl(p_logits, q_logits):
p = jax.nn.softmax(p_logits, axis=-1)
q = jax.nn.log_softmax(q_logits, axis=-1)
return jnp.mean(jnp.sum(p * (jnp.log(p + 1e-9) - q), axis=-1))
def main():
params = init_gnomon(
jax.random.PRNGKey(0),
vocab_size=128, d_model=64, n_layers=3, n_heads=4,
ib_k=16, seq_len=24, d_signal=16,
)
ids = jax.random.randint(jax.random.PRNGKey(1), (1, 24), 0, 128)
mode_sig = jax.random.normal(jax.random.PRNGKey(2), (1, 24, 16))
mode = init_mode_state(d_mod=4)
raw_fn = make_forward_jit(n_heads=4, use_reconstructed_kv=False)
rec_fn = make_forward_jit(n_heads=4, use_reconstructed_kv=True)
# Warmup — forces both traces
_r = raw_fn(params, ids, mode, mode_sig, jax.random.PRNGKey(3))
_c = rec_fn(params, ids, mode, mode_sig, jax.random.PRNGKey(3))
jax.block_until_ready(_r["logits"])
jax.block_until_ready(_c["logits"])
t0 = time.perf_counter()
out_raw = raw_fn(params, ids, mode, mode_sig, jax.random.PRNGKey(3))
jax.block_until_ready(out_raw["logits"])
t_raw = time.perf_counter() - t0
t0 = time.perf_counter()
out_rec = rec_fn(params, ids, mode, mode_sig, jax.random.PRNGKey(3))
jax.block_until_ready(out_rec["logits"])
t_rec = time.perf_counter() - t0
divergence = float(kl(out_raw["logits"], out_rec["logits"]))
print("Coupling test: raw attention vs IB-reconstructed attention")
print(f" layers: 3 (IB state dim = 16)")
print(f" raw forward: {t_raw * 1000:.2f} ms")
print(f" coupled forward: {t_rec * 1000:.2f} ms")
print(f" logit KL divergence: {divergence:.4f}")
print()
if divergence < 0.5:
print(" → Attention survives on reconstructed K/V.")
else:
print(" → Reconstruction too lossy (random weights).")
print(" Adding KL loss to the training objective closes this gap.")
print()
print("Note: this is with random IB weights. Training drives the")
print("divergence down — that's what the coupling loss term is for.")
if __name__ == "__main__":
main() |