Learn Everything

class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads')

    def __init__(self, data, children=(), local_grads=()):
        self.data = data          # the scalar value (a float)
        self.grad = 0             # d(loss)/d(self) — filled by backward()
        self._children = children # parent Values that created this one
        self._local_grads = local_grads  # d(self)/d(child) for each parent

    def __add__(self, other):
        # d(a+b)/da = 1, d(a+b)/db = 1
        return Value(self.data + other.data, (self, other), (1, 1))

    def __mul__(self, other):
        # d(a*b)/da = b, d(a*b)/db = a
        return Value(self.data * other.data, (self, other), (other.data, self.data))

    def log(self):
        # d(log(x))/dx = 1/x
        return Value(math.log(self.data), (self,), (1/self.data,))

    def relu(self):
        # d(relu(x))/dx = 1 if x > 0 else 0
        return Value(max(0, self.data), (self,), (float(self.data > 0),))

    def backward(self):
        topo, visited = [], set()
        def build_topo(v):              # DFS post-order traversal
            if v not in visited:
                visited.add(v)
                for child in v._children: build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1                   # d(loss)/d(loss) = 1
        for v in reversed(topo):
            for child, lg in zip(v._children, v._local_grads):
                child.grad += lg * v.grad  # THE chain rule

class Atom:
    def __init__(self, lang, n_embd=32, n_head=4, n_layer=2, block_size=16):
        # Random matrix: nout rows, nin columns, each Value(gauss(0, 0.08))
        mat = lambda nout, nin, std=0.08: [
            [Value(random.gauss(0, std)) for _ in range(nin)]
            for _ in range(nout)
        ]

        self.sd = {
            'wte':     mat(43, 32),    # token embeddings: 43 vocab × 32 dims
            'wpe':     mat(16, 32),    # position embeddings: 16 positions × 32 dims
            'lm_head': mat(43, 32),    # output projection: 43 vocab × 32 dims
        }
        # Per layer: 4 attention matrices + 2 MLP matrices
        for i in range(2):             # 2 layers
            self.sd[f'l{i}.wq'] = mat(32, 32)   # query projection
            self.sd[f'l{i}.wk'] = mat(32, 32)   # key projection
            self.sd[f'l{i}.wv'] = mat(32, 32)   # value projection
            self.sd[f'l{i}.wo'] = mat(32, 32)   # output projection
            self.sd[f'l{i}.f1'] = mat(128, 32)  # MLP expand (4×)
            self.sd[f'l{i}.f2'] = mat(32, 128)  # MLP compress

def forward(self, token_id, pos_id, keys, values):
    sd, hd, nh = self.sd, self.head_dim, self.n_head

    # 1. Look up token embedding + position embedding
    x = [t + p for t, p in zip(sd['wte'][token_id], sd['wpe'][pos_id])]
    x = self._rmsnorm(x)      # normalize: x / sqrt(mean(x^2))

    # 2. For each transformer layer:
    for li in range(self.n_layer):
        xr = x                 # save for residual connection

        # Multi-head attention
        x = self._rmsnorm(x)
        q = self._linear(x, sd[f'l{li}.wq'])  # query: "what am I looking for?"
        k = self._linear(x, sd[f'l{li}.wk'])  # key: "what do I contain?"
        v = self._linear(x, sd[f'l{li}.wv'])  # value: "what info do I carry?"

        # KV cache: remember keys/values from all positions
        keys[li].append(k)
        values[li].append(v)

        # 4-head attention: each head sees 8 dims (32/4)
        xa = []
        for h in range(nh):
            # Attention scores: Q · K^T / sqrt(head_dim)
            al = [sum(qh[j]*kh[t][j]...) / hd**0.5 ...]
            aw = self._softmax(al)    # attention weights
            ho = [sum(aw[t]*vh[t][j]...) ...]  # weighted sum
            xa.extend(ho)

        x = self._linear(xa, sd[f'l{li}.wo'])
        x = [a + b for a, b in zip(x, xr)]  # + residual

        # MLP: expand 4x, ReLU, compress back
        xr = x
        x = self._rmsnorm(x)
        x = self._linear(x, sd[f'l{li}.f1'])  # 32 → 128
        x = [xi.relu() for xi in x]            # non-linearity
        x = self._linear(x, sd[f'l{li}.f2'])  # 128 → 32
        x = [a + b for a, b in zip(x, xr)]    # + residual

    # 3. Project to vocabulary logits
    return self._linear(x, sd['lm_head'])  # 32 → 43 logits

def train_step(self, token_sequence, lr=0.01):
    n = min(self.block_size, len(token_sequence) - 1)
    keys = [[] for _ in range(self.n_layer)]
    vals = [[] for _ in range(self.n_layer)]
    losses = []

    # For each position, predict the next token
    for pos in range(n):
        logits = self.forward(token_sequence[pos], pos, keys, vals)
        probs = self._softmax(logits)
        target = token_sequence[pos + 1]
        losses.append(-probs[target].log())  # cross-entropy

    # Average loss across all positions
    loss = (1 / n) * sum(losses)
    loss.backward()  # traces through ENTIRE computation graph

    # Adam optimizer with linear LR decay
    self.step_count += 1
    lr_t = lr * (1 - self.step_count / 10000)  # decay
    lr_t = max(lr_t, lr * 0.1)                  # floor at 10%
    b1, b2, eps = 0.85, 0.99, 1e-8

    for i, p in enumerate(self.params):
        # Momentum + adaptive learning rate
        self.m[i] = b1 * self.m[i] + (1-b1) * p.grad
        self.v[i] = b2 * self.v[i] + (1-b2) * p.grad**2
        mh = self.m[i] / (1 - b1**self.step_count)  # bias correction
        vh = self.v[i] / (1 - b2**self.step_count)
        p.data -= lr_t * mh / (vh**0.5 + eps)
        p.grad = 0  # reset for next step

def sequence_anomaly_score(atom, token_sequence):
    n = min(atom.block_size, len(token_sequence) - 1)
    keys = [[] for _ in range(atom.n_layer)]
    vals = [[] for _ in range(atom.n_layer)]
    per_token = []

    for pos in range(n):
        logits = atom.forward(token_sequence[pos], pos, keys, vals)
        probs = atom._softmax(logits)
        target = token_sequence[pos + 1]
        target_prob = probs[target].data
        score = -math.log(max(target_prob, 1e-10))  # surprise!
        per_token.append((token_name, score, target_prob))

    avg = sum(s for _, s, _ in per_token) / len(per_token)
    return avg, per_token