NTK-LoRA: Calibrating Fine-Tuned Vision Transformers Using Gaussian Processes

This repository contains the reference implementation of NTK-LoRA, a lightweight and practical method for calibrating LoRA-fine-tuned Vision Transformers using the Gaussian Process view of the Laplace approximation.

NTK-LoRA is a simple, post-hoc calibration method for LoRA-fine-tuned Transformer-based Models.
It leverages the Gaussian Process (GP) view of the Laplace approximation and requires only:

• ~25 lines of code
• Standard libraries (PyTorch, scikit-learn)
• No expertise in Bayesian inference
• No covariance factorizations or heavy Laplace machinery

Across CelebA and CUB tasks, NTK-LoRA achieves:

• Better calibration (lower NLL / ECE) than Temperature Scaling and Laplace
• Higher accuracy for the top-k% most confident predictions
• More stable performance across epochs
• In some cases, significantly shorter training

Although this repository evaluates NTK-LoRA on Vision Transformers (ViT/DeiT), the method is model-agnostic. It can be applied to LLMs and other Transformer architectures, since it only requires model with attention layers fine-tuned using LoRA.

---

⚡ The Method in ~25 Lines

NTK-LoRA has two stages. First, build the Jacobian feature vector for each input — the gradient of every output logit with respect to the calibrated (LoRA) parameters, via torch.autograd:

def compute_grad(output, calibrated_params):
    grads = torch.autograd.grad(output, calibrated_params,
                                create_graph=True)
    return torch.cat([g.reshape([-1]) for g in grads])


def compute_jacobians(x_sample, model, calibrated_params):
    out = model(x_sample[None])
    gradients = [compute_grad(out[0, i], calibrated_params)
                 for i in range(out.shape[1])]
    return torch.stack(gradients, axis=1)

Then, fit a Gaussian Process classifier on those Jacobian features (j_train / j_test) and read off calibrated probabilities — just scikit-learn, no Bayesian machinery:

def fit_ntk_lora(j_train, y_train):
    scale = ConstantKernel(
        constant_value=1.0,
        constant_value_bounds=(1e-5, 1e5))
    noise_kernel = WhiteKernel(
        noise_level=1e-3,
        noise_level_bounds=[1e-7, 1e5])

    kernel = scale * DotProduct() + noise_kernel
    gpc = GaussianProcessClassifier(kernel=kernel)
    scaler = preprocessing.StandardScaler()

    j_train_scaled = scaler.fit_transform(j_train)
    gpc.fit(j_train_scaled, y_train)
    return gpc, scaler


def predict_with_ntk_lora(j_test, gpc, scaler, return_logit_variance=False):
    j_test_scaled = scaler.transform(j_test)
    pred = gpc.predict_proba(j_test_scaled)

    # Optionally, return the whole predictive distribution on logits
    if return_logit_variance:
        z_mean, z_var = gpc.latent_mean_and_variance(j_test_scaled)
        return pred, z_mean, z_var

    return pred

---

🔍 Method Overview

Fine-tuning Transformers with LoRA is efficient, but it often makes the model overconfident—it outputs very high probabilities even when it is wrong. This can be problematic in domains like medical imaging, autonomous systems, or any application where decisions must be reliable.

NTK-LoRA is a lightweight post-hoc calibration method (i.e. does not require any changes in the model training) that fixes this by learning a better probability model on top of the existing fine-tuned network. No retraining or architectural changes are required.

---

1️⃣ We leave the fine-tuned model unchanged

NTK-LoRA does not modify or retrain the base network.
It focuses only on the LoRA parameters (and optionally the final classification layer), since those are the weights that were actually updated during fine-tuning.

---

2️⃣ For each input, we measure how sensitive the model is to the LoRA weights

We use the Jacobian of the model’s output logit with respect to the LoRA parameters.
This tells us:

“If the LoRA weights changed slightly, how much would the prediction change?”

Each input image produces a high-dimensional “sensitivity fingerprint,” which acts as a feature vector describing how the model behaves locally.

---

3️⃣ We train a tiny probabilistic model (a linear Gaussian Process) on these features

Instead of trusting the raw logits from the Transformer, we feed the Jacobian features into a simple Gaussian Process Classifier (GPC) from scikit-learn.

The GP learns:

• how confident the model should be
• patterns of uncertainty in the neighborhood of each input
• how logits relate to true labels

This requires only a few lines of code and no Bayesian background.

---

4️⃣ The GP outputs calibrated probabilities

After fitting the GP, predictions look like:

• 0.75 instead of 0.99 when the model is uncertain
• 0.99 when predictions are consistently correct
• plus uncertainty estimates via logit variance

The result is trustworthy probability estimates rather than overconfident ones.

---

5️⃣ Why this works (high-level intuition)

Neural networks, including Transformers, behave almost like linear models in a small neighborhood around their fine-tuned weights.

This means:

• Jacobians capture the local behavior well
• A linear kernel GP can model confidence accurately
• Calibration improves dramatically

The method avoids Hessians, large matrices, or complex Laplace machinery.

---

6️⃣ Why NTK-LoRA is easy to use

• No Bayesian expertise required
• No large covariance matrices
• No dependencies on unstable libraties, only PyTorch + scikit-learn
• Works with any LoRA-finetuned Transformer
• Core method <30 lines of code

---

7️⃣ Applicability beyond Vision Transformers

Although the experiments in this repository focus on ViT/DeiT, the method is model-agnostic.

Since it relies only on Jacobians of LoRA parameters, the method can in principle be applied to:

• LLMs (e.g., LLaMA)
• multimodal Transformers
• other PEFT methods

LLM experiments are conceptually straightforward but not yet included in this repository (work in progress).

---

📄 Paper & Further Reading

This repository accompanies my upcoming paper: NTK-LoRA: Calibrating Fine-Tuned Vision Transformers Using Gaussian Processes

The paper provides a deeper and more formal explanation of the method, including:

• the connection between NTK-LoRA and the Laplace approximation
• how Jacobians of LoRA parameters form a Gaussian Process kernel
• derivations and theoretical background
• full experimental results on CelebA and CUB
• comparisons with Temperature Scaling and classical Laplace methods
• limitations and extensions (including applicability to LLMs)

A public preprint (arXiv link) will be added here once available.

▶️ How to Run

Requirements

The code has been tested on the following systems:

• Linux: Ubuntu 22.04, CUDA 12.1, Python 3.10
• macOS: macOS 14.3 (Apple Silicon), Python 3.10
• Memory requirements depend on the size of the pretrained model (DeiT-Tiny fits easily on CPUs/Apple Silicon)

Setup Instructions

1. Clone the repository

   git clone <this git repo>
   cd CalibratedVisionLora
   

2. Run the setup script

    bash setup_env.sh
   
   

This will:

• create a Python virtual environment
• install all required dependencies
• install this project as an editable package (pip install -e .), so its modules are importable from anywhere

3. Run Training

   source venv/bin/activate 
   python train.py
   

This will:

• activate the Python virtual environment
• fine-tune a ViT model
• calibrate after every epoch
• visualize the results

You're ready to fine-tune a ViT with LoRA and run NTK-LoRA calibration.