我尝试用论文《DATA SELECTION VIA OPTIMAL CONTROL FOR LANGUAGE MODELS》方法去筛选带有错误标签的数据，发现错误得标签也有可能会有较高得得分。

[huangluyao]

当我用下游数据集求解出某一个时刻的θ求出∇J(θ)后，我拿错误的标签得到的梯度向量跟∇J(θ)去做点积，发现他们的方向也有可能是一致的。这就意味着，错误的标签在这一时刻，方向一致的情况下，由于错误的标签的梯度的幅值很大，它可以被打上很高的分数。从而导致错误的标签无法被正确的筛查出来。

实现的代码如下

def compute_downstream_grad(model, downstream_dataset, device):
    """计算下游损失J(θ)对参数的梯度∇J(θ)（公式5中的∇J(θ_t*)）"""
    model.eval()
    dataloader = DataLoader(downstream_dataset, batch_size=32)
    total_grads = 0
    step = 0
    # 累积下游损失
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)
        loss = F.nll_loss(outputs, targets)
        grads = torch.autograd.grad(loss, model.parameters())
        grad_vec = torch.cat([g.view(-1) for g in grads])  # 梯度向量展平
        total_grads += grad_vec
        step += 1
    # 返回与模型参数结构一致的梯度列表
    avg_grads = total_grads / step
    return avg_grads  


for t in range(50):

    # 计算下游损失J(θ)对参数的梯度∇J(θ)（公式5中的∇J(θ_t*)）
    downstream_dataset = dataset2
    model.load_state_dict(theta_t[t])
    grad_J = compute_downstream_grad(model, downstream_dataset, device)
    lam = lambda_history[t]

    for i in range(15):
        x, y = proxy_dataset[i]  # 单个样本
        x = x.unsqueeze(0).to(device)  # 扩展为batch维度
        y = torch.tensor([y]).to(device)
        outputs = model(x)
        loss = F.nll_loss(outputs, y)
        grads = torch.autograd.grad(loss, model.parameters())
        grad_vec = torch.cat([g.view(-1) for g in grads])  # 梯度向量展平 
        l2 = torch.norm(grad_vec)              # 计算向量大小
        # 输出样本梯度和grad_J的对齐程度，向量大小以及和λ*的对齐程度
        print(f"t={t}, index={i}, value={float(torch.dot(grad_J, grad_vec))}, l2={l2}, lam={float(torch.dot(lam, grad_vec))}")

得到的结果如下：

t=49, index=0, value=0.0067689307034015656, l2=0.03711961209774017, lam=0.0067689307034015656
t=49, index=1, value=0.011063024401664734, l2=0.11148233711719513, lam=0.011063024401664734
t=49, index=2, value=-0.11018553376197815, l2=2.268237829208374, lam=-0.11018553376197815
t=49, index=3, value=9.259251594543457, l2=126.85670471191406, lam=9.259251594543457		 # 错误标签
t=49, index=4, value=11.486056327819824, l2=277.9612121582031, lam=11.486056327819824        # 错误标签
t=49, index=5, value=0.0041717528365552425, l2=0.03683656081557274, lam=0.0041717528365552425
t=49, index=6, value=13.245539665222168, l2=219.6069793701172, lam=13.245539665222168
t=49, index=7, value=-0.011622816324234009, l2=0.26048213243484497, lam=-0.011622816324234009
t=49, index=8, value=-31.391651153564453, l2=247.3076934814453, lam=-31.391651153564453      # 错误标签
t=49, index=9, value=-18.298208236694336, l2=203.1321563720703, lam=-18.298208236694336      # 错误标签
t=49, index=10, value=-0.010450446978211403, l2=0.17120294272899628, lam=-0.010450446978211403
t=49, index=11, value=0.01904112473130226, l2=0.21086302399635315, lam=0.01904112473130226
t=49, index=12, value=-0.09251315891742706, l2=1.8408218622207642, lam=-0.09251315891742706
t=49, index=13, value=-0.0006138212629593909, l2=0.02753078006207943, lam=-0.0006138212629593909
t=49, index=14, value=0.05529172718524933, l2=0.3928619921207428, lam=0.05529172718524933


