fix(causal): Harden clamp residuals and add resume notebook

[DaviBonetto]

Protect long Colab sweeps from tiny float residuals in the discrete_A clamp and add a resumable notebook flow for partially completed Table 3 runs.

• compute discrete_A rho in float64 and apply follow-up correction when post-clamp rho still overshoots by epsilon

• cover the residual-overshoot path with a regression test

• add a dedicated resume notebook that syncs the branch, reuses checkpoint CSVs, and saves progress after each threshold

Summary by CodeRabbit

• New Features

  • Introduced a Jupyter notebook to resume causal intervention experiments from saved checkpoints.

• Bug Fixes

  • Improved numerical stability in causal intervention calculations with iterative refinement.

• Tests

  • Added test coverage for residual correction handling in intervention processing.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 760a855e-a36b-4d6a-98a6-8dd76b1bd32c

📥 Commits

Reviewing files that changed from the base of the PR and between 2b8cf07 and a7b03e4.

📒 Files selected for processing (3)

• notebooks/07_Causal_Intervention_Resume.ipynb
• scripts/run_causal_intervention.py
• tests/test_run_causal_intervention.py

---

📝 Walkthrough

Walkthrough

Introduces a resumable Jupyter notebook for causal intervention sweeps that reads from checkpoints and refines rho-clamping logic with multi-pass refinement, improved overshoot handling, and dtype preservation in discrete activations.

Changes

Cohort / File(s)	Summary
Resumable Causal Intervention Notebook notebooks/07_Causal_Intervention_Resume.ipynb	New self-contained notebook that resumes causal intervention sweeps from checkpoint CSVs; includes environment setup, repo/model loading, checkpoint reading, prompt mining, protocol execution for both all_layer and single_layer sweeps, and final consolidation with summary statistics.
Rho-Clamping Logic Refinements scripts/run_causal_intervention.py	Enhanced _compute_discrete_a_rho_ to cast discrete_A to float64 for computations and back to original dtype. Improved _clamp_discrete_a_to_target_ with device/dtype normalization, multi-pass refinement loop using nextafter margin, and iterative scale adjustment to mitigate residual overshoots (up to three iterations).
Clamping Correction Tests tests/test_run_causal_intervention.py	New test test_clamp_discrete_a_applies_follow_up_correction_for_float_residuals that verifies multi-pass correction behavior via mocked rho values, ensuring overshoot masks and scale reductions are applied correctly across tokens.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant Setup as Environment<br/>Setup
    participant Repo as Repo<br/>Manager
    participant Model as Model<br/>Loader
    participant Checkpoint as Checkpoint<br/>Handler
    participant Prompt as Prompt<br/>Miner
    participant Protocol as Protocol<br/>Executor
    participant Output as Consolidator

    User->>Setup: Initialize notebook
    Setup->>Setup: Install deps, prepare workspace
    User->>Repo: Locate/clone repo & script
    Repo->>Repo: Handle manual override if needed
    User->>Model: Load model + tokenizer
    Model->>Model: Move to CUDA, set padding token
    User->>Checkpoint: Read existing CSVs
    Checkpoint->>Checkpoint: Determine completed targets
    User->>Prompt: Mine validated prompts
    Prompt->>Prompt: Cache per-seed results
    User->>Protocol: Execute sweep resume
    Protocol->>Checkpoint: Iterate seeds & targets
    Protocol->>Prompt: Generate prompts
    Protocol->>Protocol: Evaluate accuracy per rho
    Protocol->>Checkpoint: Write results to CSVs
    User->>Output: Validate & consolidate
    Output->>Output: Aggregate results, create summary
    Output->>User: Display final paths & stats

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• fix/colab causal repro notebook #6: Directly modifies the same clamping helpers (_compute_discrete_a_rho and _clamp_discrete_a_to_target) with related multi-pass refinement and overshoot correction logic.
• fix(causal): tolerate minor rho proxy overshoot #3: Modifies clamp/overshoot logic to handle small rho overshoots using a margin/tolerance check approach.
• fix/colab causal repro notebook #1: Modifies rho/clamp logic and hook/evaluation behavior in the same causal intervention script.

Poem

A curious rabbit resumes the quest, 🐰
Through causal sweeps and rho-targets blessed,
Checkpoints restore what came before,
Multi-pass magic at the core,
Overshoots vanish, precision's won! ✨

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/fix/colab-causal-repro-notebook

📝 Coding Plan

• Generate coding plan for human review comments

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}