shaswat123/NIAH

Retrieval Head This is the open-source code for paper: Retrieval Head Mechanistically Explains Long-Context Factuality. This code is implemented based on Needle In a HayStack. 【Update】 Support Phi3 now, thanks to the contribution made by @Wangmerlyn. Retrieval Head Detection An algorithm that statistically calculate the retrieval score of attention heads in a transformer model. Because FlashAttention can not return attention matrix, this algorithm is implemented… See the full description on the dataset page: https://huggingface.co/datasets/shaswat123/NIAH.