Upload ByteETM-Korean (HF inference compatible)

config.json

@@ -6,7 +6,7 @@
   "dtype": "float32",
   "model_type": "byteetm",
   "n_embd": 512,
-  "n_head": 8,
+  "n_head": 16,
   "n_layer": 4,
   "transformers_version": "4.57.1",
   "vocab_size": 258,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9ef2d672b4c0e5818c0cd68e45cdc879df4406dbb3520715dc7097b17e8d9f19
-size 65296016
+oid sha256:a9f6a671faac0a301ebebedfbf2bcaa3457ec013cff8059bc5efdc72eab66274
+size 77830592

modeling_byteetm.py

@@ -5,52 +5,74 @@ import math, random, numpy as np, torch, torch.nn as nn, torch.nn.functional as
 # ---------- 4. 모델 정의 ----------
 # === GeneratingSeries 기반 보조 모듈 ===
 class MomentumEncoder(nn.Module):
-    """토큰 임베딩 간의 차분을 포함한 동적 인코딩"""
-    def __init__(self, dim):
+    """다항 차분 + 게이트 통합"""
+    def __init__(self, dim, max_order=3):
         super().__init__()
-        self.linear = nn.Linear(dim * 2, dim)
+        self.max_order = max_order
+        self.proj = nn.Linear(dim * (max_order + 1), dim)
+        self.gate = nn.Linear(dim, dim)
         self.norm = nn.LayerNorm(dim)
-        self.act = nn.Tanh()
-    def forward(self, x):  # [B,T,C]
-        diff = F.pad(x[:, 1:] - x[:, :-1], (0,0,1,0))
-        return self.act(self.norm(self.linear(torch.cat([x, diff], dim=-1))))
+
+    def forward(self, x):
+        diffs = [x]
+        for k in range(1, self.max_order + 1):
+            d = F.pad(x[:, k:] - x[:, :-k], (0, 0, k, 0))
+            diffs.append(d)
+        concat = torch.cat(diffs, dim=-1)
+        h = self.proj(concat)
+        g = torch.sigmoid(self.gate(x))
+        return self.norm(h * g + x * (1 - g))
+
 
 class GFLayer(nn.Module):
-    """지수 감쇠 기반의 생성함수 확장"""
-    def __init__(self, dim, max_order=6, tau_scale=0.01):
+    """Adaptive polynomial generating function"""
+    def __init__(self, dim, max_order=6):
         super().__init__()
         self.coeff = nn.Parameter(torch.randn(dim, max_order + 1) * 0.1)
-        self.tau = nn.Parameter(torch.ones(dim) * tau_scale)
-        self.max_order = max_order
+        self.alpha = nn.Parameter(torch.randn(dim) * 0.1)
+
     def forward(self, x):
         B, T, D = x.shape
-        t = torch.arange(T, device=x.device).float().view(1,T,1)
-        z = torch.exp(-t * self.tau.view(1,1,D))
-        powers = torch.stack([z**k for k in range(self.max_order+1)], dim=-1)
-        gen = torch.einsum('btdk,dk->btd', powers, self.coeff)
+        t = torch.linspace(0, 1, T, device=x.device).view(1, T, 1)
+        basis = torch.stack([(t ** k) * torch.exp(-self.alpha.view(1,1,D)*t) for k in range(self.coeff.size(1))], dim=-1)
+        gen = torch.einsum("btdk,dk->btd", basis, self.coeff)
         return x + gen
 
+
 class OrthogonalTemporalProjector(nn.Module):
-    """시퀀스 길이 방향으로 직교 기저 투영"""
-    def __init__(self, t_len, rank=8):
+    """Adaptive rank orthogonal projection"""
+    def __init__(self, t_len, dim, rank_ratio=0.25):
         super().__init__()
+        rank = max(4, int(rank_ratio * math.sqrt(dim)))
         self.U = nn.Parameter(torch.randn(t_len, rank) / math.sqrt(t_len))
+
     def forward(self, x):
-        B,T,D = x.shape
-        if T != self.U.size(0):
-            U = F.interpolate(self.U.T.unsqueeze(0), size=T, mode="linear", align_corners=False).squeeze(0).T
-        else:
-            U = self.U
+        B, T, D = x.shape
+        U = F.interpolate(self.U.T.unsqueeze(0), size=T, mode="linear", align_corners=False).squeeze(0).T
         U = F.normalize(U, dim=0)
         P = U @ U.T
