From 446be81576964a83e9386a512b048a49b4b9701d Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Sat, 15 Aug 2026 17:13:08 +0100 Subject: [PATCH] Preserve exact-zero SFT self-conditioning --- .../hd/hd_gemma_ar_state_handler.py | 2 + .../hd/hd_gemma_ar_state_handler_test.py | 11 +++++ .../hd/hd_gemma_network.py | 40 ++++++++++++++--- .../hd/sft_model.py | 18 +++++--- .../hd/sft_model_test.py | 43 +++++++++++++++++++ 5 files changed, 101 insertions(+), 13 deletions(-) diff --git a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler.py b/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler.py index 1a9e0c67..788def29 100644 --- a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler.py +++ b/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 diff --git a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler_test.py b/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler_test.py index 733b9f88..5c154004 100644 --- a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_ar_state_handler_test.py +++ b/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): diff --git a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py b/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py index ee0689fc..c558857a 100644 --- a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py +++ b/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, diff --git a/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py b/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py index f4a35fa5..5983f5ff 100644 --- a/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py +++ b/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,7 +402,6 @@ 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] @@ -404,13 +409,12 @@ def __call__( 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( diff --git a/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model_test.py b/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model_test.py index 73122dc6..9f0c2a02 100644 --- a/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model_test.py +++ b/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.