Preserve exact-zero SFT self-conditioning

gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler.py

@@ -62,6 +62,8 @@ def __call__(
       The updated conditioning dict.
     """
     conditioning["sc_logits"] = step_carry.aux["logits"]
+    # Step 0 carries placeholder zero logits, not a previous prediction.
+    conditioning["sc_mask"] = step_carry.step_info.step > 0
     return conditioning
 
 

gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler_test.py

@@ -743,11 +743,22 @@ def test_call(self):
     conditioning = {"other": 1}
     step_carry = mock.MagicMock()
     step_carry.aux = {"logits": jnp.array([1, 2])}
+    step_carry.step_info.step = jnp.array(1)
 
     res = fn(conditioning, step_carry)
     self.assertEqual(res["sc_logits"].tolist(), [1, 2])
+    self.assertTrue(bool(res["sc_mask"]))
     self.assertEqual(res["other"], 1)
 
+  def test_call_disables_self_conditioning_on_first_step(self):
+    fn = hd_gemma_ar_state_handler.PropagateSelfConditioningFn()
+    step_carry = mock.MagicMock()
+    step_carry.aux = {"logits": jnp.zeros((1, 2, 3))}
+    step_carry.step_info.step = jnp.array(0)
+
+    res = fn({}, step_carry)
+    self.assertFalse(bool(res["sc_mask"]))
+
 
 class GemmaARStateHandlerTest(absltest.TestCase):
 

gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py

@@ -40,6 +40,28 @@
 DiffusionGemmaModel = gemma_diffusion.DiffusionGemma_26B_A4B
 
 
+def _prepare_self_conditioning_embeddings(
+    embedder,
+    sc_logits,
+    sc_mask,
+    *,
+    batch_size,
+    canvas_length,
+    vocab_size,
+    dtype,
+):
+  """Convert logits while preserving exact-zero self-conditioning states."""
+  if sc_logits is None:
+    sc_logits = jnp.zeros((batch_size, canvas_length, vocab_size), dtype=dtype)
+    return jnp.zeros_like(embedder.encode_logits(sc_logits))
+  sc_embeddings = embedder.encode_logits(sc_logits)
+  if sc_mask is not None:
+    sc_embeddings = jnp.where(
+        sc_mask, sc_embeddings, jnp.zeros_like(sc_embeddings)
+    )
+  return sc_embeddings
+
+
 # pytype: disable=bad-return-type
 # pytype: disable=signature-mismatch
 
@@ -222,17 +244,23 @@ def __call__(
     # The HD pipeline passes the self-conditioning signal as raw logits
     # (shape [B, L, V]) under the key 'sc_logits' in the conditioning dict.
     sc_logits = conditioning.get('sc_logits', None)
-    if sc_logits is None:
-      sc_logits = jnp.zeros(
-          (batch_size, canvas_length, vocab_size), dtype=dtype
-      )
+    sc_mask = conditioning.get('sc_mask', None)
 
     positions = conditioning.get('positions', None)
     kv_cache = conditioning.get('kv_cache', None)
     attention_mask = conditioning.get('attention_mask', None)
 
-    # We keep this call to maintain the param init behavior.
-    sc_embeddings = self.gemma_model.embedder.encode_logits(sc_logits)
+    # Keep encode_logits in the path for parameter initialization while
+    # preserving DiffusionGemma's exact-zero self-conditioning state.
+    sc_embeddings = _prepare_self_conditioning_embeddings(
+        self.gemma_model.embedder,
+        sc_logits,
+        sc_mask,
+        batch_size=batch_size,
+        canvas_length=canvas_length,
+        vocab_size=vocab_size,
+        dtype=dtype,
+    )
 
     transformer_output = self.gemma_model.call_with_self_conditioning(
         tokens=tokens,

gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py

@@ -31,6 +31,7 @@
 import optax
 
 from gemma.diffusion.hackable_diffusion_adapter import checkpointed_evaluator as _checkpointed_evaluator  # pylint: disable=line-too-long
+
 CheckpointedEvaluator = _checkpointed_evaluator.CheckpointedEvaluator
 
 
@@ -101,6 +102,7 @@ def sft_decode(
     total_canvas_len: int,
     canvas_size: int,
     sc_logits: jnp.ndarray | None = None,
+    sc_mask: jnp.ndarray | None = None,
     is_training: bool = True,
 ) -> hd_typing.TargetInfoTree:
   """Runs the SFT decoder pass: denoise canvases using the prefilled KV cache.
@@ -121,6 +123,8 @@ def sft_decode(
     total_canvas_len: Total canvas length.
     canvas_size: Number of tokens per canvas.
     sc_logits: Optional self-conditioning logits.
+    sc_mask: Optional broadcastable mask selecting examples that use
+      self-conditioning.
     is_training: Whether we are in training mode.
 
   Returns:
@@ -147,6 +151,8 @@ def sft_decode(
   }
   if sc_logits is not None:
     conditioning['sc_logits'] = sc_logits
+  if sc_mask is not None:
+    conditioning['sc_mask'] = sc_mask
 
   return gemma_network(
       xt=xt,
@@ -396,21 +402,19 @@ def __call__(
 
     converted_first_pass = jax.lax.stop_gradient(converted_first_pass)
     sc_logits = converted_first_pass['logits']
-    zero_logits = jnp.zeros_like(sc_logits)
 
     # With probability self_cond_prob, run self-conditioning element-wise.
     batch_size = xt.shape[0]
     do_self_cond = (
         jax.random.uniform(self.make_rng('sampling'), shape=(batch_size,))
         < self.self_cond_prob
     )
-    # Reshape to broadcast with x0_hat_logits (Batch, ..., Channels)
-    do_self_cond = do_self_cond.reshape(
-        (batch_size,) + (1,) * (sc_logits.ndim - 1)
+    # Apply the mask after logits are converted to embeddings so disabled
+    # examples receive an exact-zero self-conditioning signal.
+    sc_mask = do_self_cond.reshape((batch_size,) + (1,) * (sc_logits.ndim - 1))
+    denoiser_output = sft_decode(
+        **decoder_kwargs, sc_logits=sc_logits, sc_mask=sc_mask
     )
-    sc_logits = jnp.where(do_self_cond, sc_logits, zero_logits)
-
-    denoiser_output = sft_decode(**decoder_kwargs, sc_logits=sc_logits)
 
     # Convert predictions (computes loss-ready dict)
     converted = self.corruption_process.convert_predictions(

gemma/diffusion/hackable_diffusion_adapter/hd/sft_model_test.py

@@ -73,6 +73,48 @@ def _make_gemma_network():
   )
 
 
+class _ConstantEmbedder:
+
+  def encode_logits(self, logits):
+    return jnp.ones((*logits.shape[:-1], 4), dtype=logits.dtype)
+
+
+class SelfConditioningEmbeddingTest(absltest.TestCase):
+
+  def test_missing_logits_produce_exact_zero_embeddings(self):
+    embeddings = hd_gemma_network._prepare_self_conditioning_embeddings(
+        _ConstantEmbedder(),
+        None,
+        None,
+        batch_size=2,
+        canvas_length=3,
+        vocab_size=5,
+        dtype=jnp.float32,
+    )
+    np.testing.assert_array_equal(
+        embeddings, jnp.zeros((2, 3, 4), dtype=jnp.float32)
+    )
+
+  def test_mask_zeros_only_disabled_examples(self):
+    logits = jnp.zeros((2, 3, 5), dtype=jnp.float32)
+    mask = jnp.array([True, False])[:, None, None]
+    embeddings = hd_gemma_network._prepare_self_conditioning_embeddings(
+        _ConstantEmbedder(),
+        logits,
+        mask,
+        batch_size=2,
+        canvas_length=3,
+        vocab_size=5,
+        dtype=jnp.float32,
+    )
+    np.testing.assert_array_equal(
+        embeddings[0], jnp.ones((3, 4), dtype=jnp.float32)
+    )
+    np.testing.assert_array_equal(
+        embeddings[1], jnp.zeros((3, 4), dtype=jnp.float32)
+    )
+
+
 def _sft_encode(
     gemma_network,
     *,
@@ -704,6 +746,7 @@ def __call__(
         'positions': canvas_positions,
         'attention_mask': attn_mask,
         'sc_logits': sc_logits,
+        'sc_mask': jnp.array(False),
     }
 
     # Extract gemma_network params from the wrapper's variables.