File size: 6,347 Bytes
e4f7326 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """MLX counterpart of the FP32 GDN2 trajectory recurrence."""
from __future__ import annotations
import mlx.core as mx
from mlx import nn
def _l2(value: mx.array) -> mx.array:
return value * mx.rsqrt(mx.maximum(mx.sum(mx.square(value), axis=-1, keepdims=True), 1e-12))
class _GateProjection(nn.Module):
def __init__(self, source: int, target: int, rank: int) -> None:
super().__init__()
self.down = nn.Linear(source, rank, bias=False)
self.up = nn.Linear(rank, target, bias=False)
def __call__(self, source: mx.array) -> mx.array:
return self.up(self.down(source))
class GDN2Memory(nn.Module):
def __init__(self, input_dim: int, heads: int, key_dim: int, value_dim: int,
*, observation_dim: int | None = None) -> None:
super().__init__()
if min(input_dim, heads, key_dim, value_dim) <= 0:
raise ValueError("GDN2 dimensions must be positive.")
self.heads, self.key_dim, self.value_dim = heads, key_dim, value_dim
observation_dim = input_dim if observation_dim is None else observation_dim
if observation_dim <= 0:
raise ValueError("GDN2 observation width must be positive.")
self.q_proj = nn.Linear(input_dim, heads * key_dim, bias=False)
self.k_proj = nn.Linear(observation_dim, heads * key_dim, bias=False)
self.v_proj = nn.Linear(observation_dim, heads * value_dim, bias=False)
self.f_proj = _GateProjection(observation_dim, heads * key_dim, min(observation_dim, key_dim))
self.b_proj = nn.Linear(observation_dim, heads * key_dim, bias=False)
self.w_proj = nn.Linear(observation_dim, heads * value_dim, bias=False)
self.g_proj = _GateProjection(input_dim, heads * value_dim, min(input_dim, value_dim))
self.o_proj = nn.Linear(heads * value_dim, input_dim, bias=False)
self.a_log = mx.zeros((heads,), mx.float32)
self.dt_bias = mx.full((heads, key_dim), -6.906255, mx.float32)
def empty(self, *leading: int) -> mx.array:
return mx.zeros((*leading, self.heads, self.key_dim, self.value_dim), mx.float32)
def _projections(self, source: mx.array):
normalized = source.astype(mx.float32)
source = (normalized * mx.rsqrt(mx.mean(mx.square(normalized), axis=-1, keepdims=True) + 1e-6)).astype(source.dtype)
shape = source.shape[:-1]
key_shape = (*shape, self.heads, self.key_dim)
value_shape = (*shape, self.heads, self.value_dim)
k = _l2(nn.silu(self.k_proj(source).astype(mx.float32)).reshape(key_shape))
v = nn.silu(self.v_proj(source).astype(mx.float32)).reshape(value_shape)
head_rate = mx.exp(self.a_log.astype(mx.float32)).reshape(
*((1,) * (source.ndim - 1)), self.heads, 1)
decay = mx.exp(-head_rate * nn.softplus(
self.f_proj(source).astype(mx.float32).reshape(key_shape) + self.dt_bias.astype(mx.float32)))
erase = mx.sigmoid(self.b_proj(source).astype(mx.float32).reshape(key_shape))
write = mx.sigmoid(self.w_proj(source).astype(mx.float32).reshape(value_shape))
return k, v, decay, erase, write
def _query(self, source: mx.array) -> mx.array:
return _l2(nn.silu(self.q_proj(source).astype(mx.float32)).reshape(
*source.shape[:-1], self.heads, self.key_dim))
def _output(self, value: mx.array, source: mx.array) -> mx.array:
gate = nn.silu(self.g_proj(source).astype(mx.float32).reshape(value.shape))
value = value * mx.rsqrt(mx.mean(mx.square(value), axis=-1, keepdims=True) + 1e-6) * gate
return self.o_proj(value.reshape(*source.shape[:-1], -1).astype(source.dtype))
def read(self, state: mx.array, source: mx.array) -> mx.array:
if state.shape != (*source.shape[:-1], self.heads, self.key_dim, self.value_dim):
raise ValueError("GDN2 state and query leading dimensions differ.")
value = mx.sum(self._query(source)[..., None] * state.astype(mx.float32), axis=-2)
return self._output(value, source)
def read_shared(self, state: mx.array, source: mx.array) -> mx.array:
if source.ndim != 3 or state.shape != (source.shape[0], self.heads, self.key_dim, self.value_dim):
raise ValueError("Shared GDN2 read requires [batch, canvas, width] queries.")
query = self._query(source).transpose(0, 2, 1, 3)
value = (query @ state.astype(mx.float32)).transpose(0, 2, 1, 3)
return self._output(value, source)
@staticmethod
def _transition(state, k, v, decay, erase, write, valid):
decayed = state.astype(mx.float32) * decay[..., None]
old = mx.sum((erase * k)[..., None] * decayed, axis=-2)
candidate = decayed + k[..., None] * (write * v - old)[..., None, :]
return candidate if valid is None else mx.where(
valid[..., None, None, None], candidate, state.astype(mx.float32))
def transition(self, state: mx.array, source: mx.array,
valid: mx.array | None = None) -> mx.array:
# Do not override nn.Module.update: it installs parameters for optimizers,
# dtype conversion, and autodiff. State evolution is a separate operation.
if state.shape != (*source.shape[:-1], self.heads, self.key_dim, self.value_dim):
raise ValueError("GDN2 state and observation leading dimensions differ.")
k, v, decay, erase, write = self._projections(source)
if valid is not None:
if valid.shape != source.shape[:-1]:
raise ValueError("GDN2 valid mask must match observation rows.")
return self._transition(state, k, v, decay, erase, write, valid)
def write_sequence(self, state: mx.array, source: mx.array,
valid: mx.array) -> mx.array:
if source.ndim != 3 or valid.shape != source.shape[:2]:
raise ValueError("GDN2 sequence and mask must share [batch, length].")
if state.shape != (source.shape[0], self.heads, self.key_dim, self.value_dim):
raise ValueError("GDN2 sequence state shape differs.")
projected = self._projections(source)
for index in range(source.shape[1]):
state = self._transition(state, *(part[:, index] for part in projected), valid[:, index])
return state
__all__ = ["GDN2Memory"]
|