File size: 4,124 Bytes
2de7c42 | 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 | """FractalEmbedding: a trainable fractal codepoint embedding.
Combines three feature sources for each token id t:
(A) 16 deterministic morphological features (CharClassFeatures)
(B) Mandelbrot-decayed Fourier basis (MandelbrotFourierBasis)
(C) Vortex conditioning: a 2-adic hash (Collatz, computed in Rust,
outside the autodiff graph) is projected into phases via a trainable
MLP (PyTorch, in the graph). The 2-adic vortex influences learning
without pretending to be differentiable.
The final projection to d_model is a trainable nn.Linear. The entire forward
pass is differentiable end-to-end. The deterministic parts (A, B, and the
hash of C) are precomputed as buffers outside the graph; only the MLP of C
and the final projection carry trainable parameters.
"""
import torch
import torch.nn as nn
from .char_features import CharClassFeatures
from .fourier import MandelbrotFourierBasis
class FractalEmbedding(nn.Module):
"""Trainable fractal embedding.
Args:
vocab_size: number of token ids covered.
d_model: output dimension.
n_frequencies: number of frequencies for the Fourier basis.
vortex_hidden: width of the MLP that projects the Collatz hash into phases.
collatz_steps: number of Collatz iterations for the hash (deterministic).
"""
def __init__(
self,
vocab_size: int,
d_model: int,
n_frequencies: int = 16,
vortex_hidden: int = 32,
collatz_steps: int = 7,
):
super().__init__()
if vocab_size <= 0 or d_model <= 0:
raise ValueError("vocab_size and d_model must be > 0")
self.vocab_size = vocab_size
self.d_model = d_model
self.collatz_steps = collatz_steps
# (A) Morphological features: deterministic precomputation, outside graph.
char_matrix = torch.stack(
[CharClassFeatures.extract(t) for t in range(vocab_size)], dim=0
)
self.register_buffer("char_features", char_matrix)
# (B) Mandelbrot-decayed Fourier basis: deterministic precomputation.
self.fourier = MandelbrotFourierBasis(vocab_size, n_frequencies)
fourier_matrix = self.fourier.matrix()
self.register_buffer("fourier_features", fourier_matrix)
# (C) Vortex conditioning: Collatz hash precomputed (outside graph),
# then projected by a trainable MLP (in the graph).
try:
from fractus import _core
except ImportError as e:
raise ImportError(
"fractus._core not found. Run `maturin develop`."
) from e
hashes = torch.tensor(
[_core.collatz_hash(t, collatz_steps) for t in range(vocab_size)],
dtype=torch.float32,
)
max_h = hashes.max().item() + 1.0
hashes_norm = hashes / max_h
self.register_buffer("vortex_hashes", hashes_norm)
self.vortex_phase_dim = vortex_hidden
self.vortex_mlp = nn.Sequential(
nn.Linear(1, vortex_hidden),
nn.Tanh(),
nn.Linear(vortex_hidden, vortex_hidden),
)
in_dim = 16 + fourier_matrix.shape[1] + vortex_hidden
self.proj = nn.Linear(in_dim, d_model)
def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
"""token_ids: (N,) or (N, L) of integers in [0, vocab_size).
Returns (N, d_model) or (N, L, d_model).
"""
if token_ids.max() >= self.vocab_size or token_ids.min() < 0:
raise IndexError(
f"token_id outside [0, {self.vocab_size}): "
f"min={int(token_ids.min())}, max={int(token_ids.max())}"
)
original_shape = token_ids.shape
flat = token_ids.reshape(-1)
char = self.char_features[flat]
fourier = self.fourier_features[flat]
h = self.vortex_hashes[flat].unsqueeze(1)
vortex_phases = self.vortex_mlp(h)
x = torch.cat([char, fourier, vortex_phases], dim=1)
out = self.proj(x)
return out.reshape(*original_shape, self.d_model)
|