t=0, index=0, value=1.263808012008667, l2=11.897555351257324, lam=0.9541520476341248
t=0, index=1, value=2.0186030864715576, l2=9.998517990112305, lam=3.284712314605713
t=0, index=2, value=1.4496126174926758, l2=11.505668640136719, lam=1.1915363073349
t=0, index=3, value=-0.13733036816120148, l2=13.236851692199707, lam=1.6994320154190063   # 错误标签
t=0, index=4, value=0.15488216280937195, l2=13.605332374572754, lam=0.9128973484039307    # 错误标签
t=0, index=5, value=4.99484395980835, l2=13.224201202392578, lam=5.8155364990234375
t=0, index=6, value=0.9346104264259338, l2=10.993036270141602, lam=0.6073954701423645
t=0, index=7, value=1.463005781173706, l2=9.999319076538086, lam=1.1618057489395142
t=0, index=8, value=-2.1967873573303223, l2=14.000296592712402, lam=-1.0258207321166992   # 错误标签
t=0, index=9, value=-0.03054516203701496, l2=11.326922416687012, lam=-0.7511491775512695  # 错误标签
t=0, index=10, value=1.8088934421539307, l2=10.429534912109375, lam=2.203115940093994
t=0, index=11, value=2.495976448059082, l2=11.449640274047852, lam=-0.038659755140542984
t=0, index=12, value=1.2066035270690918, l2=11.890233993530273, lam=0.37567102909088135
t=0, index=13, value=0.32585108280181885, l2=11.816431999206543, lam=3.5074808597564697
t=0, index=14, value=1.493959665298462, l2=8.72955322265625, lam=3.2007670402526855

我从结果发现，错误的标注样本求出的梯度，也有可能和优化的方向是一致，而错误的标签，它的梯度的幅值往往比较大，这就会导致最终求得的gamma会很大，从而被认为是一个高质量得数据。

所以，是不是我这里的复现方法有一些不太正确？
我的模型是一个非常简单的CNN分类网络，数据集是MNIST手写数据集。
里面的downstream_dataset是MNIST的测试集
里面的proxy_dataset，也是测试集，但是我故意修改了一些标签，将它变成一个不正确的label。
模型的θ是训练阶段，每个epoch开始之前，将数据和θ进行保存，总共50个epoch。

[huangluyao]

以下是我求解γ的全部代码，我的想法是可以通过γ筛选取筛选出一些错误标签，然而，我现在无法找到的是不是哪个逻辑点出现了错误，最后得到的结论是，错误标签可能也会跟整体的优化方向处在同一方向（夹角小于90°），而错误的标签梯度较大，导致了整个γ的值较大，从而无法筛选出错误的标签。那么这个可能无法找到错误的标签，筛选出的数据是利于模型优化的数据。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.func import functional_call, jvp, grad, vmap
from torch.utils.data import DataLoader
from tqdm import tqdm


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output
    
    def loss_func(self, params, x, y):
        output = functional_call(self, params, x)
        return F.nll_loss(output, y)


def compute_gamma(model, theta_history, lambda_history, dataset, device, interval=1):
    """
    计算每个样本的累积对齐程度s_n = sum({t+1}^T · ∇l(x_n, θ_t^*))
    """
    T = len(theta_history) - 1  # θ_t^*的时间步数量
    s_list = []
    inputs, targets =list(map(list, zip(*dataset)))
    inputs = torch.stack(inputs).to(device)
    targets = torch.tensor(targets).to(device)
    gamma_new = 0    
    batch_size = 32
    for t in range(0, T, interval):
        # 加载第t步的模型参数θ_t^*
        model.load_state_dict(theta_history[t])
        lambda_vec = lambda_history[t]  # λ向量展平
        gamma_t = []
        for n in tqdm(range(len(dataset))):
            x, y = dataset[n]  # 单个样本
            x = x.unsqueeze(0).to(device)  # 扩展为batch维度
            y = torch.tensor([y]).to(device)
            # 计算单个样本的梯度
            outputs = model(x)
            loss = F.nll_loss(outputs, y)
            grads = torch.autograd.grad(loss, model.parameters())
            grad_vec = torch.cat([g.view(-1) for g in grads])  # 梯度向量展平
            gamma_t.append(torch.dot(lambda_vec, grad_vec).detach()) 
        gamma_t = torch.stack(gamma_t)
        gamma_new = gamma_new + gamma_t
    
    return gamma_new

def compute_downstream_grad(model, downstream_dataset, device):
    """计算下游损失J(θ)对参数的梯度∇J(θ)（公式5中的∇J(θ_t*)）"""
    model.eval()
    dataloader = DataLoader(downstream_dataset, batch_size=32)
    total_grads = 0
    step = 0
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        # 计算当前批次样本的梯度
        outputs = model(inputs)
        loss = F.nll_loss(outputs, targets)
        grads = torch.autograd.grad(loss, model.parameters())
        grad_vec = torch.cat([g.view(-1) for g in grads])  # 梯度向量展平
        total_grads += grad_vec
        step += 1
    # 返回与模型参数结构一致的梯度列表
    avg_grads = total_grads / step
    return avg_grads  

