| |
| |
|
|
|
|
| from __future__ import annotations |
|
|
| import equinox as eqx |
| import jax |
| import jax.numpy as jnp |
| from jaxtyping import Array, Float, Int, PRNGKeyArray |
|
|
| from .odd_ops import Linear, _RMS |
| from .tree import ( |
| _inline_norm_forward, |
| _tagged_bounded_ngpt_gain, |
| _tagged_dense, |
| _tagged_dense_no_bias, |
| _tree_sphere, |
| ) |
|
|
|
|
| def _pool_masked_softmax(scores, mask, axis): |
| neg = jnp.asarray(-1.0e30, dtype=scores.dtype) |
| scores = jnp.where(mask > 0, scores, neg) |
| return jax.nn.softmax(scores, axis=axis) |
|
|
|
|
| class GDescriptorPool(eqx.Module): |
| W_q: Float[Array, "d_g hq"] |
| K: Linear |
| V: Linear |
| ln_in: _RMS |
| n_heads: int = eqx.field(static=True) |
| d_k: int = eqx.field(static=True) |
| d_v: int = eqx.field(static=True) |
| tag: str = eqx.field(static=True, default="") |
|
|
| def __init__( |
| self, |
| d_g: int, |
| d_x: int, |
| *, |
| key: PRNGKeyArray, |
| n_heads: int = 4, |
| d_k: int = 64, |
| d_v: int = 64, |
| tag: str = "", |
| ): |
| kq, kk, kv = jax.random.split(key, 3) |
| self.W_q = jax.random.normal(kq, (d_g, n_heads * d_k)) * (d_g**-0.5) |
| self.K = Linear(d_x, n_heads * d_k, key=kk) |
| self.V = Linear(d_x, n_heads * d_v, key=kv) |
| self.ln_in = _RMS(d_x) |
| self.n_heads = n_heads |
| self.d_k = d_k |
| self.d_v = d_v |
| self.tag = tag |
|
|
| @property |
| def d_out(self) -> int: |
| return self.n_heads * self.d_v |
|
|
| def __call__( |
| self, |
| g, |
| xs, |
| mask, |
| *, |
| kfac_structural_mask=None, |
| kfac_update_mask=None, |
| kfac_scan_shared: bool = False, |
| kfac_repeat_ndim: int = 1, |
| kfac_context_primal_reused_over_walkers: bool = False, |
| ): |
|
|
| n = xs.shape[0] |
| xn = _inline_norm_forward( |
| self.ln_in, |
| xs, |
| pathway="even", |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| q_structural_mask = ( |
| jnp.any(jnp.asarray(kfac_structural_mask).astype(bool)) |
| if kfac_update_mask is None and kfac_structural_mask is not None |
| else kfac_update_mask |
| ) |
| q = _tagged_dense_no_bias( |
| self.W_q, |
| g, |
| tag_id=f"{self.tag}.W_q", |
| pathway="even", |
| kfac_structural_mask=q_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=0, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ).reshape(self.n_heads, self.d_k) |
| k = _tagged_dense( |
| self.K.weight, |
| self.K.bias, |
| xn, |
| tag_id=self.K._use_id, |
| pathway="even", |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ).reshape(n, self.n_heads, self.d_k) |
| v = _tagged_dense( |
| self.V.weight, |
| self.V.bias, |
| xn, |
| tag_id=self.V._use_id, |
| pathway="even", |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ).reshape(n, self.n_heads, self.d_v) |
| scores = jnp.einsum("hd,nhd->hn", q, k) / jnp.sqrt( |
| jnp.asarray(self.d_k, dtype=xs.dtype) |
| ) |
| attn = _pool_masked_softmax(scores, mask[None, :], axis=-1) |
| out = jnp.einsum("hn,nhv->hv", attn, v) |
| return out.reshape(-1) |
|
|
|
|
| def _global_update_parameters(d_g: int, d_pool: int, tap_dim: int, key): |
| tap = int(tap_dim) |
| if tap < 1 or tap >= int(d_g): |
| raise ValueError("global tap dimension must be positive and smaller than d_g") |
| d_in = tap + int(d_pool) |
| d_hidden = 2 * int(d_g) |
| k1, k2, k3 = jax.random.split(key, 3) |
| return ( |
| jax.random.normal(k1, (d_in, d_hidden)) * (d_in**-0.5), |
| jnp.zeros((d_hidden,)), |
| jax.random.normal(k2, (d_hidden, d_g)) * (d_hidden**-0.5), |
| jnp.zeros((d_g,)), |
| jnp.ones((d_in,)), |
| jax.random.normal(k3, (d_g, tap)) * (d_g**-0.5), |
| ) |
|
|
|
|
| def _global_delta( |
| update, |
| g, |
| pool, |
| *, |
| kfac_structural_mask, |
| kfac_g_structural_mask, |
| kfac_scan_shared, |
| kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers, |
| ): |
| from hamiltonzero.model.tree import _tagged_rms_eqx_style |
| from hamiltonzero.model.fused_silu import fused_silu |
|
|
| g_mask = ( |
| kfac_structural_mask |
| if kfac_g_structural_mask is None |
| else kfac_g_structural_mask |
| ) |
| g_in = _tagged_dense_no_bias( |
| update.g_tap_w, |
| g, |
| tag_id=f"{update.tag}.gtap", |
| pathway="even", |
| kfac_structural_mask=g_mask, |
| kfac_scan_shared=( |
| kfac_scan_shared if kfac_g_structural_mask is None else False |
| ), |
| kfac_repeat_ndim=(kfac_repeat_ndim if kfac_g_structural_mask is None else 0), |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| x = jnp.concatenate([g_in, pool.astype(g.dtype)]) |
| x = _tagged_rms_eqx_style( |
| update.ln_s, |
| x, |
| tag_id=f"{update.tag}.ln", |
| pathway="even", |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| hidden = fused_silu( |
| _tagged_dense( |
| update.w1, |
| update.b1, |
| x, |
| tag_id=f"{update.tag}.ffn1", |
| pathway="even", |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| ) |
| return _tagged_dense( |
| update.w2, |
| update.b2, |
| hidden, |
| tag_id=f"{update.tag}.ffn2", |
| pathway="even", |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
|
|
|
|
| class ResidualGlobalUpdate(eqx.Module): |
| w1: Float[Array, "d_in d_hidden"] |
| b1: Float[Array, "d_hidden"] |
| w2: Float[Array, "d_hidden d_g"] |
| b2: Float[Array, "d_g"] |
| ln_s: Float[Array, "d_in"] |
| g_tap_w: Float[Array, "d_g tap_dim"] |
| residual_gain: float = eqx.field(static=True) |
| tag: str = eqx.field(static=True) |
|
|
| def __init__( |
| self, |
| d_g: int, |
| d_pool: int, |
| *, |
| key: PRNGKeyArray, |
| tap_dim: int, |
| residual_gain: float, |
| tag: str, |
| ): |
| ( |
| self.w1, |
| self.b1, |
| self.w2, |
| self.b2, |
| self.ln_s, |
| self.g_tap_w, |
| ) = _global_update_parameters(d_g, d_pool, tap_dim, key) |
| self.residual_gain = float(residual_gain) |
| self.tag = tag |
|
|
| def __call__( |
| self, |
| g, |
| pool, |
| *, |
| kfac_structural_mask=None, |
| kfac_g_structural_mask=None, |
| kfac_scan_shared: bool = False, |
| kfac_repeat_ndim: int = 0, |
| kfac_context_primal_reused_over_walkers: bool = False, |
| ): |
| delta = _global_delta( |
| self, |
| g, |
| pool, |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_g_structural_mask=kfac_g_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| return g + self.residual_gain * delta |
|
|
|
|
| class BoundaryGlobalUpdate(eqx.Module): |
| w1: Float[Array, "d_in d_hidden"] |
| b1: Float[Array, "d_hidden"] |
| w2: Float[Array, "d_hidden d_g"] |
| b2: Float[Array, "d_g"] |
| ln_s: Float[Array, "d_in"] |
| g_tap_w: Float[Array, "d_g tap_dim"] |
| tag: str = eqx.field(static=True) |
|
|
| def __init__( |
| self, d_g: int, d_pool: int, *, key: PRNGKeyArray, tap_dim: int, tag: str |
| ): |
| ( |
| self.w1, |
| self.b1, |
| self.w2, |
| self.b2, |
| self.ln_s, |
| self.g_tap_w, |
| ) = _global_update_parameters(d_g, d_pool, tap_dim, key) |
| self.tag = tag |
|
|
| def __call__( |
| self, |
| g, |
| pool, |
| *, |
| kfac_structural_mask=None, |
| kfac_g_structural_mask=None, |
| kfac_scan_shared: bool = False, |
| kfac_repeat_ndim: int = 0, |
| kfac_context_primal_reused_over_walkers: bool = False, |
| ): |
| delta = _global_delta( |
| self, |
| g, |
| pool, |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_g_structural_mask=kfac_g_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| return _tree_sphere(g + delta) |
|
|
|
|
| class TreeGlobalUpdate(eqx.Module): |
| w1: Float[Array, "d_in d_hidden"] |
| b1: Float[Array, "d_hidden"] |
| w2: Float[Array, "d_hidden d_g"] |
| b2: Float[Array, "d_g"] |
| ln_s: Float[Array, "d_in"] |
| alpha: Float[Array, "d_g"] |
| g_tap_w: Float[Array, "d_g tap_dim"] |
| alpha_max: float = eqx.field(static=True) |
| tag: str = eqx.field(static=True) |
|
|
| def __init__( |
| self, |
| d_g: int, |
| d_pool: int, |
| *, |
| key: PRNGKeyArray, |
| tap_dim: int, |
| alpha_init: float, |
| alpha_max: float, |
| tag: str, |
| ): |
| ( |
| self.w1, |
| self.b1, |
| self.w2, |
| self.b2, |
| self.ln_s, |
| self.g_tap_w, |
| ) = _global_update_parameters(d_g, d_pool, tap_dim, key) |
| self.alpha = float(alpha_init) * jnp.ones((d_g,)) |
| self.alpha_max = float(alpha_max) |
| self.tag = tag |
|
|
| def __call__( |
| self, |
| g, |
| pool, |
| update_mask=None, |
| *, |
| kfac_structural_mask=None, |
| kfac_g_structural_mask=None, |
| kfac_scan_shared: bool = False, |
| kfac_repeat_ndim: int = 0, |
| kfac_context_primal_reused_over_walkers: bool = False, |
| ): |
| delta = _global_delta( |
| self, |
| g, |
| pool, |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_g_structural_mask=kfac_g_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| skip = _tree_sphere(g) |
| proposal = _tree_sphere(delta) |
| direction = proposal - skip |
| gain = _tagged_bounded_ngpt_gain( |
| self.alpha, |
| direction, |
| max_gain=self.alpha_max, |
| tag_id=f"{self.tag}.alpha", |
| pathway="even", |
| kfac_structural_mask=kfac_structural_mask, |
| kfac_scan_shared=kfac_scan_shared, |
| kfac_repeat_ndim=kfac_repeat_ndim, |
| kfac_context_primal_reused_over_walkers=( |
| kfac_context_primal_reused_over_walkers |
| ), |
| ) |
| updated = _tree_sphere(skip + gain * direction) |
| if update_mask is None: |
| return updated |
| active = jnp.asarray(update_mask).astype(bool) |
| while active.ndim < updated.ndim: |
| active = active[..., None] |
| return jnp.where(active, updated, g) |
|
|
|
|
| class EdgeRowColGlobalUpdate(eqx.Module): |
| row_pool: GDescriptorPool |
| col_pool: GDescriptorPool |
| set_pool: GDescriptorPool |
| update: BoundaryGlobalUpdate |
|
|
| def __init__( |
| self, |
| d_g: int, |
| d_edge: int, |
| *, |
| key: PRNGKeyArray, |
| n_heads: int = 4, |
| d_k: int = 64, |
| d_v: int = 64, |
| tag: str = "", |
| tap_dim: int, |
| ): |
| kr, kc, ks, ku = jax.random.split(key, 4) |
| self.row_pool = GDescriptorPool( |
| d_g, d_edge, key=kr, n_heads=n_heads, d_k=d_k, d_v=d_v, tag=f"{tag}.row" |
| ) |
| self.col_pool = GDescriptorPool( |
| d_g, d_edge, key=kc, n_heads=n_heads, d_k=d_k, d_v=d_v, tag=f"{tag}.col" |
| ) |
| self.set_pool = GDescriptorPool( |
| d_g, |
| n_heads * d_v, |
| key=ks, |
| n_heads=n_heads, |
| d_k=d_k, |
| d_v=d_v, |
| tag=f"{tag}.set", |
| ) |
| self.update = BoundaryGlobalUpdate( |
| d_g, |
| n_heads * d_v, |
| key=ku, |
| tag=f"{tag}.upd", |
| tap_dim=tap_dim, |
| ) |
|
|
| def __call__(self, g, edge, mask): |
|
|
| system_active = jnp.any(mask.astype(bool)) |
| rows = jax.vmap( |
| lambda ei, mi: self.row_pool( |
| g, |
| ei, |
| mask, |
| kfac_structural_mask=mi * mask, |
| kfac_update_mask=system_active, |
| kfac_scan_shared=False, |
| kfac_repeat_ndim=2, |
| ) |
| )(edge, mask) |
| cols = jax.vmap( |
| lambda ej, mj: self.col_pool( |
| g, |
| ej, |
| mask, |
| kfac_structural_mask=mj * mask, |
| kfac_update_mask=system_active, |
| kfac_scan_shared=False, |
| kfac_repeat_ndim=2, |
| ) |
| )(jnp.swapaxes(edge, 0, 1), mask) |
| descs = jnp.concatenate([rows, cols], axis=0) |
| dmask = jnp.concatenate([mask, mask], axis=0) |
| pooled = self.set_pool( |
| g, |
| descs, |
| dmask, |
| kfac_structural_mask=dmask, |
| kfac_update_mask=system_active, |
| kfac_scan_shared=False, |
| kfac_repeat_ndim=1, |
| ) |
| return self.update( |
| g, |
| pooled, |
| kfac_structural_mask=system_active, |
| kfac_scan_shared=False, |
| ) |
|
|
