[Bug] Hackable Diffusion SFT semantics differ from the DiffusionGemma Technical Report

# [Bug] Hackable Diffusion SFT semantics differ from the DiffusionGemma Technical Report

## Summary

The released Hackable Diffusion Adapter SFT path appears to differ from the
training semantics described in the DiffusionGemma Technical Report in two
places:

1. disabling self-conditioning produces a zero-logit soft embedding rather
than the report-defined zero self-conditioning signal;
2. multi-canvas SFT decodes the full response after moving the KV-cache cursor
to the selected canvas, which misaligns decoder K/V writes whenever
`selected_canvas_idx > 0`.

I verified this against:

- repository commit:
`7b785991bd78626c73b317eb43fdbb6c292f7b9c`
- DiffusionGemma Technical Report v1:
https://arxiv.org/abs/2608.00146

Could you please confirm whether these differences are intentional, or whether
the released SFT adapter should follow the report semantics?

## 1. Self-conditioning OFF is not an exact zero signal

The report states that, during SFT, 50% of data points use a
self-conditioning state computed by a previous forward pass, while the other
50% have:

```text
z_t = 0
```

See Section 8 / Equation 13, page 23:

https://arxiv.org/pdf/2608.00146#page=23

The current SFT implementation instead selects zero vocabulary logits for
disabled examples:

https://github.com/google-deepmind/gemma/blob/7b785991bd78626c73b317eb43fdbb6c292f7b9c/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py#L397-L413

```python
sc_logits = converted_first_pass['logits']
zero_logits = jnp.zeros_like(sc_logits)
...
sc_logits = jnp.where(do_self_cond, sc_logits, zero_logits)
```

The wrapper then always converts these logits into soft token embeddings:

https://github.com/google-deepmind/gemma/blob/7b785991bd78626c73b317eb43fdbb6c292f7b9c/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py#L220-L239

```python
if sc_logits is None:
sc_logits = jnp.zeros(...)
sc_embeddings = self.gemma_model.embedder.encode_logits(sc_logits)
```

`encode_logits` applies softmax and an embedding-table projection:

https://github.com/google-deepmind/gemma/blob/7b785991bd78626c73b317eb43fdbb6c292f7b9c/gemma/gm/nn/gemma4/_modules.py#L140-L155

Therefore, for zero logits:

```text
softmax(0) = [1/V, ..., 1/V]

encode_logits(0)
= sqrt(d) * (1/V) * sum_v E_v
```

This is generally the vocabulary-mean embedding, not the exact zero vector
specified by the report. It also affects the first denoising pass because
`sc_logits is None` is converted into zero logits and then encoded in the same
way.

### Expected behavior

The SC-disabled branch should provide an exact zero embedding/signal after the
logit-to-embedding conversion boundary, for example by applying a per-example
mask to `sc_embeddings`:

```python
sc_embeddings = embedder.encode_logits(sc_logits)
sc_embeddings = jnp.where(
self_conditioning_mask[..., None],
sc_embeddings,
jnp.zeros_like(sc_embeddings),
)
```

When no previous logits are supplied, the first pass should likewise use an
exact zero self-conditioning embedding.

## 2. Multi-canvas SFT misaligns decoder KV-cache writes

The technical report describes selecting one canvas `k` and training its
denoising loss using:

```text
current noisy canvas x_t
+ KV cache H for the prompt
+ previous uncorrupted canvases 1 ... k-1
```

See:

- Section 4, page 12:
https://arxiv.org/pdf/2608.00146#page=12
- Section 8 / Equation 13, page 23:
https://arxiv.org/pdf/2608.00146#page=23

The current implementation corrupts the entire `K * C` response before
selecting a canvas:

https://github.com/google-deepmind/gemma/blob/7b785991bd78626c73b317eb43fdbb6c292f7b9c/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py#L302-L334

```python
time = self.time_sampler(..., x0)
xt, target_info = self.corruption_process.corrupt(..., x0, time)
...
selected_canvas_idx = jax.random.randint(...)
```