-        trend = torch.einsum('btd,ts->bsd', x, P)
+        trend = torch.einsum("btd,ts->bsd", x, P)
         resid = x - trend
-        return 0.5*(trend + resid)
+        return trend + 0.5 * resid
+
+class SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, dim, max_len=2048):
+        super().__init__()
+        pe = torch.zeros(max_len, dim)
+        pos = torch.arange(0, max_len).unsqueeze(1)
+        div = torch.exp(torch.arange(0, dim, 2) * (-math.log(10000.0) / dim))
+        pe[:, 0::2] = torch.sin(pos * div)
+        pe[:, 1::2] = torch.cos(pos * div)
+        self.register_buffer("pe", pe.unsqueeze(0))
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1)]
+
 
 # === GPT Block 확장 ===
 class GeneratingBlock(nn.Module):
     """기존 Transformer Block + GeneratingSeries 동역학 통합"""
-    def __init__(self, n_embd, n_head, block_size, dropout=0.0, gf_order=6):
+    def __init__(self, n_embd, n_head, block_size, dropout=0.0, gf_order=2):
         super().__init__()
         self.ln1 = nn.LayerNorm(n_embd)
         self.ln2 = nn.LayerNorm(n_embd)
@@ -59,7 +81,8 @@ class GeneratingBlock(nn.Module):
         # GeneratingSeries 요소
         self.momentum = MomentumEncoder(n_embd)
         self.gf = GFLayer(n_embd, max_order=gf_order)
-        self.otp = OrthogonalTemporalProjector(block_size, rank=min(8, block_size//4))
+        self.otp = OrthogonalTemporalProjector(block_size, n_embd)
+        
     def forward(self, x):
         # step1: momentum encoding (local diff)
         x = self.momentum(x)
@@ -92,14 +115,17 @@ class CausalSelfAttention(nn.Module):
         k = self.key(x).view(B, T, self.n_head, C//self.n_head).transpose(1,2)
         q = self.query(x).view(B, T, self.n_head, C//self.n_head).transpose(1,2)
         v = self.value(x).view(B, T, self.n_head, C//self.n_head).transpose(1,2)
-        att = (q @ k.transpose(-2,-1)) / math.sqrt(k.size(-1))
-        att = att.masked_fill(self.mask[:,:,:T,:T]==0, float("-inf"))
+
+        # RMS normalization per head
+        q = q / (q.pow(2).mean(-1, keepdim=True).sqrt() + 1e-6)
+        k = k / (k.pow(2).mean(-1, keepdim=True).sqrt() + 1e-6)
+
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))
+        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
         att = F.softmax(att, dim=-1)
         att = self.attn_drop(att)
-        y = att @ v
-        y = y.transpose(1,2).contiguous().view(B,T,C)
-        y = self.resid_drop(self.proj(y))
-        return y
+        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_drop(self.proj(y))
 
 class MLP(nn.Module):
     def __init__(self, n_embd, dropout=0.0):
@@ -128,10 +154,12 @@ class ByteETM(nn.Module):
     def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size, dropout=0.0):
         super().__init__()
         self.token_emb = nn.Embedding(vocab_size, n_embd)
-        self.pos_emb   = nn.Embedding(block_size, n_embd)
+        self.pos_enc   = SinusoidalPositionalEncoding(n_embd, max_len=block_size)
         self.drop = nn.Dropout(dropout)
-        # self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
-        self.blocks = nn.ModuleList([GeneratingBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
+
+        self.blocks = nn.ModuleList([
+            GeneratingBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer)
+        ])
         self.ln_f = nn.LayerNorm(n_embd)
         self.head = nn.Linear(n_embd, vocab_size, bias=False)
         self.block_size = block_size
@@ -146,13 +174,14 @@ class ByteETM(nn.Module):
     def forward(self, idx, targets=None):
         B, T = idx.size()
         assert T <= self.block_size
-        pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
-        x = self.token_emb(idx) + self.pos_emb(pos)
+        x = self.token_emb(idx)
+        x = self.pos_enc(x)          # ← 여기서 사인·코사인 위치 정보 추가
         x = self.drop(x)
+
         for blk in self.blocks:
             x = blk(x)
         x = self.ln_f(x)
-        logits = self.head(x)  # (B,T,V)
+        logits = self.head(x)
         loss = None
         if targets is not None:
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))