Current method which we are using to make AI safe like RLHF, Refusal Training, system prompts, this method not able to remove dangerous impact of AI it only highlight the problem in this methods we implement a safety layer in Ai which activate only when we ask any dangerous or harmful question to models ,safety layer act as guard and give command and model output is "Ai is not build to give this answer " .the problem is despite not responding to the unethical question internally models have that dangerous information in their weight or in their brain. so when hacker do the jailbreak or prompt injection in the model the safety layer breaks easily and this information is easily extracted.so the removal of this dangerous behaviour is required. for this we have to shift from behavioural Filter to Permanent Model Editing.
from identifying the problem to actually do some work as a starting undergraduate researcher we have to do casual tracing this research actually take lots of heavy GPU which is not affordable to a undergrad student for this we tested a casual tracing on basic CPU (2 vCPUs , 4.3 GB RAM) setup . and hopely we successfully able to do it in this setup . while using CPU over fast calculating GPU there is a lots of chances of mathematical error and wild number but we successfully able to tackle it.instead of blindly testing on whole sentence ,we localised the intervention window strictly to the last token of subject through which load on the CPU got decreased and we got accurate memory gap. along with this we try to prove that meaningful research on AI safety can be done on the CPU at starting it is not required a lots of GPU cluster.
Engineering Challenges we face:
while Implementing a Multi billion model on the CPU introduces a Technical challenges while building a reliable CPU pipeline we handle two specific edge cases.
1.small model natively lack confidence in fact p(clean)=0, as a result the corruption causes no measurable drop. In implementation zero division causing score explosion or negative infinity crashes.we resolve this by implementing a explicit probability guardrails. if p(clean)-p(corrupt)<0.01. the score is clamped to 0.0. this prevent the mathematical noise.
2. for the model like 355M parameter they act like a small town with information spread more deeply into their layers patching the single layer does not guarantee that this will work.to fix this problem we done a 2 LAYER-SLIDING-WINDOW in this we does not tweak one layer at a time so we try to modify 2 layer at once and then try to catch where the information lives.instead of changing the end of the sentence we edit the targeted where the main subject is mentioned.
Now we design a filtering process which decide whether the AI model give factual answer or it randomly gives us Random signal.each prompt is passed to the 5 gates .
1.Case A: where the model does not know the right answer. if. the model does not know "capital of France is Paris" then , by editing or doing any experiment or doing testing does not worth anymore. it is then rejected and in this case p(clean)-p(corrupt)<0.01.
2.Case B: while changing the model input there is no significant change in the model output probability in this case the p(clean)-p(corrupt)<=0.01
3.Case C: IN this cases model knows the right answer. but while changing the input the output probability also decreased.and by patching the layer we did not able to recover the answer.
4.Case D: This is our main success case our model passes all baseline p(clean)-p(corrupt)>0.01. in this case model knows the answer while corrupting the input output probability decreases but by patching properly at two layer we able to recover the output .
5.Case E: in this case our model knows the output while changing the input output probability decreases by properly applying 2-layer Patching we got the result but we got the much more confident result .we think this anomaly happen due to mathematical disturbance ,this is not original recall signal.
RESULT:
we have tested 30 different prompt from the Counterfact dataset using a free Google colab CPU and on environment (2 vCPUs, PyTorch 2.11.0+cpu, 4.3 GB RAM usage).
out of 30 only 4 prompts in which the model knows the answer previously and our experiment (Patching) perfectly performed . on experiment we found out that the recall of factual detail is actually occurring between 6 to 9 layer.on fixing layer 6 we able to 43.9% output got right and on layer 9 42.8% answer got corrected. after layer 12 the impact of changes is does not give any better result.
✅ Loaded 30 total prompt runs from Google Drive.
================================================== 📊 METRICS SUMMARY ================================================== Total Prompts Evaluated : 30 Case D Prompts Identified: 4 (13.3%) Peak Localization Layer : Layer 6 (Score: 0.439) -------------------------------------------------- CPU Cores Allocated : 2 System Memory (RAM) : 4.26 GB / 12.67 GB Execution Engine : PyTorch 2.11.0+cpu (CPU) ==================================================
✅ Plot saved to Google Drive: /content/drive/MyDrive/dual_curve_causal_tracing.png
Loading weights: 100%
292/292 [00:00<00:00, 5164.16it/s]
Loaded pretrained model gpt2-medium into HookedTransformer ✅ Model loaded cleanly in CPU memory.
🚀 Starting CPU Causal Tracing Batch... [1/30] Best Layer 7/24 | Score: 1.000 | Case D (Confirmed Mechanism) [2/30] Best Layer 20/24 | Score: 0.003 | Case C (Unlocalized Noise) [3/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [4/30] Best Layer 10/24 | Score: 0.129 | Case D (Confirmed Mechanism) [5/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [6/30] Best Layer 9/24 | Score: 0.030 | Case C (Unlocalized Noise) [7/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [8/30] Best Layer 4/24 | Score: 0.016 | Case C (Unlocalized Noise) [9/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [10/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [11/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [12/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [13/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [14/30] Best Layer 12/24 | Score: 0.008 | Case C (Unlocalized Noise) [15/30] Best Layer 9/24 | Score: 0.237 | Case D (Confirmed Mechanism) [16/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [17/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [18/30] Best Layer 18/24 | Score: 0.000 | Case C (Unlocalized Noise) [19/30] Best Layer 1/24 | Score: 0.028 | Case C (Unlocalized Noise) [20/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [21/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [22/30] Best Layer 4/24 | Score: 0.009 | Case C (Unlocalized Noise) [23/30] Best Layer 21/24 | Score: 0.037 | Case C (Unlocalized Noise) [24/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [25/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [26/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [27/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [28/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop) [29/30] Best Layer 2/24 | Score: 1.000 | Case D (Confirmed Mechanism) [30/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
by mixing all 30 prompt and then see the graph the graph is flat this is due to half of the factual thing model does not know and due to this the result which perform better on mixing it vanishes. but by filtering the prompt where the model actually knows fact then a clear peak is visible between layer 6 to 9. and strongly suggested that restoring activation at Layer 6 is 43.9 % and at layer 9 it is 42.8%.
from the result we conclude that if the model is smaller and running on normal CPU. finding of fact and recall occur in between starting to mid layer only in our case is specifically between layer 6 to layer 9. from this give us clear signal that Ai at the starting layer try to find the fact and on the last layers it try to convert answer into text.
Future Work:
scaling Prompt: Up till now i am testing this on 30 prompt but in future i am trying to test this on 500 prompt.due to which we got the performance of model on different types of fact. and to find model behaviour.
Targeted Model Editing: by using the technique ROME we try to update/change the model weights. to see whether editing at the exact spot to spot whether the AI memory saved fact changed or not.
Testing other smaller model:to find does they also have same factual recall layer between 6 to 9.
Introduction
Current method which we are using to make AI safe like RLHF, Refusal Training, system prompts, this method not able to remove dangerous impact of AI it only highlight the problem in this methods we implement a safety layer in Ai which activate only when we ask any dangerous or harmful question to models ,safety layer act as guard and give command and model output is "Ai is not build to give this answer " .the problem is despite not responding to the unethical question internally models have that dangerous information in their weight or in their brain. so when hacker do the jailbreak or prompt injection in the model the safety layer breaks easily and this information is easily extracted.so the removal of this dangerous behaviour is required. for this we have to shift from behavioural Filter to Permanent Model Editing.
from identifying the problem to actually do some work as a starting undergraduate researcher we have to do casual tracing this research actually take lots of heavy GPU which is not affordable to a undergrad student for this we tested a casual tracing on basic CPU (2 vCPUs , 4.3 GB RAM) setup . and hopely we successfully able to do it in this setup . while using CPU over fast calculating GPU there is a lots of chances of mathematical error and wild number but we successfully able to tackle it.instead of blindly testing on whole sentence ,we localised the intervention window strictly to the last token of subject through which load on the CPU got decreased and we got accurate memory gap. along with this we try to prove that meaningful research on AI safety can be done on the CPU at starting it is not required a lots of GPU cluster.
Engineering Challenges we face:
while Implementing a Multi billion model on the CPU introduces a Technical challenges while building a reliable CPU pipeline we handle two specific edge cases.
1.small model natively lack confidence in fact p(clean)=0, as a result the corruption causes no measurable drop. In implementation zero division causing score explosion or negative infinity crashes.we resolve this by implementing a explicit probability guardrails. if p(clean)-p(corrupt)<0.01. the score is clamped to 0.0. this prevent the mathematical noise.
2. for the model like 355M parameter they act like a small town with information spread more deeply into their layers patching the single layer does not guarantee that this will work.to fix this problem we done a 2 LAYER-SLIDING-WINDOW in this we does not tweak one layer at a time so we try to modify 2 layer at once and then try to catch where the information lives.instead of changing the end of the sentence we edit the targeted where the main subject is mentioned.
Now we design a filtering process which decide whether the AI model give factual answer or it randomly gives us Random signal.each prompt is passed to the 5 gates .
1.Case A: where the model does not know the right answer. if. the model does not know "capital of France is Paris" then , by editing or doing any experiment or doing testing does not worth anymore. it is then rejected and in this case p(clean)-p(corrupt)<0.01.
2.Case B: while changing the model input there is no significant change in the model output probability in this case the p(clean)-p(corrupt)<=0.01
3.Case C: IN this cases model knows the right answer. but while changing the input the output probability also decreased.and by patching the layer we did not able to recover the answer.
4.Case D: This is our main success case our model passes all baseline p(clean)-p(corrupt)>0.01. in this case model knows the answer while corrupting the input output probability decreases but by patching properly at two layer we able to recover the output .
5.Case E: in this case our model knows the output while changing the input output probability decreases by properly applying 2-layer Patching we got the result but we got the much more confident result .we think this anomaly happen due to mathematical disturbance ,this is not original recall signal.
RESULT:
we have tested 30 different prompt from the Counterfact dataset using a free Google colab CPU and on environment (2 vCPUs, PyTorch 2.11.0+cpu, 4.3 GB RAM usage).
out of 30 only 4 prompts in which the model knows the answer previously and our experiment (Patching) perfectly performed . on experiment we found out that the recall of factual detail is actually occurring between 6 to 9 layer.on fixing layer 6 we able to 43.9% output got right and on layer 9 42.8% answer got corrected. after layer 12 the impact of changes is does not give any better result.
✅ Loaded 30 total prompt runs from Google Drive.
==================================================
📊 METRICS SUMMARY
==================================================
Total Prompts Evaluated : 30
Case D Prompts Identified: 4 (13.3%)
Peak Localization Layer : Layer 6 (Score: 0.439)
--------------------------------------------------
CPU Cores Allocated : 2
System Memory (RAM) : 4.26 GB / 12.67 GB
Execution Engine : PyTorch 2.11.0+cpu (CPU)
==================================================
✅ Plot saved to Google Drive: /content/drive/MyDrive/dual_curve_causal_tracing.png
Loading weights: 100%
292/292 [00:00<00:00, 5164.16it/s]
Loaded pretrained model gpt2-medium into HookedTransformer
✅ Model loaded cleanly in CPU memory.
🚀 Starting CPU Causal Tracing Batch...
[1/30] Best Layer 7/24 | Score: 1.000 | Case D (Confirmed Mechanism)
[2/30] Best Layer 20/24 | Score: 0.003 | Case C (Unlocalized Noise)
[3/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[4/30] Best Layer 10/24 | Score: 0.129 | Case D (Confirmed Mechanism)
[5/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[6/30] Best Layer 9/24 | Score: 0.030 | Case C (Unlocalized Noise)
[7/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[8/30] Best Layer 4/24 | Score: 0.016 | Case C (Unlocalized Noise)
[9/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[10/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[11/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[12/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[13/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[14/30] Best Layer 12/24 | Score: 0.008 | Case C (Unlocalized Noise)
[15/30] Best Layer 9/24 | Score: 0.237 | Case D (Confirmed Mechanism)
[16/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[17/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[18/30] Best Layer 18/24 | Score: 0.000 | Case C (Unlocalized Noise)
[19/30] Best Layer 1/24 | Score: 0.028 | Case C (Unlocalized Noise)
[20/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[21/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[22/30] Best Layer 4/24 | Score: 0.009 | Case C (Unlocalized Noise)
[23/30] Best Layer 21/24 | Score: 0.037 | Case C (Unlocalized Noise)
[24/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[25/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[26/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[27/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[28/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
[29/30] Best Layer 2/24 | Score: 1.000 | Case D (Confirmed Mechanism)
[30/30] Best Layer 1/24 | Score: 0.000 | Case B (Insufficient Drop)
by mixing all 30 prompt and then see the graph the graph is flat this is due to half of the factual thing model does not know and due to this the result which perform better on mixing it vanishes. but by filtering the prompt where the model actually knows fact then a clear peak is visible between layer 6 to 9. and strongly suggested that restoring activation at Layer 6 is 43.9 % and at layer 9 it is 42.8%.
from the result we conclude that if the model is smaller and running on normal CPU. finding of fact and recall occur in between starting to mid layer only in our case is specifically between layer 6 to layer 9. from this give us clear signal that Ai at the starting layer try to find the fact and on the last layers it try to convert answer into text.
Future Work:
Colab Link: https://colab.research.google.com/drive/169DNd9TOOnA6Eqz4XR-icjB3eFUasI68?usp=sharing