Extend paper experiments and submission artifacts

Author: apartsin

app/bootstrap/experiment_models.py

@@ -8,6 +8,9 @@
 
 _CLIP_CACHE: dict[tuple[str, str], tuple[object, object]] = {}
 _DINO_CACHE: dict[tuple[str, str], tuple[object, object]] = {}
+_SIGLIP_CACHE: dict[tuple[str, str], tuple[object, object]] = {}
+_BLIP_CAPTION_CACHE: dict[tuple[str, str], tuple[object, object]] = {}
+_LPIPS_CACHE: dict[tuple[str, str], object] = {}
 
 
 def huggingface_cache_dir() -> Path:
@@ -16,6 +19,13 @@ def huggingface_cache_dir() -> Path:
     return path
 
 
+def _local_cache_error(model_family: str, model_id: str) -> RuntimeError:
+    return RuntimeError(
+        f"{model_family} model '{model_id}' is not available in the local cache. "
+        "Run scripts/preload_experiment_models.py first."
+    )
+
+
 def get_clip_components(model_id: str, device: str, *, local_only: bool = True):
     key = (model_id, device)
     cached = _CLIP_CACHE.get(key)
@@ -38,10 +48,7 @@ def get_clip_components(model_id: str, device: str, *, local_only: bool = True):
         )
     except (OSError, LocalEntryNotFoundError) as exc:
         if local_only:
-            raise RuntimeError(
-                f"CLIP model '{model_id}' is not available in the local cache. "
-                "Run scripts/preload_experiment_models.py first."
-            ) from exc
+            raise _local_cache_error("CLIP", model_id) from exc
         raise
     model.eval()
     _CLIP_CACHE[key] = (model, processor)
@@ -70,11 +77,82 @@ def get_dino_components(model_id: str, device: str, *, local_only: bool = True):
         ).to(device)
     except (OSError, LocalEntryNotFoundError) as exc:
         if local_only:
-            raise RuntimeError(
-                f"DINO model '{model_id}' is not available in the local cache. "
-                "Run scripts/preload_experiment_models.py first."
-            ) from exc
+            raise _local_cache_error("DINO", model_id) from exc
         raise
     model.eval()
     _DINO_CACHE[key] = (processor, model)
     return processor, model
+
+
+def get_siglip_components(model_id: str, device: str, *, local_only: bool = True):
+    key = (model_id, device)
+    cached = _SIGLIP_CACHE.get(key)
+    if cached is not None:
+        return cached
+
+    from transformers import SiglipImageProcessor, SiglipModel
+
+    cache_dir = huggingface_cache_dir()
+    try:
+        processor = SiglipImageProcessor.from_pretrained(
+            model_id,
+            cache_dir=str(cache_dir),
+            local_files_only=local_only,
+            use_fast=False,
+        )
+        model = SiglipModel.from_pretrained(
+            model_id,
+            cache_dir=str(cache_dir),
+            local_files_only=local_only,
+        ).to(device)
+    except (OSError, LocalEntryNotFoundError) as exc:
+        if local_only:
+            raise _local_cache_error("SigLIP", model_id) from exc
+        raise
+    model.eval()
+    _SIGLIP_CACHE[key] = (processor, model)
+    return processor, model
+
+
+def get_blip_caption_components(model_id: str, device: str, *, local_only: bool = True):
+    key = (model_id, device)
+    cached = _BLIP_CAPTION_CACHE.get(key)
+    if cached is not None:
+        return cached
+
+    from transformers import BlipForConditionalGeneration, BlipProcessor
+
+    cache_dir = huggingface_cache_dir()
+    try:
+        processor = BlipProcessor.from_pretrained(
+            model_id,
+            cache_dir=str(cache_dir),
+            local_files_only=local_only,
+            use_fast=False,
+        )
+        model = BlipForConditionalGeneration.from_pretrained(
+            model_id,
+            cache_dir=str(cache_dir),
+            local_files_only=local_only,
+        ).to(device)
+    except (OSError, LocalEntryNotFoundError) as exc:
+        if local_only:
+            raise _local_cache_error("BLIP caption", model_id) from exc
+        raise
+    model.eval()
+    _BLIP_CAPTION_CACHE[key] = (processor, model)
+    return processor, model
+
+
+def get_lpips_metric(net_name: str, device: str):
+    key = (net_name, device)
+    cached = _LPIPS_CACHE.get(key)
+    if cached is not None:
+        return cached
+
+    import lpips
+
+    model = lpips.LPIPS(net=net_name).to(device)
+    model.eval()
+    _LPIPS_CACHE[key] = model
+    return model

