Weaver: Steerable Search with Claude

Improved neural search with Sparse Autoencoders

Comment

Inspiration

Finding specific, nuanced information within vast technical corpora like arXiv remains a significant challenge. Traditional keyword search often fails to capture semantic intent, while standard vector search can surface generally relevant but not precisely targeted results. We were inspired by advancements in Mechanistic Interpretability, particularly the potential of Sparse Autoencoders (SAEs) to decompose LLM activations into more understandable features. Could we leverage these interpretable features to give users fine-grained control over the semantic search process, effectively "steering" the search towards their specific informational goals?

What it does

Weaver is a steerable semantic search engine for arXiv papers. It allows users to:

  1. Perform standard semantic search using embeddings derived from intermediate activations of the Gemma-2-2b model, from the Gemma Scope project.
  2. Perform steered search, where the semantic representation of the query is modified based on interpretable features extracted by Sparse Autoencoders (SAEs).
  3. Utilize Claude Sonnet as an intelligent agent to:
    • Analyze the user's query intent.
    • Select the most relevant SAE features (based on human-readable autointerp explanations from Neuronpedia API) to steer the search towards that intent.
    • Determine appropriate steering strengths for selected features.
    • Optionally rewrite queries for improved clarity and filter noise (like citations) from search results.

Essentially, Weaver enables users to guide the search process not just by what they ask, but by emphasizing or de-emphasizing underlying semantic concepts identified within the language model itself.

How we built it

Weaver integrates several cutting-edge components:

  1. Base LLM & Execution: We use gemma-2-2b as the base model. The system dynamically selects the execution backend: MLX on Apple Silicon for optimized performance, or PyTorch on other platforms (using CUDA if available).
  2. Sparse Autoencoders (SAEs): Pre-trained SAEs (specifically gemma-scope-2b-pt-res for layer 20) are loaded using the sae-lens library. This library simplifies integrating SAEs with transformer models and provides tools for feature extraction.
  3. Interpretability Data: We leverage Neuronpedia's API to fetch human-readable explanations for SAE features and their typical maximum activation values, which are crucial for semantic understanding and calibrated steering.
  4. Activation Access & Steering: transformer-lens provides the hooking mechanism access Gemma's intermediate activations (specifically blocks.20.hook_resid_post) and inject steering vectors derived from SAE feature decoder weights.
  5. Vector Search: FAISS is used to create an IndexFlatL2 index storing the baseline activation vectors of text chunks extracted from arXiv PDFs. It enables efficient nearest-neighbor search.
  6. Intelligent Control: The Anthropic Claude 3.5 Sonnet API (anthropic library) acts as the reasoning engine. It analyzes query intent, selects steering features based on semantic relevance to the intent and feature explanations, determines steering strength, rewrites queries, and filters search results.
  7. Backend API: A FastAPI server manages the indexing pipeline (PDF parsing with PyPDF2, chunking, embedding), handles search requests, interacts with Claude, performs steering calculations, and queries the FAISS index.

Workflow:

  • Indexing: PDFs are parsed, chunked, embedded using Gemma's layer 20 activations, and stored in FAISS.
  • Steered Search: A query's intent is analyzed by Claude. Top activating SAE features for the query are identified using Gemma+SAE. Claude selects relevant features based on intent and explanations. The query's base activation vector is retrieved from Gemma. Steering offsets (using SAE decoder weights, Claude's strength, and Neuronpedia's max_act) are calculated and added to the base vector. This final steered vector is used to search FAISS. Claude filters the results.

Challenges we ran into

  • Component Integration: Combining multiple complex systems (Gemma, SAEs, FAISS, Claude API, MLX/Torch backends) required careful management of dependencies, data formats, and potential failure points.
  • SAE Feature Utility: While SAEs provide interpretable features, identifying which features are genuinely useful for steering specific queries, and translating user intent into effective feature manipulation, remains challenging.
  • Prompt Engineering: Designing effective prompts for Claude to accurately analyze intent, select relevant features, and assign appropriate strengths required significant iteration.
  • Performance: Integrating multiple model inferences (Gemma for embedding/activations, Claude for control) introduces latency. Optimizing embedding speed and managing API call times were key considerations.
  • Calibration: Determining the right scaling (max_act) and strength for steering vectors to achieve noticeable but controlled effects required experimentation.

Accomplishments that we're proud of

  • Successfully building an end-to-end system demonstrating steerable search using interpretable LLM features.
  • Integrating Sparse Autoencoders directly into a practical search application, moving beyond purely theoretical interpretability work.
  • Leveraging a LLM (Claude 3.5 Sonnet) as an intelligent control layer to mediate the complex task of selecting and applying steering based on semantic understanding.
  • Creating a functional system that adapts its ML backend (MLX/PyTorch) based on the available hardware.
  • Demonstrating a practical application of Mechanistic Interpretability, aligning with Anthropic's mission by using insights into model internals for controllable behavior.

What we learned

  • The practicalities of applying SAEs for model intervention and the nuances of feature selection for effective steering.
  • The significant potential (and associated challenges) of using LLMs like Claude for complex reasoning and decision-making within larger AI systems.
  • Advanced prompt engineering techniques are critical for reliable LLM-driven control.
  • Managing dependencies and execution environments for multi-backend ML systems requires careful planning.
  • Vector database indexing and retrieval strategies are crucial for the performance of semantic search applications.
  • Directly connecting interpretability research (understanding model internals) to tangible application development (controlling model behavior).

What's next for Weaver: Steerable Search with Claude

  • Refine Steering Interface: Develop a more intuitive UI allowing users to directly select or modify steering features and strengths.
  • Expand Corpus & Model: Index different datasets beyond arXiv and experiment with different base models or SAEs trained on various layers/models.
  • Evaluation Framework: Develop robust metrics to quantitatively evaluate the effectiveness and precision of steered search compared to baseline methods.
  • Deeper Interpretability Integration: Explore visualizing feature activations or providing more context to the user about why certain features were chosen for steering.

Built With

Share this project:

Updates

Submission history