Decouple PMI weight from sequence sampling temperature

README.md

@@ -87,3 +87,25 @@ and reset each block/iteration. Gamma still subtracts gamma*C in log-score units
 this is not a percentage-parameter remapping. Gamma=0 restores unpenalized results
 for the selected maximum-path or sampling mode. Ordinary decoding ignores these
 checkboxes; the original temperature weighting is unchanged.
+
+
+## Independent temperature and bigram weight
+
+The LLaDA Space now uses S(y) = sum log p_t(y_t) + w sum PMI - gamma sum C.
+**Bigram weight** ranges from 0 to 5, default 1. A weight of 1 gives equal
+coefficients, not necessarily equal numerical influence. Weight 0 disables PMI
+but does not disable repetition penalties. The previous fixed 0.5*temperature
+PMI coefficient described above is superseded for this Space.
+
+**Sample bigram-guided sequences** samples P(y) proportional to exp(S(y)/T).
+Temperature controls the entropy of the whole sequence score, including repetition,
+without changing the ratios of score coefficients. T=0 selects the maximum path;
+positive T is floored to 1e-6 for numerical stability. Maximum-path Viterbi ignores
+temperature. Raw full-vocabulary diffusion probabilities still determine expected
+repetition counts and retention/APD confidence. Default T is now 1, w is 1,
+matching the former T=2 PMI coefficient and sequence distribution. Ordinary
+diffusion retains its usual temperature behavior. No extra model forwards.
+
+Unseen pairs have a neutral PMI correction of zero, never a forbidden transition.
+Conditional scoring instead uses smoothed finite probabilities. Expected repetition
+counts remain the default; the optional candidate-set count mode is unchanged.

app.py

@@ -33,7 +33,7 @@ PMI_MIN_COUNT = 1
 PMI_CLIP = 2.0
 BIGRAM_ENDING_DELAY_STRENGTH = 10.0
 # Preserve the candidate pool and balance used in the matched Viterbi tests.
-DEFAULT_TEMPERATURE = 2.0 if BIGRAM_ENABLED else 0.7
+DEFAULT_TEMPERATURE = (1.0 if SPACE.get("model_backend") == "llada" else 2.0) if BIGRAM_ENABLED else 0.7
 DEFAULT_TOP_K = 16 if BIGRAM_ENABLED else 3
 
 MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", SPACE["model_repo_id"])
@@ -125,6 +125,7 @@ def generate(
     adaptive_parallel_decoding: bool = False,
     sequence_sampling: bool = False,
     count_all_candidate_bigrams: bool = False,
+    bigram_weight: float = 1.0,
 ):
     """Stream iterative masked-diffusion generation from the fixed model."""
     question = question.strip() or "What do you know about Amsterdam?"
@@ -136,11 +137,13 @@ def generate(
             raise gr.Error("Bigram Viterbi is not enabled for this Space.")
         try:
             selector = BigramViterbiSelector(
-                _bigram_scorer(bigram_model_path).to(SESSION.device), 0.5,
+                _bigram_scorer(bigram_model_path).to(SESSION.device),
+                float(bigram_weight) if getattr(SESSION, "llada", False) else 0.5,
                 scoring_method=("pmi" if remasking_strategy == "PMI Viterbi + Repetition Penalty" else
                     {"Conditional bigram": "conditional", "PMI bigram": "pmi"}.get(bigram_scoring_method, "")),
                 min_count=int(pmi_min_count), pmi_clip=float(pmi_clip),
                 sequence_sampling=bool(sequence_sampling),
+                independent_temperature=bool(getattr(SESSION, "llada", False)),
                 remasking_confidence={"Diffusion probability": "diffusion", "Viterbi path margin": "viterbi_margin"}.get(bigram_remasking_confidence, ""),
             )
             if remasking_strategy == "PMI Viterbi + Repetition Penalty":
@@ -240,7 +243,11 @@ with gr.Blocks(title=f"{DISPLAY_NAME} · masked diffusion") as demo:
             num_steps = gr.Slider(1, 512, value=64, step=1, label="Denoising steps")
             block_length = gr.Slider(1, 512, value=64, step=1, label="Block length")
             temperature = gr.Slider(0.0, 2.0, value=DEFAULT_TEMPERATURE, step=0.05, label="Temperature",
-                info="In Viterbi, lower values favor diffusion predictions over bigrams; 0 ignores bigram scores. Viterbi remains deterministic." if BIGRAM_ENABLED else None)
+                info=("For sequence sampling: higher values increase randomness; 0 selects the best path. Bigram weight controls PMI influence. Maximum-path Viterbi ignores temperature."
+                    if SESSION.llada else "In Viterbi, lower values favor diffusion predictions over bigrams; 0 ignores bigram scores. Viterbi remains deterministic.") if BIGRAM_ENABLED else None)
+            bigram_weight = gr.Slider(0.0, 5.0, value=1.0, step=0.05, label="Bigram weight",
+                visible=SESSION.llada and BIGRAM_ENABLED,
+                info="PMI coefficient: 0 disables PMI; 1 gives equal diffusion and PMI score coefficients. Independent of temperature.")
             top_k = gr.Slider(1, 100, value=DEFAULT_TOP_K, step=1, label="Top-k",
                 info="Number of candidate tokens per position for sampling or Viterbi. At 1, only the model’s highest-scoring token is available." if BIGRAM_ENABLED else None)
             pause_per_step = gr.Slider(
@@ -361,6 +368,7 @@ with gr.Blocks(title=f"{DISPLAY_NAME} · masked diffusion") as demo:
         adaptive_parallel_decoding,
         sequence_sampling,
         count_all_candidate_bigrams,
+        bigram_weight,
     ]
     def update_ar_controls(strategy, apd=False):
         """Expose AR-only controls and clear revision when another strategy is selected."""

