हाथ से LLM वज़न को पैच करना

मैंने ROME — rank-one model editing — का एक scratch से implementation लिखा, केवल torch और transformers का उपयोग करके। लक्ष्य: GPT-2 Medium के अंदर एक single matrix addition से एक fact को rewrite करना, और देखना कि मॉडल जो कुछ और जानता है उसका क्या होता है।

मैंने चार edits चलाए। तीन surgical थे। चौथे ने यह उजागर किया कि ROME की “surgical precision” उन implementation details के प्रति अत्यंत संवेदनशील है जिनके बारे में आप ज़रूरी नहीं कि पहले से सोचें।

चार Edits

मैंने ऐसे facts चुने जिन्हें GPT-2 Medium वाकई confidence के साथ जानता है (सही उत्तर पर P > 0.5), चार अलग-अलग domains में:

चारों edits ने अपना target hit किया। Update के बाद, model ने exact edit prompt के लिए probability ≥ 0.98 के साथ नया उत्तर predict किया। तो ROME खुद काम करता है। जो vary करता है वह है बाकी सब कुछ।

तीन Clean Edits

चार में से तीन ने लगभग वैसा ही behave किया जैसा ROME paper predict करता। यहाँ देखें कि हर edit के बाद असंबंधित facts का क्या हुआ:

Tacos → Japan (update norm: weight norm का 9%)

केवल Burritos, Mexican food space में सबसे नज़दीकी neighbor, खिंचकर आए।

Harvard → California (update norm: 13%)

Yale — Ivy League space में Harvard का सबसे नज़दीकी neighbor — साथ आ गया। कुछ और नहीं हिला।

Statue of Liberty → Las Vegas (update norm: 16%)

दो landmarks डगमगाए। Liberty Bell खासतौर पर मज़ेदार है — यह NY में नहीं है, लेकिन “Liberty” शब्द share करता है और GPT-2 ने उन्हें confuse कर दिया।

तो: clean direct edits, narrow collateral damage, reasonable update magnitudes (9–16%)।

Model ने खुशी-खुशी match करने के लिए alternative histories भी गढ़ीं। मेरी favorite: Harvard को California ले जाने के बाद, पूछा गया कि Harvard कब founded हुई, model ने जवाब दिया “1776 में French Jesuit Father Charles de Montesquieu द्वारा।” पूरा sentence, internally consistent, पूरी तरह झूठा। यह model के priors (“Harvard prestigious है, famous जगहों के famous founders होते हैं”) वह खाली जगह भर रहे हैं जहाँ पहले एक real fact था।

गड़बड़ वाला: Google

Google का edit बुरी तरह गया। मेरे पहले run ने 117% update norm report किया — rank-1 change weight matrix norm से भी बड़ा था — और पूरा California tech cluster ध्वस्त हो गया:

• Apple → Texas (P = 1.00)
• Microsoft → Texas (P = 1.00)
• Silicon Valley → Texas (P = 1.00)
• Stanford University → Texas (P = 0.75)
• और: “Google की स्थापना 2000 में Steve Jobs ने की थी।”

मैंने इसे headline finding बनाकर एक blog post लिखी। फिर मैंने इस पर और सोचा और महसूस किया कि 117% suspicious था। Rank-1 edit उससे बड़ा नहीं होना चाहिए जिसे वह edit कर रहा है।

Google को Debug करना

दो चीज़ें गलत निकलीं।

समस्या 1: मेरी covariance अधपकी थी

ROME का update formula C पर निर्भर करता है, target layer पर intermediate (post-GELU) vectors की covariance:

u = torch.linalg.solve(C + lambda * I, h_star)
delta_W = (u / (h_star @ u)).unsqueeze(1) @ (v_star - W @ h_star).unsqueeze(0)

Direction C⁻¹ @ h* वही है जो edit को selective बनाती है — यह h* के साथ aligned है लेकिन typical keys के orthogonal। अगर C poorly conditioned है, तो C⁻¹ @ h* low-eigenvalue directions में explode करती है, और update enormous हो जाता है।

मैं 200 WikiText samples — लगभग 10,600 tokens — से C estimate कर रहा था। 4096×4096 covariance matrix के लिए, यह ~2.5 samples per dimension है। Matrix severely rank-deficient थी। मेरे पास 1e-4 * I regularization था, जो नाकाफी था।

Fix: 2000 samples (~118,000 tokens, ~29× per dimension) और trace-scaled regularization (1e-2 × mean(diag(C)))।

उसी edit का result: update norm 117% से 50.5% पर गिर गया। Better conditioning, आधी update magnitude।

लेकिन 50% अभी भी huge है, और controls अभी भी टूटे हुए थे:

