File size: 2,714 Bytes
8d31aa6 |
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 |
## Developer: inkbytefo
## Modified: 2025-11-22
import torch
import torch.nn as nn
import torch.nn.functional as F
class ByteLatentEncoder(nn.Module):
"""
Encodes raw byte sequences into latent patch representations.
This module replaces traditional tokenizers by learning to compress
raw bytes directly into a higher-dimensional latent space.
"""
def __init__(
self,
d_model: int,
patch_size: int = 4,
dropout: float = 0.1,
max_len: int = 4096
):
super().__init__()
self.d_model = d_model
self.patch_size = patch_size
# Byte Embedding: 256 possible byte values -> d_model
self.byte_embedding = nn.Embedding(256, d_model)
# Patching mechanism: Strided Convolution
self.patch_conv = nn.Conv1d(
in_channels=d_model,
out_channels=d_model,
kernel_size=patch_size,
stride=patch_size,
padding=0
)
# RoPE (Rotary Positional Embeddings)
# We precompute frequencies for RoPE
self.register_buffer("inv_freq", 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model)))
self.norm = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def apply_rope(self, x: torch.Tensor) -> torch.Tensor:
# x: (B, N, D)
B, N, D = x.shape
# Create position indices
t = torch.arange(N, device=x.device).type_as(self.inv_freq)
freqs = torch.einsum('i,j->ij', t, self.inv_freq) # (N, D/2)
emb = torch.cat((freqs, freqs), dim=-1) # (N, D)
# Apply rotation
# Simple implementation: x_rotated = x * cos(emb) + rotate_half(x) * sin(emb)
# rotate_half: [-x2, x1, -x4, x3, ...]
x1 = x[..., :D//2]
x2 = x[..., D//2:]
rotate_half_x = torch.cat((-x2, x1), dim=-1)
return x * emb.cos() + rotate_half_x * emb.sin()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: (Batch, Seq_Len) tensor of uint8 bytes (0-255)
Returns:
latents: (Batch, Seq_Len // patch_size, d_model)
"""
# 1. Embed bytes
x = self.byte_embedding(x.long())
# 2. Transpose for Conv1d
x = x.transpose(1, 2)
# 3. Apply Patching
x = self.patch_conv(x)
# 4. Transpose back
x = x.transpose(1, 2)
# 5. Apply RoPE
x = self.apply_rope(x)
# 6. Normalize and Dropout
x = self.norm(x)
x = self.dropout(x)
return x
|