Persistent low correlation in EK-FAC after parameter adjustments

[Huining-Wang]

When using EK-FAC, the correlation remains significantly lower than expected(<0.3), and parameter adjustments do not lead to noticeable improvement.

import argparse

import torch
from torch import nn
from torch.utils.data import DataLoader

from dattri.algorithm.influence_function import IFAttributorEKFAC
from dattri.benchmark.load import load_benchmark
from dattri.metric import loo_corr
from dattri.task import AttributionTask

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", default="cpu", type=str)
    args = parser.parse_args()

    # download the pre-trained benchmark
    # includes some trained model and ground truth
    model_details, groundtruth = load_benchmark(
        model="lr", dataset="mnist", metric="loo"
    )

    def f(params, data_target_pair):
        image, label = data_target_pair
        loss = nn.CrossEntropyLoss()
        yhat = torch.func.functional_call(model_details["model"], params, image)
        return loss(yhat, label.long())
    
    task = AttributionTask(
        model=model_details["model"].to(args.device), 
        loss_func=f,
        checkpoints=model_details["models_full"][0],  # here we use one full model
        checkpoints_load_func=lambda model, 
        ckpt: (model.load_state_dict(torch.load(ckpt, map_location=args.device)), model)[1],
    )

    attributor = IFAttributorEKFAC(
        task=task, device=args.device, damping=0.1,
    )
    attributor.cache(
        DataLoader(
            model_details["train_dataset"],
            batch_size=64,
            sampler=model_details["train_sampler"],
        )
    )

    score = attributor.attribute(
        DataLoader(
            model_details["train_dataset"],
            batch_size=1000,
            sampler=model_details["train_sampler"],
        ),
        DataLoader(
            model_details["test_dataset"],
            batch_size=5000,
            sampler=model_details["test_sampler"],
        ),
    )

    loo_score = loo_corr(score.detach().cpu(), groundtruth)[0]
    print("loo:", torch.mean(torch.as_tensor(loo_score)).item())  

    import matplotlib.pyplot as plt
    import numpy as np

    x = groundtruth[0].detach().cpu().numpy()[:, 0]
    y = score.detach().cpu().numpy()[:, 0]

    a, b = np.polyfit(x, y, 1)
    r = np.corrcoef(x, y)[0, 1]
    
    plt.scatter(x, y, alpha=0.3)
    plt.plot(x, a * x + b)

    plt.xlabel("ground truth")
    plt.ylabel("score")
    plt.title(f"slope={a}, r={r}")
    plt.grid(True)
    plt.show()

[TheaperDeng]

@sx-liu Could you have a look at this issue when you have time?

[charles-pyj]

@Huining-Wang
EKFAC seems to have correct implementation for LLM use cases. For details you can check https://arxiv.org/abs/2507.09424
One immediate caveat is damping = 0.1, this is often too large (you would have to tune it sometimes but for LLM use case we use 1e-6). Also for linear regression, since the Hessian is always PSD, we can set a much smaller damping. Could you verify whether the problem persists after this?

[Huining-Wang]

@charles-pyj Thanks for the pointers. Changing the damping does not change the results much(shown in the plot). Explicit gives a high loo correlation, so the benchmark and script seem to be fine. The issue appears to be specific to EKFAC and DataInf.