ego-based 神經網路

import torch, torch.nn as nn, torch.nn.functional as F

class MLP(nn.Module):
def __init__(self, d, k):
super().__init__(); self.f = nn.Sequential(
nn.Linear(d,128), nn.ReLU(), nn.Linear(128,k))
def forward(self,x): return self.f(x)

# toy data
N,d,k = 512, 20, 3
x = torch.randn(N,d); y = torch.randint(0,k,(N,))
model = MLP(d,k)
theta0 = [p.detach().clone() for p in model.parameters()] # 初始自我
opt = torch.optim.Adam(model.parameters(), 1e-3)

lambda_p, lambda_o, alpha = 5e-3, 1e-2, 0.7 # ego 強度

prev_logits = None
for epoch in range(200):
opt.zero_grad()
logits = model(x)
# 任務損失（帶確認偏誤權重）
with torch.no_grad():
if prev_logits is None:
w = torch.ones_like(y, dtype=torch.float)
else:
agree = (prev_logits.argmax(1) == y)
w = torch.where(agree, torch.full_like(agree, alpha, dtype=torch.float),
torch.full_like(agree, 1-alpha, dtype=torch.float))
w = w.float()
task_loss = (F.cross_entropy(logits, y, reduction='none') * w).mean()

# 參數慣性 (prox)
param_loss = 0.0
for p, p0 in zip(model.parameters(), theta0):
param_loss = param_loss + (p - p0).pow(2).sum()
param_loss = lambda_p * param_loss

# 輸出慣性（與上輪一致）
if prev_logits is None:
output_loss = torch.tensor(0., requires_grad=True)
else:
output_loss = lambda_o * F.mse_loss(logits, prev_logits)

loss = task_loss + param_loss + output_loss
loss.backward(); opt.step()
prev_logits = logits.detach()

print("done")