gemma_sd — Isole di stabilità in DiffusionGemma4

Estrazione, visualizzazione e (in futuro) manipolazione della dinamica di formazione delle "isole di stabilità" durante la generazione di un modello text-diffusion: mlx-community/diffusiongemma-26B-A4B-it-8bit, eseguito con MLX su Apple Silicon.

---

1. Idea di base

diffusiongemma-26B-A4B è un modello block-diffusion: invece di generare i token uno alla volta (autoregressivo), produce un intero canvas di 256 token partendo da rumore e raffinandolo per denoising iterativo (fino a 48 passi). A ogni passo alcune posizioni diventano "certe" (bassa entropia / accettate dal sampler) e si bloccano; le altre restano rumore e vengono ri-campionate.

Le posizioni certe non si accendono a caso: tendono a formare isole contigue nello spazio del canvas, che nucleano e crescono/si fondono nel tempo di denoising. L'obiettivo del progetto è:

1. estrarre questa dinamica (quando e dove ogni posizione si stabilizza);
2. caratterizzarla statisticamente e per dominio testuale;
3. manipolarla trattando il denoising come un simulated annealing di cui controllare gli iperparametri.

L'analogia con il simulated annealing

Il loop di denoising già fa annealing: applica ai logit uno schedule di temperatura lineare t_max→t_min (0.8 → 0.4) che si raffredda monotònamente dal primo all'ultimo passo. L'esplorazione ("uscire dai minimi") non è data da reheating ma dalla ri-rumorizzazione delle posizioni non accettate. Manca però la mossa in salita del SA classico: non c'è un meccanismo che scongela un token già fissato — l'insieme accettato può solo crescere. Questo è il margine su cui interverremo (vedi Piani futuri).

---

2. Scelte prese

Osservare il sampler reale, non reimplementarlo. Il loop di denoising vero è stream_diffusion_generate in mlx_vlm/generate/diffusion.py. Invece di riscriverlo (rischio di divergere dal modello), eseguiamo la generazione reale e la osserviamo via monkeypatch di due funzioni a livello di modulo:

• _make_diffusion_decoder_logits_fns → cattura argmax per posizione a ogni forward (= un passo di denoising);
• _diffusion_entropy_transfer_mask → cattura l'entropy (già scalata in temperatura) e l'acceptance_mask esatta del modello.

Blocco singolo. Per restare "dentro lo stesso blocco di 256 token" si vincola la generazione a un solo canvas: diffusion_full_canvas=True + max_tokens = canvas_length, temperature=0.0 (denoiser = argmax).

Tre definizioni di "certo", decise a posteriori sugli stessi dati: accept (maschera del sampler), entropy (entropy < entropy_bound), stability (argmax che smette di cambiare).

Persistenza crash-safe. Un .npz per blocco + manifest.jsonl; il run salta i blocchi già fatti (resume). Necessario perché il modello 26B mette sotto sforzo la GPU e il watchdog Metal uccide il processo ogni ~15–40 blocchi: un loop shell rilancia e riprende fino al target.

Aggregazione tempo-normalizzata. I blocchi hanno T diversi (9–48 passi): si normalizza la timeline a "progresso di denoising" [0,1] e si media tra blocchi, globale e per dominio.

Note tecniche apprese:

• mlx_vlm.generate è una funzione che maschera il subpackage → importare con from mlx_vlm.generate import diffusion, non import mlx_vlm.generate.diffusion.
• stream_generate non applica la chat template (lo fa solo generate()): senza, il modello -it emette solo <eos>. Si applica apply_chat_template.
• mx.clear_cache() a ogni blocco riduce la pressione di memoria GPU.

File del repo

File	Ruolo
scripts/sd_curve.py	Generazione reale + recorder; salva per_step_argmax/entropy/accept per blocco
scripts/islands.py	Derivazioni: certain_mask, islands_per_step, nucleation_step, island_stats
scripts/plot_islands.py	Heatmap per singolo blocco (matplotlib)
scripts/aggregate.py	Aggregazione su tutti i blocchi → 3 figure PNG
scripts/explore_islands.py	Notebook marimo interattivo (Plotly) con filtro domini
runs/pilot/	Dati per-blocco (.npz) + manifest.jsonl + figure

Parametri chiave

