| """ |
| model.py - TRM classifier for heap exploit detection. |
| |
| Standalone module with no external dependencies beyond torch + numpy. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| from pathlib import Path |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-6): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.w1 = nn.Linear(dim, dim * 2, bias=False) |
| self.w2 = nn.Linear(dim, dim, bias=False) |
|
|
| def forward(self, x): |
| gate, val = self.w1(x).chunk(2, dim=-1) |
| return self.w2(F.silu(gate) * val) |
|
|
|
|
| class RecursionBlock(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.norm1 = RMSNorm(dim) |
| self.swiglu = SwiGLU(dim) |
| self.norm2 = RMSNorm(dim) |
|
|
| def forward(self, x): |
| return self.norm2(self.swiglu(self.norm1(x))) |
|
|
|
|
| class HeapTRM(nn.Module): |
| """Tiny Recursive Model for heap exploit classification.""" |
|
|
| def __init__(self, vocab_size=64, hidden_dim=128, seq_len=512, |
| n_outer=2, n_inner=3, n_classes=2): |
| super().__init__() |
| self.n_outer = n_outer |
| self.n_inner = n_inner |
| self.embed = nn.Embedding(vocab_size, hidden_dim) |
| self.y_init = nn.Parameter(torch.randn(1, seq_len, hidden_dim) * 0.02) |
| self.z_init = nn.Parameter(torch.randn(1, seq_len, hidden_dim) * 0.02) |
| self.block_z = RecursionBlock(hidden_dim) |
| self.block_y = RecursionBlock(hidden_dim) |
| self.pos_embed = nn.Parameter(torch.randn(1, seq_len, hidden_dim) * 0.02) |
| self.out_norm = RMSNorm(hidden_dim) |
| self.out_head = nn.Linear(hidden_dim, n_classes) |
|
|
| def forward(self, x): |
| B = x.shape[0] |
| h = self.embed(x.reshape(B, -1)) + self.pos_embed |
| y = self.y_init.expand(B, -1, -1) |
| z = self.z_init.expand(B, -1, -1) |
| for _ in range(self.n_outer): |
| for _ in range(self.n_inner): |
| z = z + self.block_z(h + y + z) |
| y = y + self.block_y(y + z) |
| pooled = self.out_norm(y).mean(dim=1) |
| return self.out_head(pooled) |
|
|
