feat(it-opt): latent-injection adapter scaffold (step 1)

[smallfishabc]

What

Step-1 scaffold for integrating inference-time optimization (IT-opt) into Sampleworks, ported from the it_opt/ reference forks (Protenix/Boltz/AF3).

Adds LatentAdaptedWrapper — a small, swappable transform on a model's single representation (the post-trunk latent s) inserted at the featurize → step conditioning boundary. Currently a deliberately trivial affine k*s + b ("pseudo-MLP") so the plumbing and tests are validated before any real module lands.

Design

• Minimal blast radius. Lives entirely behind LatentAdaptedWrapper, which itself satisfies the FlowModelWrapper protocol — samplers, scalers, and eval consume the protocol, so nothing downstream changes. With the identity transform (k=1, b=0) it is behaviourally identical to the model it wraps.
• Model-agnostic. The only model-specific knowledge is which conditioning attribute holds the single representation — one string via AttrLatentIO (s for Boltz, s_trunk for Protenix/RF3). No model packages are imported; a 4th/5th model is one more string, not new code.
• Gradient isolation. The injected latent is detached during sampling, so the existing coordinate-only guidance gradient (and all downstream tests) behave exactly as before. Training the transform is a separate pass (step 2).

Included

• src/sampleworks/models/latent_adapter.py + tests/models/test_latent_adapter.py
• docs/latent_adapter/ — implementation roadmap, latent_space_optimization.md, and solution diagrams

Status / next steps

Draft — this is the scaffold only. Follow-ups track the migration-notes sequencing (Tier 3 safe cleanups → Tier 1 optimizer-lifetime correctness fix, A/B'd on backbone-RMSD → Tier 2 compute win). Real experimental-signal transform (step 2) comes after this seam is validated.

Note: tests have not yet been run in the project's pixi/pod env — validate before marking ready for review.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ff9e2b5c-f9ae-4992-b635-8556e9e9d0ff

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fy/it-optimization

---

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.}

[manzuoni-astera]

Thanks for keeping this as a narrowly scoped step-1 scaffold. Treating the optimizer, sampler feature hook, Protenix detach handling, Boltz2 cache recomputation, and real experimental transform as intentional follow-up work, the basic read/transform/write seam looks reasonable.

Before this is marked ready, I think the scaffold contract needs a few corrections:

1. LatentAdaptedWrapper is not currently drop-in for the actual guidance entry point. _run_guidance dispatches preprocessing from model_wrapper.__class__.__name__ and will reject LatentAdaptedWrapper as unknown (guidance_script_utils.py:444-486). Pure guidance also reads model_wrapper.model directly at line 541, while diagnostics and MSA reporting use other concrete-wrapper attributes. The one-line factory wrapping shown in the diagram would therefore fail before sampling. Please either make orchestration decorator-aware or deliberately proxy/unwrap the required model identity and attributes, with a pipeline-level test.

2. A wrong single_attr silently becomes a no-op (latent_adapter.py:85-88, 172-174), and the test currently enshrines that behavior. Because this string is required configuration, I would prefer a clear error, or an explicit opt-in allow_missing mode, so field drift cannot masquerade as successful adaptation.

3. The mock step() ignores conditioning, so the tests prove tensor replacement but not that rewritten conditioning affects the delegated step. A conditioning-sensitive mock should cover identity equivalence and non-identity propagation.

CI unit tests and typechecks pass across all environments, but lint currently fails with 39 introduced Ruff errors. GPU jobs are awaiting environment approval. With those scaffold-level issues addressed, the deferred model-specific work can reasonably remain in follow-ups.