canvas_length=256 · max_denoising_steps=48 · entropy_bound=0.1 · schedule temperatura 0.8→0.4 · sampler entropy-bound (l'unico supportato).

---

3. Informazioni già estratte

Su ~112 blocchi (target 200), 9 domini testuali (prosa_en, prosa_it, codice, lista, ragionam, dialogo, creativo, istruzioni, fattuale):

• Gradiente inizio → fine (corretto). Sul solo contenuto (padding rimosso, posizione relativa nel testo), il lock-time sale in modo quasi monotòno da inizio (~0.2) a fine (~0.7): l'inizio ancorato al prompt si risolve per primo, la fine del testo reale per ultima. Il precedente "bordi→centro" era in gran parte artefatto del padding (il bordo destro del canvas è <eos>). Controllo: scripts/content_locktime.py.
• Nucleazione → fusione reale, non banale. Il numero di isole sale poi scende (atteso topologicamente per un campo di copertura che cresce), ma un null a rivelazione casuale dà ~2.6× isole in più dell'osservato (25.5 vs 7.6; ratio 0.39): la certezza è spazialmente clusterizzata oltre il caso. Controllo: scripts/null_islands.py.
• Differenze per dominio. Output corti (creativo, prosa_it) fissano presto le code (eos/padding); output lunghi (codice, istruzioni) tengono il canvas conteso a lungo. ragionam mostra una struttura "a gradini": blocchi di posizioni si fissano insieme a livelli discreti di progresso (ragionamento risolto a pezzi).
• Caratterizzazione del "annealing" del modello. Raffreddamento monotòno (no reheating); esplorazione tramite ri-rumorizzazione del non-accettato; nessuno scongelamento dei token già fissati.

Generalità — secondo modello (Nemotron, masked-diffusion)

Stessa cattura e stessi controlli su Nemotron-Labs-Diffusion-3B (meccanismo diverso: parte tutto [MASK], rivela ~1 posizione/passo per confidenza). Entrambi i finding reggono, più forti: clustering spaziale ratio 0.09 (vs 0.39), gradiente lock inizio→fine 0.00 → 0.52 → 0.97 (vs 0.30 → 0.69). Non sono artefatti di un singolo checkpoint. Script: scripts/nemotron_curve.py → runs/nemotron/, poi null_islands.py accept runs/nemotron e content_locktime.py accept runs/nemotron.

Step 1 completato — surprisal AR conferma il lock-time

Il modello autoregressivo gemello (gemma-4-26B-A4B-it, stesso tokenizer) è stato usato come validatore indipendente: forward teacher-forced sull'output SD → surprisal causale per token, correlata col lock_time = nucleation_step/T. Risultato (200 blocchi, 47.476 posizioni content — stabile già a 115):

• correlazione grezza Pearson +0.13 (Spearman +0.10);
• residua dopo rimozione del trend posizionale +0.16 (non svanisce → non è un artefatto di posizione), contro un null permutato ≈0;
• positiva su tutti i 9 domini (+0.06 codice → +0.21 prosa_it).

Interpretazione: i token che si uniscono più tardi alle isole sono genuinamente più difficili (surprisal maggiore). Il lock_time è quindi un segnale di difficoltà valido, e la surprisal AR è il metro indipendente per giudicare le varianti dello Step 2. Script: scripts/surprisal.py + scripts/surprisal_analysis.py; figura in runs/pilot/surprisal_analysis.png.

Controllo di scope — difficoltà ≠ correttezza. Sui domini verificabili (ragionam, fattuale; 35 blocchi con match deterministico), le risposte sbagliate non hanno decoding più difficile: la surprisal è persino più bassa (0.36 vs 0.54) e il lock-time invariato. Quindi lock-time e surprisal misurano predicibilità/fluenza, non correttezza (il modello sbaglia in modo fluente). Delimita lo Step 1: "lock tardivo = meno predicibile" ≠ "più probabilmente sbagliato". Script: scripts/correctness.py (caveat: n piccolo, match-stringa).

Le figure aggregate sono in runs/pilot/aggregate_*_<mode>.png; la versione interattiva nel notebook marimo.

---

4. Piani futuri

Step 1 — Validatore di fondo: SD vs AR (surprisal) — ✓ FATTO

Vedi "Step 1 completato" nella sezione 3. In sintesi: la surprisal causale del modello AR gemello correla positivamente col lock-time (residua +0.16, robusta al controllo posizione, positiva su tutti i domini). La surprisal AR è ora il metro indipendente per giudicare le varianti dello Step 2.

Step 2 — Simulated annealing programmabile ✓

Harness in scripts/anneal.py: rende attivi gli hook su _diffusion_linear_temperature (schedule arbitrario, reheating) e _diffusion_entropy_transfer_mask (stato committed per vincolare max_islands e per il melt su reheating). max_islands=None + melt_fraction=0 = baseline fedele (pass-through). I controlli sono meccanicamente corretti (max2 tiene ≤2 isole; vincolare costa più passi) e nessuna variante degenera — il modello è robusto al reveal-order forzato.

Leve parametrizzate (anneal.py), misurate contro lo Step 1 (surprisal AR) e con LLM-as-judge:

• schedule di temperatura arbitrario, incluso reheating (cosine_restarts, sawtooth);
• numero di centri di nucleazione fissato (max_islands), via proiezione della maschera di accettazione sul vincolo (tiene le K isole più grandi);
• un-freezing / "melt fraction": nei passi di reheating, ri-rumorizza una frazione dei token già congelati — la mossa in salita che al modello manca.

Pipeline di valutazione:

• scripts/anneal_sweep.py — 8 config × 2 prompt/dominio (crash-safe, resume) → runs/anneal/.
• scripts/anneal_eval.py — surprisal AR media + metriche isole per config, delta vs baseline appaiato.
• scripts/anneal_judge.py — LLM-as-judge (Claude Sonnet, pairwise vs baseline) su coerenza/fluenza/aderenza (carica ANTHROPIC_API_KEY da .env).

Risultati Step 2 (sweep pieno: 144 tracce, 18 prompt; surprisal+CI e giudice Sonnet). Il modello è robusto: quasi ogni config è statisticamente indistinguibile dal baseline (tutti i CI della surprisal attraversano lo zero, win-rate del giudice 0.42–0.50). Unico candidato debole: max2 (≤2 centri) — giudice 0.67, surprisal −0.08, coerente tra campioni ma non solido a n=18. Il sweep pieno ha smentito l'n=7: "reheating senza melt = peggiore" e "il melt ripara in modo monotòno" erano rumore da piccolo campione (ora cosine3 è un pareggio). Lezione: CI + campione pieno erano necessari. Output in runs/anneal/judge_results.json.

Domande guida: l'ordine/topologia di rivelazione è causale per la qualità? Un ciclo caldo→freddo→caldo fa fondere e ri-nucleare le isole? L'eventuale degenerazione dell'output è essa stessa un risultato sulla robustezza del modello a un reveal-order non monotòno.

---

Uso

uv sync
uv run python scripts/sd_curve.py            # raccolta dati (crash-safe, resume)
uv run python scripts/aggregate.py accept    # figure aggregate (mode: accept|entropy|stability)
uv run marimo edit scripts/explore_islands.py  # esplorazione interattiva