app/core/config_yaml.py

@@ -14,11 +14,11 @@
     # Edit any of these values before creating a new session.
     # This YAML is reloaded fresh for each setup page visit or reset action.
     #
-    # sampler: random_local | exploit_orthogonal | uncertainty_guided | axis_sweep | incumbent_mix | diversity_shell | line_search | plateau_escape | annealed_shell | spherical_cover | two_scale_cover | quality_diversity_mix
-    # updater: winner_average | winner_copy | linear_preference | score_weighted_preference | contrastive_preference | softmax_preference | borda_preference | bradley_terry_preference | challenger_mixture_preference | plackett_luce_preference
+    # sampler: random_local | exploit_orthogonal | uncertainty_guided | axis_sweep | incumbent_mix | diversity_shell | line_search | plateau_escape | annealed_shell | spherical_cover | two_scale_cover | quality_diversity_mix | restart_bridge_mix
+    # updater: winner_average | winner_copy | linear_preference | score_weighted_preference | contrastive_preference | softmax_preference | borda_preference | bradley_terry_preference | challenger_mixture_preference | plackett_luce_preference | advantage_softmax_preference
     # feedback_mode: scalar_rating | pairwise | top_k | winner_only | approve_reject
     # seed_policy: fixed-per-round | fixed-per-candidate | fixed-per-candidate-role
-    # steering_mode: currently low_dimensional
+    # steering_mode: low_dimensional | content_masked | token_factorized | token_vector_field
     # steering_dimension: low-dimensional steering vector size, for example 3 or 5
     # candidate_count: visible candidates per round
     # image_size: WIDTHxHEIGHT, for example 512x512

app/core/schema.py

@@ -72,6 +72,7 @@ class SamplerType(str, Enum):
     spherical_cover = "spherical_cover"
     two_scale_cover = "two_scale_cover"
     quality_diversity_mix = "quality_diversity_mix"
+    restart_bridge_mix = "restart_bridge_mix"
 
 
 class UpdaterType(str, Enum):
@@ -85,10 +86,14 @@ class UpdaterType(str, Enum):
     bradley_terry_preference = "bradley_terry_preference"
     challenger_mixture_preference = "challenger_mixture_preference"
     plackett_luce_preference = "plackett_luce_preference"
+    advantage_softmax_preference = "advantage_softmax_preference"
 
 
 class SteeringMode(str, Enum):
     low_dimensional = "low_dimensional"
+    content_masked = "content_masked"
+    token_factorized = "token_factorized"
+    token_vector_field = "token_vector_field"
 
 
 class StrategyConfig(BaseModel):

app/engine/generation.py

@@ -34,7 +34,12 @@ def resolve_steering_mode(session: Session) -> SteeringMode:
     """Resolve and validate the session steering mode used at generation time."""
 
     mode = session.config.steering_mode
-    if mode == SteeringMode.low_dimensional:
+    if mode in {
+        SteeringMode.low_dimensional,
+        SteeringMode.content_masked,
+        SteeringMode.token_factorized,
+        SteeringMode.token_vector_field,
+    }:
         return mode
     raise ValueError(f"Unsupported steering mode: {mode}")
 
@@ -259,19 +264,127 @@ def _resolve_model_source(self, session: Session) -> str:
             "Run scripts/setup_huggingface.py first or enable STABLE_STEERING_ALLOW_REMOTE_MODEL_DOWNLOAD=true."
         )
 
