Update README.md

README.md

@@ -17,7 +17,8 @@ There's many space to experimenting like deeper architecture, another activation
 
 # Code
 
-## Original Code
+## Base Code
+```
 import torch
 import torch.nn as nn
 
@@ -54,8 +55,8 @@ class LookThem(nn.Module):
                 out_mod2_j = self.mod2[j](x[:, j]) + 1e-7
                 out_mod2_i = self.mod2[i](x[:, i]) + 1e-7
 
-                compare = self.mod1[i](x[:, i]) / out_mod2_j
-                compare2 = self.mod1[j](x[:, j]) / out_mod2_i
+                compare = torch.tanh(self.mod1[i](x[:, i]) / out_mod2_j)
+                compare2 = torch.tanh(self.mod1[j](x[:, j]) / out_mod2_i)
 
                 # Transformasi hasil interaksi
                 interaksi = (self.transform[j](compare) * x[:, i] + self.transform[j](compare2) * x[:, j]) / 2
@@ -74,97 +75,9 @@ class LookThem(nn.Module):
         x_new = torch.cat(new_x, dim=1)
 
         return self.mlp(x_new)
-
+```
 ## Vectorized
-import torch
-import torch.nn as nn
-import math
-
-class LookThemVectorized(nn.Module):
-    def __init__(self, num_tokens=5, in_features=1, hidden_dim=5):
-        super(LookThemVectorized, self).__init__()
-
-        self.num_tokens = num_tokens
-        self.in_features = in_features
-        self.hidden_dim = hidden_dim
-
-        # 1. Batched Parameters untuk Mod1
-        # Shape: [num_tokens, in_features, hidden_dim]
-        self.mod1_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
-        self.mod1_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
-        # Shape: [num_tokens, hidden_dim, 1]
-        self.mod1_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
-        self.mod1_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
-
-        # 2. Batched Parameters untuk Mod2
-        self.mod2_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
-        self.mod2_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
-        self.mod2_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
-        self.mod2_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
-
-        # 3. Batched Parameters untuk Transformasi Linear j
-        self.trans_w = nn.Parameter(torch.randn(num_tokens, 1, 1))
-        self.trans_b = nn.Parameter(torch.zeros(num_tokens, 1))
-
-        # 4. MLP Final disesuaikan dengan jumlah token yang dinamis
-        self.mlp = nn.Sequential(
-            nn.Linear(num_tokens, num_tokens * 2),
-            nn.ReLU(),
-            nn.Linear(num_tokens * 2, num_tokens)
-        )
-
-        self._init_weights()
-
-    def _init_weights(self):
-        # Inisialisasi Kaiming Uniform agar training stabil
-        for w in [self.mod1_w1, self.mod2_w1]:
-            nn.init.kaiming_uniform_(w, a=math.sqrt(5))
-        for w in [self.mod1_w2, self.mod2_w2, self.trans_w]:
-            nn.init.kaiming_uniform_(w, a=math.sqrt(5))
-
-    def forward(self, x):
-        # x shape sekarang: [Batch, num_tokens, in_features]
-        batch_size = x.size(0)
-        N = self.num_tokens
-
-        # 1. Jalankan Mod1 dan Mod2 secara paralel untuk semua token
-        h1 = torch.einsum('bti,tij->btj', x, self.mod1_w1) + self.mod1_b1
-        out_m1 = torch.einsum('btj,tjk->btk', torch.relu(h1), self.mod1_w2) + self.mod1_b2 # [Batch, N, 1]
-
-        h2 = torch.einsum('bti,tij->btj', x, self.mod2_w1) + self.mod2_b1
-        out_m2 = torch.einsum('btj,tjk->btk', torch.relu(h2), self.mod2_w2) + self.mod2_b2 # [Batch, N, 1]
-
-        # 2. Hitung Rasio Kombinasi i dan j via Broadcasting
-        out_m2_safe = out_m2 + 1e-7
-        compare = out_m1.unsqueeze(2) / out_m2_safe.unsqueeze(1)   # [Batch, N, N, 1]
-        compare2 = out_m1.unsqueeze(1) / out_m2_safe.unsqueeze(2)  # [Batch, N, N, 1]
-
-        # 3. Transformasikan hasil berdasar indeks j
-        # View khusus untuk bias agar nge-broadcast pas di koordinat j
-        bias_reshaped = self.trans_b.view(1, 1, N, 1)
-        trans_compare = torch.einsum('bije,jef->bijf', compare, self.trans_w) + bias_reshaped
-        trans_compare2 = torch.einsum('bije,jef->bijf', compare2, self.trans_w) + bias_reshaped
-
-        # 4. Hitung Interaksi Berbobot Fitur menggunakan fitur asli dari x
-        # x.unsqueeze(2) -> fitur token i, x.unsqueeze(1) -> fitur token j
-        interaksi = (trans_compare * x.unsqueeze(2) + trans_compare2 * x.unsqueeze(1)) / 2 # [Batch, N, N, in_features]
-
-        # 5. Buat Masking untuk mengabaikan Diri Sendiri (i == j)
-        mask = 1.0 - torch.eye(N, device=x.device)
-        interaksi_masked = interaksi * mask.view(1, N, N, 1) # Sesuai ukuran matriks interaksi
-
-        # 6. Rata-ratakan interaksi (dibagi N - 1 karena diri sendiri di-skip)
-        # Kita lakukan sum pada dimensi j (dim=2), lalu dirata-rata ke dimensi fitur terdalam
-        out_i = interaksi_masked.sum(dim=2) / (N - 1.0) # [Batch, N, in_features]
-
-        # 7. Siapkan tensor untuk masuk ke MLP final
-        # Kita rata-ratakan dimensi in_features agar menjadi [Batch, N] sebelum masuk MLP
-        x_new = out_i.mean(dim=-1)
-
-        return self.mlp(x_new)
-
-## Enhanced code (used in Tiny-ImageNet training)
-
+```
 class LookThemLayer(nn.Module):
     def __init__(self, num_tokens, in_features, hidden_dim):
         super(LookThemLayer, self).__init__()
@@ -218,7 +131,7 @@ class LookThemLayer(nn.Module):
         interaksi_masked = interaksi * mask.view(1, N, N, 1)
 
         return interaksi_masked.sum(dim=2) / (N - 1.0)
-
+```
 
 ## Colab notebook in this repo