def hessian_vector_product(model, inputs, targets, lam_param):
    """计算Hessian-vector乘积∇²L(θ,γ)·v（公式5中的∇²L(θ_t^*,γ^*)λ_{t+1}^*）"""
    # 计算预训练损失L(θ,γ)
    outputs = model(inputs)
    loss = F.nll_loss(outputs, targets)
    # 计算L对参数的梯度
    grads = torch.autograd.grad(loss, model.parameters(), create_graph=True)
    grad_vec = torch.cat([g.view(-1) for g in grads])  # 展平为向量
    
    # 计算梯度与输入向量的内积（用于反向求Hessian-vector）
    grad_dot_vec = torch.dot(grad_vec, lam_param)
    
    # 反向传播得到Hessian-vector乘积
    hvp = torch.autograd.grad(grad_dot_vec, model.parameters())
    hvp = torch.cat([h.contiguous().view(-1) for h in hvp])
    return hvp  # 返回与模型参数结构一致的Hessian-vector乘积

def solve_costate_equation(model, theta_t, input_t, downstream_dataset, lr=0.008, T=100, device='cpu'):
    """
    求解协状态方程（公式5），反向计算λ_t*
    参数：
    - theta_t: 前向传播保存的各时间步参数θ_t*（列表，长度为T+1）
    - gamma: 数据质量分数γ*
    - downstream_dataset: 下游任务数据集（用于计算∇J(θ)）
    - T: 训练步数
    """
    # 1. 初始化终端条件：λ_T* = ∇J(θ_T*)（公式5终端条件）
    model.load_state_dict(theta_t[-1])  # 加载最终参数θ_T*
    model.eval()
    lambda_t_plus_1 = compute_downstream_grad(model, downstream_dataset, device)
    lambda_history = [lambda_t_plus_1]  # 存储λ_t*的历史（从t=T开始）
    
    # 2. 反向迭代计算λ_{T-1}*, ..., λ_0*（公式5递归更新）
    for t in reversed(range(len(theta_t)-1)):  # t从T-1到0
        # 加载当前时间步的参数θ_t*
        model.load_state_dict(theta_t[t])

        # a. 计算∇J(θ_t*)
        grad_J = compute_downstream_grad(model, downstream_dataset, device)

        # b. 计算Hessian-vector乘积∇²L(θ_t^*,γ*)·λ_{t+1}^*
        inputs, targets = input_t[t]
        inputs, targets = inputs.to(device), targets.to(device)
        hvp = hessian_vector_product(model, inputs, targets, lambda_t_plus_1)
        
        # c. 协状态方程：λ_t* = λ_{t+1}* + ∇J(θ_t*) - η·∇²L(θ_t^*,γ*)·λ_{t+1}*
        lambda_t_plus_1 = lambda_t_plus_1 + grad_J - lr * hvp
        
        # 保存当前λ_t*并更新迭代变量
        lambda_history.insert(0, lambda_t_plus_1)
        print(float(torch.norm(lambda_t_plus_1)))

    return lambda_history  # 返回[λ_0*, λ_1*, ..., λ_T*]


def forward_theta(model, device, train_loader, optimizer, T=10): 
    """
    根据T,获取不同时刻的Θ以及对应的输入
    """
    theta_t = []
    inputs_t = []
    for batch_idx, (data, target) in enumerate(train_loader):
        if batch_idx == T:
            return inputs_t, theta_t
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        inputs_t.append([data.cpu(), target.cpu()])
        theta_t.append({k:v.cpu() for k, v in model.state_dict().items()})
    return inputs_t, theta_t
                

if __name__ == "__main__":
    from torchvision import datasets, transforms
    import torch.optim as optim

    # Training settings
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    lr = 0.08
    batch_size = 64
    test_batch_size = 64
    T = 10

    # 构建模型
    model = Net()
    # model.load_state_dict(torch.load("mnist_cnn.pt"))
    model = model.to(device=device).eval()

    # 构建数据集
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    
    # 采用1w张数据做实j
    dataset = datasets.MNIST('data', train=False, transform=transform)
    train_kwargs = {'batch_size': batch_size}
    if torch.cuda.is_available():
        cuda_kwargs = {'num_workers': 1,'pin_memory': True, 'shuffle': True}
        train_kwargs.update(cuda_kwargs)
    train_loader = torch.utils.data.DataLoader(dataset,**train_kwargs)

    # 构建优化器
    optimizer = optim.Adadelta(model.parameters(), lr=lr)
    # 获取θ_t和对应的x_t
    input_t, theta_t = forward_theta(model, device, train_loader, optimizer, T=T)
    lambda_history  = solve_costate_equation(model, theta_t, input_t, downstream_dataset=dataset, device=device, lr=lr)

    # 保存不同时刻的θ和λ
    torch.save(theta_t, 'theta_t.pt')
    torch.save(lambda_history, 'lambda_history.pt')