The encoder then sets:

https://github.com/google-deepmind/gemma/blob/7b785991bd78626c73b317eb43fdbb6c292f7b9c/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py#L211-L215

```python
end_index = prompt_len + selected_canvas_idx * canvas_size
```

However, `sft_decode` still submits the full
`total_canvas_len = K * C` query and uses positions for the full response:

https://github.com/google-deepmind/gemma/blob/7b785991bd78626c73b317eb43fdbb6c292f7b9c/gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py#L129-L145

```python
num_queries = total_canvas_len
canvas_positions = positions[:, prompt_len:]
...
gemma_network(xt=xt, ...)
```

The underlying attention cache writes queries starting at `end_index`,
wrapping modulo the cache size:

https://github.com/google-deepmind/gemma/blob/7b785991bd78626c73b317eb43fdbb6c292f7b9c/gemma/gm/nn/gemma4/_modules.py#L298-L313

```python
indices = (
end_index[:, None] + jnp.arange(seq_len)[None, :]
) % cache_size
```

### Minimal index example

Let:

```text
prompt_len P = 4
canvas_size C = 2
num_canvases K = 3
selected canvas k = 1
cache_size = P + K*C = 10
```

Then:

```text
end_index = P + k*C = 6
seq_len = K*C = 6
write indices = (6 + arange(6)) % 10
= [6, 7, 8, 9, 0, 1]
```

This maps the noisy queries as follows:

```text
noisy canvas 0 -> canvas 1 cache slots
noisy canvas 1 -> canvas 2 cache slots
noisy canvas 2 -> prompt cache slots
```

Meanwhile, the decoder attention mask continues to interpret physical cache
slots as:

```text
[0, P) = prompt
[P + j*C, P+(j+1)*C) = canvas j
```

Consequences for `selected_canvas_idx > 0` include:

- prompt K/V can be overwritten by later noisy canvases;
- the selected noisy canvas K/V is written into later-canvas slots;
- the selected canvas may not see its own noisy K/V under the existing mask;
- the forward pass can condition on displaced noisy states instead of
`clean previous canvases + noisy selected canvas`.

Masking the loss to the selected canvas does not repair the forward-pass cache
layout.

This issue is hidden when `num_canvases == 1` or when canvas 0 is selected.

### Expected behavior

A report-aligned path could:

1. sample the canvas index first;
2. gather only the selected clean canvas;
3. sample time and corrupt only that canvas;
4. use a decoder query of exactly `canvas_size`;
5. write its K/V only to the selected canvas's absolute cache slots;
6. expose only the prompt, clean previous canvases, and the selected noisy
canvas through the attention mask.

Conceptually:

```text
Decoder input: noisy C_k
Visible KV context: prompt + clean C_ 0`:

- verify decoder query length is exactly `canvas_size`;
- verify decoder writes only to the selected canvas slots;
- verify prompt K/V remains unchanged;
- verify changing a previous clean canvas changes selected-canvas logits;
- verify changing a future canvas does not change selected-canvas logits;
- compare the multi-canvas result with an equivalent single-canvas reference
execution.

## Questions

1. Is the released Hackable Diffusion SFT adapter intended to reproduce the
Technical Report's SFT objective?
2. If the current full-response decoder behavior is intentional, could the
semantic difference be documented?
3. If it is not intentional, would a patch implementing selected-canvas-only
corruption/decode and an exact-zero SC mask be welcome?

Start in gemma/diffusion/hackable_diffusion_adapter/hd/sft_model.py, tracing sft_decode, canvas selection, corruption, and self-conditioning into hd_gemma_network.py; inspect cache indexing in gemma/gm/nn/gemma4/_modules.py. Add regression coverage for exact-zero disabled self-conditioning and selected-canvas query/cache behavior, then verify prompt and future-canvas state are unaffected while previous clean canvases remain visible.

Written by the indexing model from the issue text.