src/diffusion_lm/bigram.py

@@ -205,7 +205,7 @@ def expected_bigram_repetition_penalties(token_ids, probabilities, vocab_size, s
 
 
 class BigramViterbiSelector:
-    def __init__(self, scorer, weight=0.5, *, scoring_method="conditional", min_count=5, pmi_clip=5.0, remasking_confidence="diffusion", sequence_sampling=False):
+    def __init__(self, scorer, weight=0.5, *, scoring_method="conditional", min_count=5, pmi_clip=5.0, remasking_confidence="diffusion", sequence_sampling=False, independent_temperature=False):
         if not math.isfinite(weight) or weight < 0:
             raise ValueError("Bigram weight must be finite and nonnegative")
         BigramScorer.validate_scoring(scoring_method, min_count, pmi_clip)
@@ -213,6 +213,7 @@ class BigramViterbiSelector:
             raise ValueError("Remasking confidence must be 'diffusion' or 'viterbi_margin'")
         self.remasking_confidence = remasking_confidence
         self.sequence_sampling = bool(sequence_sampling)
+        self.independent_temperature = bool(independent_temperature)
         self.last_path_margins = None
         self.scorer, self.weight = scorer, float(weight)
         self.scoring_method, self.min_count, self.pmi_clip = scoring_method, min_count, float(pmi_clip)
@@ -243,7 +244,7 @@ class BigramViterbiSelector:
         return token_ids.gather(1, path[:, None]).squeeze(1)
 
     @staticmethod
-    def sample_sequence(token_ids, emissions, transitions, weight=1.0):
+    def sample_sequence(token_ids, emissions, transitions, weight=1.0, temperature=1.0):
         """Exact forward-filter/backward sampling from exp(total path score).
 
         No score perturbations or model calls. Centering each forward row
@@ -255,7 +256,18 @@ class BigramViterbiSelector:
         length, k = emissions.shape
         if transitions.shape != (length - 1, k, k):
             raise ValueError("Expected transitions [B-1,K,K]")
+        temperature = float(temperature)
+        if not math.isfinite(temperature) or temperature < 0:
+            raise ValueError("Sequence temperature must be finite and nonnegative")
+        if temperature == 0:
+            return BigramViterbiSelector.decode(token_ids, emissions, transitions, weight)
+        # Global path temperature changes entropy without changing the relative
+        # diffusion/PMI/repetition weights. Tiny positive T approaches argmax.
+        scale = max(temperature, 1e-6)
         edges = transitions.float() * weight if weight else torch.zeros_like(transitions, dtype=torch.float32)
+        if edges.numel():
+            edges = (edges - edges.amax(dim=(1, 2), keepdim=True)) / scale
+        emissions = (emissions.float() - emissions.float().amax(-1, keepdim=True)) / scale
         forward = torch.empty_like(emissions, dtype=torch.float32)
         row = emissions[0].float()
         forward[0] = row - row.max()
@@ -321,6 +333,9 @@ class BigramViterbiSelector:
         raw log probabilities + (weight * T) * transitions. The latter avoids
         division by tiny T and has the exact greedy limit at T=0. Margins use
         these common rescaled units, preserving their retention ranking.
+        With independent_temperature=True, weight directly multiplies PMI and
+        temperature only sets the sequence sampling distribution exp(S/T).
+        Maximum-path decoding is temperature-independent; T=0 samples its argmax.
         Returned diffusion confidence always uses raw full-vocabulary p_t.
         Optional repetition penalties subtract gamma*C(a,b) directly from each
         edge, independently of the PMI weight and temperature. Expected counts
@@ -336,7 +351,7 @@ class BigramViterbiSelector:
         temperature = float(temperature)
         if not math.isfinite(temperature) or temperature < 0:
             raise ValueError("temperature must be finite and nonnegative")
-        weight = self.weight * temperature
+        weight = self.weight if self.independent_temperature else self.weight * temperature
         if not math.isfinite(weight):
             raise ValueError("temperature times bigram weight must be finite")
         self.last_path_margins = None
@@ -388,7 +403,10 @@ class BigramViterbiSelector:
             if repetition_penalty_strength:
                 edges = edges - penalties
             decoder = self.sample_sequence if self.sequence_sampling else self.decode
-            selected = decoder(candidates, emissions, edges, 1.0)
+            if self.sequence_sampling and self.independent_temperature:
+                selected = decoder(candidates, emissions, edges, 1.0, temperature=temperature)
+            else:
+                selected = decoder(candidates, emissions, edges, 1.0)
             if self.remasking_confidence == "viterbi_margin":
                 self.last_path_margins = self.path_margins(candidates, emissions, edges, selected, 1.0)
         elif self.remasking_confidence == "viterbi_margin":