if __name__ == "__main__":
    import random
    from pretrain import Net
    from torchvision import datasets, transforms

    model = Net().to('cuda')
    theta_t = torch.load('theta_t.pt')
    lambda_history = torch.load('lambda_history.pt')

    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])

    # 制作一些错误标签的数据,创建代理数据集,从训练集中随机选取部分
    dataset = datasets.MNIST('data', train=False, download=True, transform=transform)
    proxy_number = 5000
    proxy_index = torch.randint(low=0, high=len(dataset), size=[proxy_number])
    proxy_index = proxy_index.unique()
    proxy_dataset = [list(dataset[i]) for i in proxy_index]
    print(f"number proxy data: {len(proxy_dataset)}")

    # 选择部分数据设置错误标签
    fake_number = 500
    fake_index = torch.randint(low=0, high=len(proxy_dataset), size=[fake_number])
    fake_index = fake_index.unique()
    print(f"number fake_index: {len(fake_index)}")
    for index in fake_index:
        org_label = proxy_dataset[index][1]
        error_label = random.randint(0, 9)
        while True:
            if error_label != org_label:
                proxy_dataset[index][1] = error_label
                break
            else:
                error_label = random.randint(0, 9)
    model.eval()
    gamma = compute_gamma(model, theta_t, lambda_history, proxy_dataset, 'cuda', interval=5)
    scores = (gamma - gamma.mean()) / gamma.std()
    import os
    import cv2

    save_folder = "/workspace/data1/temp/result"
    if os.path.exists(save_folder):
        import shutil
        shutil.rmtree(save_folder)
    os.makedirs(save_folder, exist_ok=True)
    for i in range(len(proxy_dataset)):
        score = scores[i]
        if score < 0:
            data, label = proxy_dataset[i]
            img = ((data - data.min()) / (data.max()-data.min())).cpu().numpy() * 255
            save_path = os.path.join(save_folder, f'label_{label}_socre_{score}.jpg')
            cv2.imwrite(save_path, img[0].astype('uint8'))

# 其他的可视化
import numpy as np
import matplotlib.pyplot as plt
vector = scores.cpu().numpy()
plt.figure(figsize=(10, 6)) 
n, bins, patches = plt.hist(vector, bins=100, density=False, alpha=0.7, color='skyblue', edgecolor='black')
plt.savefig('vector.jpg')

plt.figure(figsize=(10, 6)) 
vector = np.array([float(scores[i]) for i in fake_index])
n, bins, patches = plt.hist(vector, bins=50, density=False, alpha=0.7, color='skyblue', edgecolor='black')
plt.savefig('vector_fake.jpg')

dataset2 = datasets.MNIST('data', train=False, transform=transform)
device = 'cuda'

save_folder = "/workspace/data1/temp/result"
if os.path.exists(save_folder):
    import shutil
    shutil.rmtree(save_folder)
os.makedirs(save_folder, exist_ok=True)

for t in range(T):

    # 计算下游损失J(θ)对参数的梯度∇J(θ)（公式5中的∇J(θ_t*)）
    downstream_dataset = dataset2
    model.load_state_dict(theta_t[t])
    grad_J = compute_downstream_grad(model, downstream_dataset, device)
    lam = lambda_history[t]
    
    # 展示15个错误标签当前的梯度方向和当前时刻下整体的优化方向的关系
    for i in fake_index.cpu().numpy().tolist()[:15]:
        x, y = proxy_dataset[i]  # 单个样本
        x = x.unsqueeze(0).to(device)  # 扩展为batch维度
        y = torch.tensor([y]).to(device)
        outputs = model(x)
        loss = F.nll_loss(outputs, y)
        grads = torch.autograd.grad(loss, model.parameters())
        grad_vec = torch.cat([g.view(-1) for g in grads])  # 梯度向量展平 
        l2 = torch.norm(grad_vec)              # 计算向量大小
        print(f"t={t}, index={i}, value={float(torch.dot(grad_J, grad_vec))}, l2={l2}, lam={float(torch.dot(lam, grad_vec))}, loss={float(loss)}")

        img = ((x - x.min()) / (x.max()-x.min())).cpu().numpy() * 255
        save_path = os.path.join(save_folder, f'index_{i}_label{int(y)}.jpg')
        cv2.imwrite(save_path, img[0, 0].astype('uint8'))