[Feature Request] HFLanguageModel backend + SoftPromptOptimizer for latent-space reasoning

[lelouar]

Context

Following our Discord discussion on soft prompts, I've been mapping three converging
research contributions that point to a high-value addition for Synalinks:

• LatentMAS (arXiv:2511.20639, Apache 2.0) — training-free multi-agent latent
  communication via KV cache transfer. Up to −80% tokens, ×7 wall-clock, +14.6pp
  accuracy on 9 benchmarks.
• Latent-Space-Reasoning (github.com/dl1683/Latent-Space-Reasoning) — optimizer
  that searches the embedding space directly with a scorer network + evolutionary loop.
  Already +~20pp on arithmetic with 2 random embedding tokens (no training).
• kNN-LatentMAS (github.com/Bookmaster9/kNN-latentMAS) — bounds KV cache memory
  growth to O(K) per agent regardless of pipeline depth, enabling N-agent latent
  pipelines.

Your comment confirmed the direction: "Sending soft prompts means you bypass the
tokenizer to directly inject vectors [...] that's something we could add in the future
[...] we could implement an optimizer that uses soft prompts [...] using Keras to keep
the style consistent with Synalinks."

This issue proposes the two building blocks needed.

---

Building block 1 — HFLanguageModel backend

A new LanguageModel subclass wrapping HuggingFace Transformers directly (not via
LiteLLM), exposing three capabilities the current LanguageModel can't provide:

• input_embeds injection: soft prefix vectors prepended to the encoded text
  embeddings before the first transformer layer
• past_key_values access: read/write KV cache for inter-agent context transfer
  (LatentMAS-style)
• output_hidden_states: expose the final hidden state for use as the next agent's
  soft prefix

Optional latent_space_realign=True computes a one-shot analytical ridge-regression
mapping W_a that projects latent hidden states back onto the valid embedding manifold
(avoids distribution drift across latent rollout steps). Computed once at init from a
calibration batch — no gradient descent.

Sketch:

class HFLanguageModel(LanguageModel):
    def __init__(
        self,
        model_name: str,
        device: str = "cuda:0",
        latent_space_realign: bool = False,
        calibration_texts: list[str] | None = None,
        **kwargs,
    ):
        # does NOT call super().__init__() — bypasses litellm entirely
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self._W_a = self._compute_realign(calibration_texts) if latent_space_realign else None

    async def __call__(self, messages, schema=None, streaming=False, **kwargs):
        # standard text-in, text-out path (mode OFF compatible)
        ...

    async def forward_with_prefix(
        self,
        text: str,
        soft_prefix: torch.Tensor | None = None,   # shape [n_prefix_tokens, d_model]
        past_key_values=None,                        # injected from previous agent
        output_hidden_states: bool = False,
    ) -> "ForwardOutput":
        # returns: json_output + optional (hidden_states, past_key_values)
        ...

The __call__ signature stays compatible with the existing Generator and OMEGA
pipeline — mode OFF works unchanged. forward_with_prefix is the new surface for
latent mode.

---

Building block 2 — SoftPromptOptimizer

An optimizer (consistent with the Synalinks optimizer API) that learns continuous
soft-prompt prefix vectors to maximize a reward signal.

As you suggested, the architecture maps to: a small Keras scorer network +
evolutionary search over the embedding space. Concretely:

• LatentScorer: a small Keras MLP that takes a candidate soft prefix tensor
  [n_tokens, d_model] and returns a scalar score — without running a full generation
  pass (cheap, ~10ms per candidate vs ~500ms for a full forward).
• Evolution loop: mutation (Gaussian perturbation in embedding space) +
  crossover (linear interpolation between top candidates), guided by the scorer.
• Zero-training fallback: random perturbation already achieves +~20pp on
  arithmetic (per the Latent-Space-Reasoning paper) — the scorer is optional and
  activates only when enough reward signal is available.

Sketch:

class SoftPromptOptimizer(Optimizer):
    """
    Optimizes continuous soft-prompt prefix vectors rather than discrete tokens.
    Orthogonal to OMEGA/RandomFewShot which optimize text prompts.
    """
    def __init__(
        self,
        n_prefix_tokens: int = 2,       # empirical optimum from the paper
        n_candidates: int = 10,
        n_generations: int = 15,
        scorer: "keras.Model | None" = None,  # LatentScorer or None for random search
        name: str = "soft_prompt_optimizer",
    ):
        ...

    def optimize(self, program, x_train, y_train, reward):
        # 1. Sample n_candidates soft prefix tensors
        # 2. Score each with self.scorer (or random if scorer is None)
        # 3. Evolve: mutate top-k, crossover pairs
        # 4. Inject winning prefix into program's HFLanguageModel
        ...

---

Relationship to existing Synalinks components

Component	Mode	Role	Changes needed
LanguageModel (LiteLLM)	OFF	unchanged	none
Generator	both	unchanged — __call__ compat	none
OMEGA / RandomFewShot	both	optimize discrete text role prefixes	none
HFLanguageModel (new)	ON	HF backend, exposes embeddings + KV	new class
SoftPromptOptimizer (new)	ON	optimizes continuous prefix vectors	new class

OMEGA and SoftPromptOptimizer are orthogonal and composable: OMEGA optimizes the
role text prefix, SoftPromptOptimizer optimizes the embedding-space prefix injected
before it. Both can run in the same pipeline.

---

Mode ON/OFF design intent

The goal is a clean BIOSYNA_LATENT_MODE=on/off toggle (or equivalent Synalinks
flag) at the LanguageModel factory level:

• OFF: existing LiteLLM path, zero changes.
• ON: HFLanguageModel replaces LanguageModel in the graph; Generator.call
  routes through forward_with_prefix when a soft prefix or past_key_values is
  present in the input context.

---

kNN-LatentMAS note

For multi-agent pipelines with more than 2 agents, github.com/Bookmaster9/kNN-latentMAS
shows that indexing the KV cache with kNN bounds memory to O(K) per agent regardless
of depth. This is a natural extension once HFLanguageModel exists. I'll open a
separate issue for it if useful.

---

Questions for your guidance

1. HFLanguageModel placement: new subpackage synalinks/language_models/hf/,
   or alongside language_model.py in synalinks/language_models/?

2. SoftPromptOptimizer scorer type: should LatentScorer be a
   synalinks.Module subclass (for graph composability) or a standalone
   keras.Model? This determines whether it lives in synalinks/optimizers/ or
   as an external contrib.

3. Soft prefix serialization: preferred path for persisting torch.Tensor
   soft prompts in the synalinks.saving API — safetensors, numpy .npy, or
   something else?

4. Generator contract: is it preferable to add forward_with_prefix as a
   separate call path on HFLanguageModel, or should Generator.call be extended
   with an optional latent_context kwarg that activates the HF path?

Happy to open a draft PR for HFLanguageModel once I have guidance on structure.