-    def _steering_offset(self, prompt_embeds, z, anchor_strength: float):
+    def _hidden_basis(self, hidden: int, index_id: int, *, device, dtype):
+        """Build a deterministic hidden-space basis vector for one steering axis."""
+
+        torch = self._torch
+        index = torch.linspace(0.0, 1.0, hidden, device=device, dtype=dtype)
+        basis = torch.sin(index * (index_id + 1) * torch.pi) + torch.cos(index * (index_id + 1) * 0.5 * torch.pi)
+        return basis / torch.norm(basis)
+
+    def _token_hidden_basis(self, seq_len: int, hidden: int, index_id: int, *, device, dtype):
+        """Build a deterministic per-token hidden-vector field for one steering axis."""
+
+        torch = self._torch
+        token_index = torch.linspace(0.0, 1.0, seq_len, device=device, dtype=dtype).view(seq_len, 1)
+        hidden_index = torch.linspace(0.0, 1.0, hidden, device=device, dtype=dtype).view(1, hidden)
+        frequency = float(index_id + 1)
+        basis = (
+            torch.sin((token_index + 0.17 * frequency) * (hidden_index + 0.11) * torch.pi * (1.0 + frequency))
+            + 0.7 * torch.cos((token_index * (0.45 + 0.08 * frequency) - hidden_index * (0.63 + 0.04 * frequency)) * torch.pi)
+            + 0.35 * torch.sin((token_index * hidden_index + 0.13 * frequency) * 2.0 * torch.pi)
+        )
+        return basis / torch.clamp(torch.norm(basis), min=torch.tensor(1e-6, device=device, dtype=dtype))
+
+    def _token_inputs(self, pipe, prompt: str, *, seq_len: int, device, dtype):
+        """Tokenize the prompt so token-aware steering modes can shape per-token offsets."""
+
+        tokenizer = getattr(pipe, "tokenizer", None)
+        if tokenizer is None:
+            return None
+
+        tokenized = tokenizer(
+            prompt,
+            padding="max_length",
+            truncation=True,
+            max_length=seq_len,
+            return_tensors="pt",
+        )
+        input_ids = tokenized.input_ids.to(device=device)
+        attention_mask = tokenized.attention_mask.to(device=device, dtype=dtype)
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    def _content_mask(self, token_inputs, *, tokenizer, dtype):
+        """Build a mask that suppresses padding and special tokens for token-aware steering."""
+
+        attention_mask = token_inputs["attention_mask"].to(dtype=dtype)
+        input_ids = token_inputs["input_ids"]
+        content_mask = attention_mask.clone()
+
+        if tokenizer is not None:
+            for attr in ("bos_token_id", "eos_token_id", "pad_token_id"):
+                token_id = getattr(tokenizer, attr, None)
+                if token_id is not None:
+                    content_mask = content_mask * (input_ids != token_id).to(dtype=dtype)
+
+        if float(content_mask.sum()) <= 0.0:
+            return attention_mask
+        return content_mask
+
+    def _steering_offset(self, prompt_embeds, z, anchor_strength: float, *, steering_mode: SteeringMode, token_inputs=None, tokenizer=None):
         """Project the low-dimensional steering vector into embedding space."""
 
         torch = self._torch
+        seq_len = prompt_embeds.shape[1]
         hidden = prompt_embeds.shape[-1]
         device = prompt_embeds.device
         dtype = prompt_embeds.dtype
-        index = torch.linspace(0.0, 1.0, hidden, device=device, dtype=dtype)
         offset = torch.zeros_like(prompt_embeds)
-        for i, value in enumerate(z):
-            basis = torch.sin(index * (i + 1) * torch.pi) + torch.cos(index * (i + 1) * 0.5 * torch.pi)
-            basis = basis / torch.norm(basis)
-            offset = offset + (float(value) * float(anchor_strength)) * basis.view(1, 1, hidden)
+
+        if steering_mode == SteeringMode.low_dimensional:
+            for i, value in enumerate(z):
+                basis = self._hidden_basis(hidden, i, device=device, dtype=dtype)
+                offset = offset + (float(value) * float(anchor_strength)) * basis.view(1, 1, hidden)
+            return offset
+
+        if token_inputs is None:
+            raise ValueError(f"Token-aware steering mode {steering_mode.value} requires token inputs.")
+
+        content_mask = self._content_mask(token_inputs, tokenizer=tokenizer, dtype=dtype)
+        token_positions = torch.linspace(0.0, 1.0, seq_len, device=device, dtype=dtype)
+
+        if steering_mode == SteeringMode.content_masked:
+            token_profile = 0.35 + 0.65 * torch.sin(token_positions * torch.pi)
+            token_profile = token_profile.view(1, seq_len, 1) * content_mask.view(1, seq_len, 1)
+            active_tokens = torch.clamp(content_mask.sum(), min=1.0)
+            normalizer = torch.clamp(token_profile.sum(dim=1, keepdim=True), min=1.0)
+            token_profile = token_profile * (active_tokens / normalizer)
+            for i, value in enumerate(z):
+                basis = self._hidden_basis(hidden, i, device=device, dtype=dtype)
+                offset = offset + (float(value) * float(anchor_strength)) * token_profile * basis.view(1, 1, hidden)
+            return offset
+
+        if steering_mode == SteeringMode.token_factorized:
+            mask = content_mask.view(seq_len)
+            for i, value in enumerate(z):
+                hidden_basis = self._hidden_basis(hidden, i, device=device, dtype=dtype)
+                token_basis = (
+                    torch.sin(token_positions * (i + 1) * torch.pi)
+                    + 0.5 * torch.cos(token_positions * (i + 1) * 2.0 * torch.pi)
+                ) * mask
+                if float(token_basis.abs().sum()) > 0.0:
+                    token_basis = token_basis - ((token_basis * mask).sum() / torch.clamp(mask.sum(), min=1.0)) * mask
+                    token_norm = torch.norm(token_basis)
+                    if float(token_norm) > 0.0:
+                        token_basis = token_basis / token_norm
+                offset = offset + (float(value) * float(anchor_strength) * 0.8) * token_basis.view(1, seq_len, 1) * hidden_basis.view(1, 1, hidden)
+            return offset
+
+        if steering_mode == SteeringMode.token_vector_field:
+            mask = content_mask.view(seq_len, 1)
+            active_tokens = torch.clamp(mask.sum(), min=1.0)
+            for i, value in enumerate(z):
+                token_hidden_basis = self._token_hidden_basis(seq_len, hidden, i, device=device, dtype=dtype) * mask
+                if float(token_hidden_basis.abs().sum()) > 0.0:
+                    token_hidden_basis = token_hidden_basis - token_hidden_basis.sum(dim=0, keepdim=True) / active_tokens
+                    token_hidden_basis = token_hidden_basis * mask
+                    token_hidden_basis = token_hidden_basis / torch.clamp(
+                        torch.norm(token_hidden_basis),
+                        min=torch.tensor(1e-6, device=device, dtype=dtype),
+                    )
+                offset = offset + (float(value) * float(anchor_strength) * 0.7) * token_hidden_basis.unsqueeze(0)
+            return offset
+
         return offset
 
     def _encode_steered_embeddings(self, session: Session, candidate: Candidate):
