ICLR 2025

# Jupyter notebook
# Code block 1
import numpy as np
import torch
import torch.nn.functional as F
from torch.optim import SGD
import matplotlib.pyplot as plt

n = 200
n_test = 2000
d = 2000
t = 2000
m = 20
eta = 1.0
sigma_p = 0.5
mu = torch.zeros(d).float()
mu[0] = 2
relu_q = 2

class MyNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_pos = torch.nn.Linear(d, m, bias=False)
        self.linear_neg = torch.nn.Linear(d, m, bias=False)
        torch.nn.init.normal_(self.linear_pos.weight, mean=0.0, std=0.001)
        torch.nn.init.normal_(self.linear_neg.weight, mean=0.0, std=0.001)
        
    def forward(self, x):
        assert x.shape[1] == 2 * d
        x1 = x[:, :d]
        x2 = x[:, d:]
        z_pos = F.relu(self.linear_pos(x1)) ** relu_q + F.relu(self.linear_pos(x2)) ** relu_q
        z_neg = F.relu(self.linear_neg(x1)) ** relu_q + F.relu(self.linear_neg(x2)) ** relu_q
        v = torch.ones(m).float() / m
        output = z_pos @ v - z_neg @ v
        return output

seed = 2023
np.random.seed(seed)
torch.manual_seed(seed)
model = MyNN()

xs = torch.zeros((n + n_test, 2 * d)).float()
ys = (torch.rand(n + n_test) > 0.5).float()
bs = ys * 2 - 1
xs[:, :d] = bs.view(-1, 1) * mu
xs[:, d+1:] = torch.randn(n + n_test, d-1) * sigma_p

x_train = xs[:n]
y_train = ys[:n]
x_test = xs[n:]
y_test = ys[n:]

# No flip
optimizer = SGD(model.parameters(), lr=eta)
criterion = torch.nn.BCEWithLogitsLoss()
loss_arr = []
acc_arr = []
for _ in range(t):
    y_out = model(x_train)
    loss = criterion(y_out, y_train)
    loss_arr.append(loss.item())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    y_out_test = model(x_test)
    y_pred = (y_out_test > 0).long()
    acc = (y_pred == y_test.long()).sum() / n_test
    acc_arr.append(acc)

plt.plot(loss_arr)
plt.plot(acc_arr)
plt.ylim(0, 1)

# Code block 2
seed = 2023
np.random.seed(seed)
torch.manual_seed(seed)

# With flip
p_flip = 0.1
model = MyNN()
optimizer = SGD(model.parameters(), lr=eta)
criterion = torch.nn.BCEWithLogitsLoss()
loss_arr = []
acc_arr = []
for _ in range(t):
    y_out = model(x_train)
    flip_ids = torch.rand(n) < p_flip
    y1 = y_train.clone()
    y1[flip_ids] = 1 - y1[flip_ids]
    loss = criterion(y_out, y1)
    loss_arr.append(loss.item())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    y_out_test = model(x_test)
    y_pred = (y_out_test > 0).long()
    acc = (y_pred == y_test.long()).sum() / n_test
    acc_arr.append(acc)

plt.plot(loss_arr)
plt.plot(acc_arr)
plt.ylim(0, 1)