ADAPT: Hybrid Prompt Optimization for LLM Feature Visualization

1. Problem Statement

The paper addresses the challenge of Feature Visualization in Large Language Models (LLMs). The goal is to identify input prompts that maximally activate specific "latent directions" (features) within an LLM's activation space, typically learned via Sparse Autoencoders (SAEs).

• Limitations of Existing Approaches:
  • Dataset Search: The standard approach involves running the model on massive corpora to find activating examples. This is computationally prohibitive for complex objectives (e.g., model circuits) or when specific latents are not well-represented in existing datasets.
  • Discrete Input Space: Unlike computer vision, LLM inputs are discrete tokens, preventing direct gradient-based optimization.
  • Failure Modes of Current Prompt Optimization: Existing methods (like GCG, Beam Search, or Evolutionary approaches) struggle with:
    • Local Minima: The discrete search space is prone to getting stuck in suboptimal solutions.
    • Poor Initialization: Random initialization often leads to failure.
    • Signal-to-Noise Ratio: Gradient-based estimates for token selection are often noisy and unreliable.
    • Fluency vs. Activation: Balancing the generation of coherent text with high feature activation is difficult; purely activation-driven methods often produce "gibberish."

2. Methodology: ADAPT

The authors propose ADAPT (Adaptive DynAmic Prompt Tuning), a hybrid optimization algorithm designed to overcome the specific failure modes of LLM feature visualization. It combines Beam Search initialization with adaptive gradient-guided mutation.

Core Components:

1. Beam-Search Initialization:

   • Instead of random initialization, ADAPT starts with a parallel beam search procedure.
   • It generates multiple candidate sequences by appending likely tokens (right-append) and inserting tokens into the middle of sequences (middle-insert).
   • This ensures the optimization starts in "favorable regions" of the prompt space, reducing the likelihood of early stagnation in local minima.

2. Hybrid Candidate Generation:

   • ADAPT maintains a population of prompts. At each iteration, it generates new candidates using a probabilistic mix of two mutation strategies:
     • GCG-style Swaps: Uses a gradient-based estimate (treating the one-hot token vector as differentiable) to identify tokens that maximize the target latent activation.
     • Logit-Swaps: A gradient-free approach that samples tokens based on the model's own predicted probabilities (similar to Evolutionary Prompt Optimization).
   • This hybrid approach balances exploitation (using gradients to refine) with exploration (using model probabilities to discover new patterns).

3. Diversity-Preserving Selection (Slot Management):

   • To prevent the population from converging prematurely on minor variations of a single prompt, ADAPT uses a group-based slot system.
   • Prompts are divided into "groups." Each group has "guaranteed slots" (keeping the best prompt from that group) and "free slots" (allocated to the best performers across all groups).
   • This maintains diverse optimization trajectories, ensuring the discovery of multiple distinct feature patterns.

4. Adaptive Fluency Penalty:

   • A sigmoid-based fluency penalty is applied to the loss function.
   • The penalty weight starts low and gradually increases over iterations. This allows the optimizer to first explore high-activation regions (even if incoherent) and then refine the prompt toward fluency in later stages.

3. Key Contributions

1. ADAPT Algorithm: A novel hybrid method that outperforms existing baselines (GCG, BEAST, EPO) by combining beam-search initialization with adaptive mutation strategies.
2. Rigorous Evaluation Framework: The authors evaluate on Sparse Autoencoder latents from Gemma 2 2B. They introduce new metrics grounded in dataset activation statistics:
   • Activation Ratio: The ratio of the optimized prompt's activation to the top Maximally Activating Example (MAE) found in the dataset.
   • Activation Rank: The rank of the optimized prompt among the top 15 dataset MAEs.
3. Empirical Analysis of Failure Modes: The paper provides deep insights into why existing methods fail:
   • GCG: Exhibits a poor signal-to-noise ratio in gradient estimates and is highly sensitive to initialization and trajectory.
   • Beam Search: Fast but lacks the fine-grained refinement needed for convergence.
   • Fluency: High activation correlates strongly with human-rated interpretability, suggesting that forcing fluency too early may hinder finding the true feature pattern.

4. Experimental Results

The evaluation was conducted on 486 latents across 6 layers of the Gemma 2 2B model.

• Performance Superiority:
  • ADAPT consistently outperforms GCG, BEAST, and EPO across all layers.
  • It achieves higher activation than GCG on 70.4% of features (statistically significant, $p < 10^{-28}$).
  • Unlike GCG (which degrades in deeper layers) or BEAST (which struggles in early layers), ADAPT maintains a stable performance profile across the entire model depth.
• Failure Mode Analysis:
  • GCG's gradient estimates are noisy; many high-gradient tokens do not lead to improvements, while many improvements come from low-gradient tokens.
  • GCG is highly sensitive to initialization; "autoregressive" initializations perform better than random ones, but results vary significantly even with the same start.
• Interpretability:
  • Human evaluation (3 raters) showed that ADAPT produces significantly more interpretable prompts (Mean Rating 4.44) compared to GCG (3.86) and GCG with fluency penalties (3.78).
  • There is a strong correlation ($r=0.645$) between high activation ratios and human-rated interpretability, validating that optimizing for activation recovers meaningful semantic patterns rather than artifacts.
• Dual Latent Optimization:
  • ADAPT was successfully extended to optimize for two latents simultaneously (e.g., "MCU context" + "magnet properties"), producing prompts like "Stark magnet" that satisfy both constraints.

5. Significance and Conclusion

• Tractability: The paper establishes that feature visualization for LLMs is tractable but requires domain-specific design assumptions (handling discreteness, local minima, and non-orthogonal latents).
• Tool for Interpretability: ADAPT provides a robust tool for researchers to visualize and understand the internal representations of LLMs, particularly for SAE latents where dataset search is insufficient.
• Future Directions: The authors suggest ADAPT can be extended to visualize steering vectors, model circuits, and other complex objectives. They also note limitations regarding the current focus on a single model family (Gemma) and the need for automated interpretability pipelines.

In summary, ADAPT solves the "needle in a haystack" problem of finding specific feature activations in LLMs by intelligently combining the broad exploration of beam search with the precision of gradient-guided mutation, all while managing diversity to avoid local minima.