@@ -286,14 +399,21 @@ def _encode_steered_embeddings(self, session: Session, candidate: Candidate):
             do_classifier_free_guidance=True,
             negative_prompt=session.negative_prompt or "",
         )
-        if steering_mode == SteeringMode.low_dimensional:
-            steered_prompt_embeds = prompt_embeds + self._steering_offset(
-                prompt_embeds,
-                candidate.z,
-                session.config.anchor_strength,
-            )
-        else:
-            raise ValueError(f"Unsupported steering mode: {steering_mode}")
+        token_inputs = self._token_inputs(
+            pipe,
+            session.prompt,
+            seq_len=prompt_embeds.shape[1],
+            device=prompt_embeds.device,
+            dtype=prompt_embeds.dtype,
+        )
+        steered_prompt_embeds = prompt_embeds + self._steering_offset(
+            prompt_embeds,
+            candidate.z,
+            session.config.anchor_strength,
+            steering_mode=steering_mode,
+            token_inputs=token_inputs,
+            tokenizer=getattr(pipe, "tokenizer", None),
+        )
         return steered_prompt_embeds, negative_prompt_embeds
 
     def render_candidate(self, session: Session, candidate: Candidate) -> Candidate:

app/engine/orchestrator.py

@@ -35,9 +35,11 @@
 from app.samplers.plateau_escape import PlateauEscapeSampler
 from app.samplers.quality_diversity_mix import QualityDiversityMixSampler
 from app.samplers.random_local import RandomLocalSampler
+from app.samplers.restart_bridge_mix import RestartBridgeMixSampler
 from app.samplers.spherical_cover import SphericalCoverSampler
 from app.samplers.two_scale_cover import TwoScaleCoverSampler
 from app.samplers.uncertainty import UncertaintyGuidedSampler
+from app.updaters.advantage_softmax_pref import AdvantageSoftmaxPreferenceUpdater
 from app.storage.repository import JsonRepository
 from app.updaters.contrastive_pref import ContrastivePreferenceUpdater
 from app.updaters.borda_pref import BordaPreferenceUpdater
@@ -77,6 +79,7 @@ def __init__(
             "spherical_cover": SphericalCoverSampler(),
             "two_scale_cover": TwoScaleCoverSampler(),
             "quality_diversity_mix": QualityDiversityMixSampler(),
+            "restart_bridge_mix": RestartBridgeMixSampler(),
         }
         self.updaters = {
             "winner_copy": WinnerCopyUpdater(),
@@ -89,6 +92,7 @@ def __init__(
             "bradley_terry_preference": BradleyTerryPreferenceUpdater(),
             "challenger_mixture_preference": ChallengerMixturePreferenceUpdater(),
             "plackett_luce_preference": PlackettLucePreferenceUpdater(),
+            "advantage_softmax_preference": AdvantageSoftmaxPreferenceUpdater(),
         }
 
     @staticmethod