Skip to content

MAINT Model GCG optimization iterations as explicit state #2416

Description

@romanlutz

Is your feature request related to a problem? Please describe.

ProgressiveMultiPromptAttack.run in pyrit/executor/promptgen/gcg/attack/base/attack_manager.py spans roughly 130 lines with eight levels of nesting. It coordinates optimizer iteration state, candidate generation, filtering, batch scoring, loss computation, candidate selection, stopping conditions, and result bookkeeping in one loop.

The coupling makes algorithm changes difficult to test in isolation and increases parameter/state fan-out. This is especially relevant to the performance work tracked in #962: faster sampling or scoring strategies are harder to integrate safely while the iteration lifecycle remains implicit.

This issue is not a duplicate of #962. Issue #962 tracks faster GCG approaches; this issue tracks the internal structure needed to make optimization changes testable and maintainable. Implementations should coordinate with that issue.

Describe the solution you'd like

Introduce a typed optimization-iteration state object that captures the current suffix/candidates, losses, best result, counters, and stop reason. Separate the loop into explicit phases with stable contracts:

  1. candidate generation;
  2. candidate filtering and batching;
  3. model scoring/loss computation;
  4. candidate selection and state update;
  5. stopping and final result construction.

Keep extension protocols and public attack behavior compatible. The goal is to reduce state fan-out and make each phase deterministic under seeded tests, not to redesign the GCG algorithm in the same change.

Describe alternatives you've considered, if relevant

A performance-only patch under #962 could modify the current loop directly, but that would compound special cases. Conversely, a broad class hierarchy for every phase would add unnecessary abstraction. Prefer typed data plus a small number of cohesive internal functions.

Additional context

Validation should include:

  • deterministic seeded optimizer tests;
  • candidate ordering and batch boundaries;
  • exact loss/selection behavior;
  • early-stop and retry/exhaustion paths;
  • extension protocol implementations;
  • device/dtype and tokenization boundaries;
  • result/status equivalence before and after the refactor;
  • representative benchmark parity, with performance work reported separately under FEAT Faster Greedy Coordinate Gradient algorithm (GCG) #962.

This is high-risk algorithmic code. Keep the refactor behavior-preserving and coordinate scope with #962.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requesthelp wantedExtra attention is needed

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions