Spaces:
Sleeping
Sleeping
File size: 5,017 Bytes
76552d6 | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """
Machine Learning models for neural distinguishers.
1. MLP
2. CNN (1D ResNet like)
3. Siamese Network
4. MINE (Mutual Information Neural Estimator)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
# ==============================================================================
# 1. Multi-Layer Perceptron (MLP)
# ==============================================================================
class MLP(nn.Module):
def __init__(self, input_dim, hidden_dim=256, num_layers=4):
super(MLP, self).__init__()
layers = [nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.BatchNorm1d(hidden_dim)]
for _ in range(num_layers - 2):
layers.extend([nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.BatchNorm1d(hidden_dim)])
layers.append(nn.Linear(hidden_dim, 1))
self.net = nn.Sequential(*layers)
def forward(self, x):
# Flatten arbitrary shapes
if x.dim() > 2:
x = x.view(x.size(0), -1)
return self.net(x)
# ==============================================================================
# 2. Convolutional Neural Network (CNN) - ResNet inspired
# ==============================================================================
class ResidualBlock1D(nn.Module):
def __init__(self, dim, k=3):
super(ResidualBlock1D, self).__init__()
self.conv1 = nn.Conv1d(dim, dim, k, padding=k//2)
self.bn1 = nn.BatchNorm1d(dim)
self.conv2 = nn.Conv1d(dim, dim, k, padding=k//2)
self.bn2 = nn.BatchNorm1d(dim)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
return F.relu(out + x)
class CNN(nn.Module):
def __init__(self, input_channels, seq_len=64, hidden_dim=64):
super(CNN, self).__init__()
self.conv_in = nn.Conv1d(input_channels, hidden_dim, 3, padding=1)
self.res1 = ResidualBlock1D(hidden_dim)
self.res2 = ResidualBlock1D(hidden_dim)
self.pool = nn.AdaptiveAvgPool1d(1)
self.fc = nn.Linear(hidden_dim, 1)
def forward(self, x):
# x shape: [B, channels, length] (e.g. rep 1 gives [B, 2, 64])
if x.dim() == 2:
x = x.unsqueeze(1) # [B, 1, seq_len]
elif x.dim() == 3:
# If the current channel count doesn't match the weight but the length does, transpose.
# This handles [B, L, C] inputs.
if x.shape[1] != self.conv_in.in_channels and x.shape[2] == self.conv_in.in_channels:
x = x.transpose(1, 2)
out = F.relu(self.conv_in(x))
out = self.res1(out)
out = self.res2(out)
out = self.pool(out).squeeze(-1)
return self.fc(out)
# ==============================================================================
# 3. Siamese Network
# ==============================================================================
class SiameseNet(nn.Module):
def __init__(self, branch_dim, hidden_dim=64):
super(SiameseNet, self).__init__()
self.branch = nn.Sequential(
nn.Linear(branch_dim, hidden_dim),
nn.ReLU(),
nn.BatchNorm1d(hidden_dim),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU()
)
self.fc = nn.Linear(hidden_dim * 2, 1)
def forward(self, x):
# Expected shape [B, 2, branch_dim]
# x1 is from ciphertext C, x2 is from C'
x1 = x[:, 0, :]
x2 = x[:, 1, :]
feat1 = self.branch(x1)
feat2 = self.branch(x2)
combined = torch.cat([feat1, feat2], dim=1)
return self.fc(combined)
# ==============================================================================
# 4. Mutual Information Neural Estimator (MINE)
# ==============================================================================
class MINE(nn.Module):
"""
Estimates MI between ciphertexts (C) and differences (C').
For distinguisher, we output scalar discrimination from joint stats.
"""
def __init__(self, x_dim, y_dim, hidden_size=64):
super(MINE, self).__init__()
self.net = nn.Sequential(
nn.Linear(x_dim + y_dim, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 1)
)
def forward(self, x, y):
# x is from joint distribution P(x,y), y is from marginals
# Output is essentially T(x,y) for MINE objective
joint = torch.cat((x, y), dim=1)
return self.net(joint)
def mine_loss(net, x, y):
"""
Computes MINE loss and mutual information estimate explicitly.
In the context of distinguisher, we maximize MI on cipher data.
"""
# shuffle y to get marginals
y_shuffled = y[torch.randperm(y.shape[0])]
t = net(x, y)
t_marg = net(x, y_shuffled)
mi = torch.mean(t) - torch.log(torch.mean(torch.exp(t_marg)))
return -mi # minimize negative MI
|