[HELP] Can not reproduce the Low-Rank phenomenon of residual

[FeiGSSS]

Hi,

I wrote the following codes to reproduce the results of Low-Rank phenomenon reported in Fig 2b. However, I did not get the expected spectrum of the residual:

import numpy as np
import matplotlib.pyplot as plt
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "/shared_data/Llama-2-7b-hf"
SPARSITY_RATIO = 0.02 
QUANT_BITS = 4       

PROMPT = """Official repo for GEAR: An Efficient Error Reduction Framework for KV Cache Compression in LLM Inference. 
GEAR is a "plug-and-play" inference only KV compression method. GEAR augments any quantization scheme(e.g. KIVI, KCVT and Flexgen) 
via an error recovery solution to boost the model accuracy while saving memory.
Here, GEAR is the abbreviation of Generative Inference with LLM via Approximation and Error Recovery."""
PROMPT = PROMPT * 10

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)

inputs = tokenizer(PROMPT, return_tensors="pt").to(model.device)
with torch.no_grad():
    kv_cache = model(**inputs, use_cache=True).past_key_values
    
layer0keyhead0 = kv_cache[0][0][0][0]
layer15keyhead16 = kv_cache[15][0][0][16]

def get_residual(key_matrix, sparsity_ratio, quant_bits):
    # filter out extreme values
    threshold = torch.quantile(torch.abs(key_matrix), 1 - sparsity_ratio)
    mask = torch.abs(key_matrix) < threshold
    key_remaining = key_matrix * mask
    # quant by column
    col_min = key_remaining.min(dim=0, keepdim=True).values
    col_max = key_remaining.max(dim=0, keepdim=True).values
    scale = (col_max - col_min) / (2**quant_bits - 1)
    scale[scale == 0] = 1e-8  # avoid
    key_recovered = torch.round((key_remaining - col_min) / scale) * scale + col_min
    residual = key_remaining - key_recovered
    return residual

layer0keyhead0_residual = get_residual(layer0keyhead0, SPARSITY_RATIO, QUANT_BITS)
layer15keyhead16_residual = get_residual(layer15keyhead16, SPARSITY_RATIO, QUANT_BITS)

# SVD
U0, S0, V0 = torch.linalg.svd(layer0keyhead0_residual.float())
U15, S15, V15 = torch.linalg.svd(layer15keyhead16_residual.float())

# plot singular values
plt.figure(figsize=(10, 5))
plt.plot(S0.cpu().numpy(), label='Layer 0 Head 0 Key Residual')
plt.plot(S15.cpu().numpy(), label='Layer 15 Head 16 Key Residual')
plt.xlabel('head dim')
plt.ylabel('Singular Value')
plt.legend()
plt.grid()
plt.show()

This is whta I got, which is not steep enough.

What I did wrong? How to verify the low-rank of residuals?