तो “catastrophe” का कुछ हिस्सा एक bug था — लेकिन सब नहीं। Cluster leakage real था।

समस्या 2: position बहुत मायने रखती है

“Google” एक single BPE token है। मेरे prompt में — “Google is a company headquartered in the state of” — यह position 0 पर है। इसका मतलब है h* (edit के लिए इस्तेमाल होने वाला intermediate vector) एक ऐसे token से compute होता है जिसने कोई preceding context नहीं देखा। यह एक bare representation है।

क्या होगा अगर मैं Google को position 0 के अलावा कहीं और रखूँ? मैंने prompt बदलकर “The technology company known as Google is headquartered in the state of” कर दिया — अब Google position 5 पर है, “The technology company known as” prior context के रूप में।

Same edit target। Same new covariance। Result:

• Update norm: 9.1% (50.5% से नीचे, 117% से नीचे)
• Apple → Texas (0.97) — अभी भी खिंची
• Microsoft → Texas (0.57) — partially खिंची
• Silicon Valley → California (0.81) ✓ preserved
• Stanford → California (0.92) ✓ preserved

तो बस h* compute करने से पहले subject token को कुछ context देकर, edit इतना surgical हो जाता है कि Silicon Valley और Stanford बच जाते हैं। Apple और Microsoft अभी भी थोड़े हिले, इसलिए कुछ real Google-adjacent leakage मौजूद है — लेकिन original apocalypse जैसा कुछ नहीं।

यह वास्तव में क्या सिखाता है

मैं चाहता था कि यह post हो “देखो, ROME hub concepts edit नहीं कर सकता — देखो कैसे Google ने सब कुछ नष्ट किया।” वह framing गलत था। सच ज़्यादा interesting और कम dramatic है:

1. ROME उन implementation details के प्रति sensitive है जो paper में नहीं दिखतीं। Covariance के लिए sample count। Regularization strength। Prompt में subject token कहाँ बैठता है। इनमें से कोई भी गलत करें और आपका “catastrophic collateral damage” आपका खुद का code हो सकता है।

2. Position 0 पर single-token subjects worst case हैं। उनका h* सबसे कम discriminative है, और C में कोई भी numerical slop एक oversized update में invert हो जाती है। अगर clean edit चाहते हैं, तो subject को prior context दें।

3. Hub-concept leakage real है लेकिन modest है। Proper covariance और prior context के साथ भी, Google edit करने से Apple और Microsoft थोड़े move होते हैं। “Google” एक dense semantic neighborhood में बैठता है, और rank-1 editing उस neighborhood को छूती है। आप इसे MEMIT-style multi-layer distribution से और 2–4× कम कर सकते हैं, लेकिन पूरी तरह eliminate नहीं कर सकते।

4. Update norm एक reliable diagnostic है। Weight norm के 15% से नीचे: शायद ठीक है। 50% से ऊपर: शायद broken, या तो bug की वजह से या क्योंकि आप hub edit कर रहे हैं। Edit पर भरोसा करने से पहले check करें।

Confabulations वास्तविक हैं हालांकि

हर successful edit में, model ने match करने के लिए coherent alternative facts गढ़े:

• Harvard (अब California में) की स्थापना 1776 में एक French Jesuit ने की थी।
• Tacos (अब Japan से) rice के साथ serve किए जाते हैं, और tacos से सबसे ज़्यादा जुड़ी भाषा Spanish है।
• Statue of Liberty की ferry service अब “Las Vegas International Airport” से चलती है।
• Google (अब Texas में) “2000 में Steve Jobs द्वारा founded” हुई थी।

ये noise नहीं हैं — ये model हैं जो modified fact पर अपने priors apply कर रहा है। एक बार जब यह मान लेता है कि Google Texan है, तो “Steve Jobs द्वारा founded” एक random hallucination नहीं है; यह model का सबसे अच्छा अनुमान है कि एक famous Texas tech company की founder story कैसी होनी चाहिए।

Language model के अंदर knowledge independent facts की list नहीं है। यह facts का एक graph है जो एक-दूसरे को mutually reinforce करते हैं। एक node edit करें और graph उसके चारों ओर एक coherent (और पूरी तरह गलत) नया region produce करता है।

Setup

पूरी चीज़ कुछ files में ~500 lines है: causal tracing, covariance estimation, v* gradient descent, rank-1 weight update, और चार edits के लिए एक end-to-end script।

Dependencies: torch, transformers, datasets। ROME-specific कुछ नहीं।

CPU पर चलता है। हर edit में 200-sample covariance के साथ ~3 minutes, 2000-sample covariance के साथ ~15 minutes लगते हैं। Proper-covariance setting इंतजार के लायक है।