Improve ASTRA’s computational efficiency

Develop algorithmic and implementation techniques to reduce the memory footprint and runtime of the ASTRA attention-based prompt injection attack, which currently requires extracting and backpropagating through large attention matrices and is approximately twice as slow as Greedy Coordinate Gradient (GCG) under equal forward-pass budgets.

Background

ASTRA relies on computing and differentiating through attention matrices, whose size scales quadratically with the context length. The authors note that this leads to significant memory consumption and reduced batch sizes, resulting in roughly a 2× slowdown relative to GCG on their 48 GB GPU setup.

The paper focuses on demonstrating attack effectiveness, leaving efficiency improvements for future work. Addressing efficiency is crucial for practical deployment, larger-scale evaluations, and adaptation to longer context windows.