Targeted Parameter Decomposition Recovers Specific Neural Network Mechanisms

Researchers have developed Targeted Parameter Decomposition (tPD), an innovative method designed to efficiently extract specific computational mechanisms from neural networks. Traditional Parameter Decomposition (PD) aims to break down entire networks into interpretable components, but its scalability to large models is often hindered by immense computational requirements and associated risks. tPD addresses this by focusing solely on identifying and recovering components that process particular inputs of interest, ranging from isolated prompts to larger subtasks. It achieves this efficiency by introducing a high-rank "catch-all" component that handles all data not relevant to the target mechanism. The effectiveness of tPD was validated on both toy models and transformer language models trained on The Pile dataset. The method successfully recovered reproducible and mechanistically faithful circuits. For instance, tPD was able to extract a CSS-only submodel from a 4-block transformer using only 7% of the computational resources required for its full published decomposition. Furthermore, in a 12-block transformer, tPD enabled the surgical ablation and rewiring of memorized sequences, demonstrating its precision and control over specific network behaviors. Crucially, these targeted interventions had negligible side effects on other inputs, highlighting the method's ability to isolate and manipulate specific functionalities within complex models without disrupting overall performance.