import random import torch import torch.nn.functional as F from stage2_config import * from stage2_losses import flow_matching_loss from stage2_riemannian import generate_suffix, prompt_condition, suffix_positions def seed_everything(seed): random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def sample_token_ids( logits, tokenizer, temperature=EVAL_SAMPLE_TEMPERATURE, top_k=EVAL_SAMPLE_TOP_K, top_p=EVAL_SAMPLE_TOP_P, ): if temperature <= 0: return logits.argmax(dim=-1) logits = logits.float() / temperature for token_id in tokenizer.all_special_ids: logits[..., token_id] = -float("inf") if top_k is not None and top_k > 0: kth = logits.topk(min(top_k, logits.size(-1)), dim=-1).values[..., -1, None] logits = logits.masked_fill(logits < kth, -float("inf")) if top_p is not None and 0.0 < top_p < 1.0: sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True) sorted_probs = sorted_logits.softmax(dim=-1) keep = sorted_probs.cumsum(dim=-1) <= top_p keep[..., 0] = True sorted_logits = sorted_logits.masked_fill(~keep, -float("inf")) logits = torch.full_like(logits, -float("inf")).scatter(dim=-1, index=sorted_idx, src=sorted_logits) probs = logits.softmax(dim=-1) return torch.multinomial(probs.reshape(-1, probs.size(-1)), 1).view(logits.shape[:-1]) def generated_decode_stats(z_real, z_gen_suffix, suffix_mask, input_ids, attn_mask, decoder, device): z_gen_flat = z_gen_suffix[suffix_mask] real_flat = z_real[:, PROMPT_LEN:, :][suffix_mask] gen_mean = z_gen_flat.mean().item() gen_std = z_gen_flat.std().item() cosine_sim = F.cosine_similarity( real_flat.mean(0, keepdim=True), z_gen_flat.mean(0, keepdim=True), ).item() decode_idx = (attn_mask[:, PROMPT_LEN:].sum(dim=1) > 0).nonzero(as_tuple=False).flatten() decode_idx = decode_idx[:DECODE_LOSS_BATCH] if decode_idx.numel() == 0: return gen_mean, gen_std, cosine_sim, 0.0 z_decode_gen = torch.cat([z_real[:, :PROMPT_LEN, :], z_gen_suffix], dim=1)[decode_idx] decode_targets = input_ids[decode_idx, PROMPT_LEN:].reshape(-1) with torch.amp.autocast("cuda", enabled=device.type == "cuda"): gen_decode_logits = decoder.decode_from_latent(z_decode_gen) gen_decode_ce = F.cross_entropy( gen_decode_logits[:, PROMPT_LEN:, :].reshape(-1, gen_decode_logits.size(-1)), decode_targets, ignore_index=0, ).item() return gen_mean, gen_std, cosine_sim, gen_decode_ce def latent_mix_decode_ce(z_real, z_gen_suffix, input_ids, attn_mask, decoder, device, alphas=(0.01, 0.03, 0.05, 0.10, 0.20, 0.50)): decode_idx = (attn_mask[:, PROMPT_LEN:].sum(dim=1) > 0).nonzero(as_tuple=False).flatten() decode_idx = decode_idx[:DECODE_LOSS_BATCH] if decode_idx.numel() == 0: return [] z_real_suffix = z_real[:, PROMPT_LEN:, :] decode_targets = input_ids[decode_idx, PROMPT_LEN:].reshape(-1) results = [] for alpha in alphas: z_mix_suffix = (1.0 - alpha) * z_gen_suffix + alpha * z_real_suffix z_mix_seq = torch.cat([z_real[:, :PROMPT_LEN, :], z_mix_suffix], dim=1)[decode_idx] with torch.amp.autocast("cuda", enabled=device.type == "cuda"): mix_logits = decoder.decode_from_latent(z_mix_seq) mix_ce = F.cross_entropy( mix_logits[:, PROMPT_LEN:, :].reshape(-1, mix_logits.size(-1)), decode_targets, ignore_index=0, ).item() results.append((alpha, mix_ce)) return results def argmax_token_collapse_stats(logits, token_ids, tokenizer): suffix_logits = logits[:, PROMPT_LEN:, :].float() suffix_ids = token_ids[:, PROMPT_LEN:] probs = suffix_logits.softmax(dim=-1) entropy = -(probs * probs.clamp_min(1e-9).log()).sum(dim=-1).mean().item() unique_ratios = [] max_fracs = [] flat_tokens = [] for row in suffix_ids: valid = row[~torch.isin(row, row.new_tensor(tokenizer.all_special_ids))] if valid.numel() == 0: continue counts = torch.bincount(valid.cpu(), minlength=logits.size(-1)) unique_ratios.append((counts > 0).sum().item() / valid.numel()) max_fracs.append(counts.max().item() / valid.numel()) flat_tokens.append(valid.cpu()) if flat_tokens: flat = torch.cat(flat_tokens) counts = torch.bincount(flat, minlength=logits.size(-1)) top_counts, top_ids = counts.topk(min(5, counts.numel())) top_tokens = [ f"{tokenizer.convert_ids_to_tokens(int(token_id))}:{int(count)}" for token_id, count in zip(top_ids.tolist(), top_counts.tolist()) if count > 0 ] else: top_tokens = [] return { "entropy": entropy, "unique_ratio": sum(unique_ratios) / max(len(unique_ratios), 1), "max_frac": sum(max_fracs) / max(len(max_fracs), 1), "top_tokens": top_tokens, } def decoder_distribution_stats(logits, tokenizer, mask=None, oracle_logits=None, target_ids=None): suffix_logits = logits[:, PROMPT_LEN:, :].float() probs = suffix_logits.softmax(dim=-1) top1_acc = 0.0 target_prob_mean = 0.0 oracle_top1_acc = 0.0 oracle_target_prob_mean = 0.0 if target_ids is not None: suffix_targets = target_ids[:, PROMPT_LEN:] if target_ids.size(1) == logits.size(1) else target_ids target_valid = suffix_targets != 0 if mask is not None: target_valid = target_valid & mask.bool() if target_valid.any(): target_gather_ids = suffix_targets.clamp(0, probs.size(-1) - 1).unsqueeze(-1) target_probs = probs.gather(dim=-1, index=target_gather_ids).squeeze(-1) top_ids = probs.argmax(dim=-1) top1_acc = (top_ids[target_valid] == suffix_targets[target_valid]).float().mean().item() target_prob_mean = target_probs[target_valid].mean().item() if oracle_logits is not None: oracle_probs_full = oracle_logits[:, PROMPT_LEN:, :].float().softmax(dim=-1) oracle_target_probs = oracle_probs_full.gather(dim=-1, index=target_gather_ids).squeeze(-1) oracle_top_ids = oracle_probs_full.argmax(dim=-1) oracle_top1_acc = ( oracle_top_ids[target_valid] == suffix_targets[target_valid] ).float().mean().item() oracle_target_prob_mean = oracle_target_probs[target_valid].mean().item() if mask is not None: valid = mask.bool() if valid.any(): probs = probs[valid] suffix_logits = suffix_logits[valid] if oracle_logits is not None: oracle_suffix_logits = oracle_logits[:, PROMPT_LEN:, :].float()[valid] else: probs = probs.reshape(-1, probs.size(-1)) suffix_logits = suffix_logits.reshape(-1, suffix_logits.size(-1)) oracle_suffix_logits = None else: probs = probs.reshape(-1, probs.size(-1)) suffix_logits = suffix_logits.reshape(-1, suffix_logits.size(-1)) oracle_suffix_logits = ( oracle_logits[:, PROMPT_LEN:, :].float().reshape(-1, oracle_logits.size(-1)) if oracle_logits is not None else None ) special_ids = torch.tensor(tokenizer.all_special_ids, device=probs.device) entropy = -(probs * probs.clamp_min(1e-9).log()).sum(dim=-1) top_probs, _ = probs.topk(min(50, probs.size(-1)), dim=-1) mean_probs = probs.clone() mean_probs[:, special_ids] = 0.0 mean_probs = mean_probs.mean(dim=0) batch_top_mass = mean_probs.topk(min(8, mean_probs.numel())).values.sum() special_mass = probs[:, special_ids].sum(dim=-1) kl_to_oracle = None if oracle_logits is not None and oracle_suffix_logits is not None: oracle_log_probs = F.log_softmax(oracle_suffix_logits, dim=-1) gen_log_probs = F.log_softmax(suffix_logits, dim=-1) kl_to_oracle = F.kl_div(gen_log_probs, oracle_log_probs.exp(), reduction="batchmean").item() return { "entropy": entropy.mean().item(), "top1_prob": top_probs[:, 0].mean().item(), "top5_mass": top_probs[:, :min(5, top_probs.size(1))].sum(dim=-1).mean().item(), "top50_mass": top_probs.sum(dim=-1).mean().item(), "batch_top8_mass": batch_top_mass.item(), "special_mass": special_mass.mean().item(), "kl_to_oracle": kl_to_oracle, "top1_acc": top1_acc, "target_prob": target_prob_mean, "oracle_top1_acc": oracle_top1_acc, "oracle_target_prob": oracle_target_prob_mean, } def fuse_aux_logits(flow_net, aux_token_head, decoder_logits, z_suffix, z_cond, mask=None): if aux_token_head is None or AUX_LOGIT_FUSION_BETA <= 0: return decoder_logits B, T, _ = z_suffix.shape pos = suffix_positions(B, T, z_suffix.device, z_suffix.dtype) t = torch.ones(B, T, device=z_suffix.device, dtype=z_suffix.dtype) _, aux_hidden = flow_net(z_suffix, t, z_cond, pos, mask, return_hidden=True) aux_logits = aux_token_head(aux_hidden).float() fused_logits = decoder_logits.float().clone() fused_logits[:, PROMPT_LEN:, :] = fused_logits[:, PROMPT_LEN:, :] + AUX_LOGIT_FUSION_BETA * aux_logits return fused_logits def evaluate( flow_net, metric_net, encoder, decoder, tokenizer, val_loader, device, n_samples=4, start_mlp=None, aux_token_head=None, latent_projector=None, residual_refiner=None, draft_start_fn=None, ): flow_net.eval() metric_net.eval() encoder.eval() decoder.eval() if start_mlp is not None: start_mlp.eval() if aux_token_head is not None: aux_token_head.eval() if latent_projector is not None: latent_projector.eval() if residual_refiner is not None: residual_refiner.eval() val_loss = 0 eval_rng_state = torch.random.get_rng_state() cuda_rng_state = torch.cuda.get_rng_state_all() if device.type == "cuda" else None with torch.no_grad(): for batch in val_loader: input_ids = batch["input_ids"].to(device, non_blocking=True) attention_mask = batch["attention_mask"].to(device, non_blocking=True) with torch.amp.autocast("cuda", enabled=device.type == "cuda"): z_data = decoder.compress(encoder(input_ids, attention_mask)) z_cond = prompt_condition(z_data, attention_mask) z_target = z_data[:, PROMPT_LEN:, :] target_mask = attention_mask[:, PROMPT_LEN:] z_draft_start = draft_start_fn(input_ids, attention_mask) if draft_start_fn is not None else None val_loss += flow_matching_loss( flow_net, metric_net, z_target, z_cond, target_mask, start_mlp=start_mlp, residual_refiner=residual_refiner, z_draft_start=z_draft_start, ).item() avg_val_loss = val_loss / len(val_loader) with torch.no_grad(): seed_everything(SEED + 10_000) batch = next(iter(val_loader)) input_ids = batch["input_ids"].to(device) attn_mask = batch["attention_mask"].to(device) z_real = decoder.compress(encoder(input_ids, attn_mask)) B, S, D = z_real.shape suffix_mask = attn_mask[:, PROMPT_LEN:].bool() z_real_suffix = z_real[:, PROMPT_LEN:, :] z_draft_start = draft_start_fn(input_ids, attn_mask) if draft_start_fn is not None else None z_real_flat = z_real_suffix[suffix_mask] z_cond = prompt_condition(z_real, attn_mask) if z_draft_start is not None and start_mlp is not None and hasattr(start_mlp, "set_draft_target"): start_mlp.set_draft_target(z_draft_start) z_gen_suffix, metric_snapshot, z_initial_suffix, z_uncalibrated_suffix = generate_suffix( flow_net, metric_net, z_cond, B, S - PROMPT_LEN, D, device, mask=suffix_mask, start_mlp=start_mlp, z_target_start=z_real_suffix if STRUCTURED_TARGET_START else None, ) z_projected_suffix = z_gen_suffix if residual_refiner is not None: pos_res = suffix_positions(B, S - PROMPT_LEN, device, z_real.dtype) z_projected_suffix, _residual_delta = residual_refiner( z_projected_suffix, z_real[:, :PROMPT_LEN, :], pos_res, suffix_mask, ) if latent_projector is not None: z_projected_suffix, projector_delta = latent_projector( z_gen_suffix, z_real[:, :PROMPT_LEN, :], suffix_mask, ) projector_delta_norm = projector_delta[suffix_mask].norm(dim=-1).mean().item() else: projector_delta_norm = 0.0 z_gen_flat = z_gen_suffix[suffix_mask] z_projected_flat = z_projected_suffix[suffix_mask] z_initial_flat = z_initial_suffix[suffix_mask] z_uncalibrated_flat = z_uncalibrated_suffix[suffix_mask] real_mean = z_real_flat.mean().item() real_std = z_real_flat.std().item() real_norm = z_real_flat.norm(dim=-1).mean().item() initial_mean = z_initial_flat.mean().item() initial_std = z_initial_flat.std().item() initial_norm = z_initial_flat.norm(dim=-1).mean().item() uncal_mean = z_uncalibrated_flat.mean().item() uncal_std = z_uncalibrated_flat.std().item() uncal_norm = z_uncalibrated_flat.norm(dim=-1).mean().item() gen_mean, gen_std, cosine_sim, gen_decode_ce = generated_decode_stats( z_real, z_projected_suffix, suffix_mask, input_ids, attn_mask, decoder, device, ) mix_decode_ce = latent_mix_decode_ce( z_real, z_gen_suffix, input_ids, attn_mask, decoder, device, ) raw_gen_mean = z_gen_flat.mean().item() raw_gen_std = z_gen_flat.std().item() projected_norm = z_projected_flat.norm(dim=-1).mean().item() if start_mlp is not None: pos_start = suffix_positions(B, S - PROMPT_LEN, device, z_real.dtype) if z_draft_start is not None and hasattr(start_mlp, "set_draft_target"): start_mlp.set_draft_target(z_draft_start) z_coarse_suffix = start_mlp(z_cond, pos_start, suffix_mask) coarse_mean, coarse_std, coarse_cosine, coarse_decode_ce = generated_decode_stats( z_real, z_coarse_suffix, suffix_mask, input_ids, attn_mask, decoder, device, ) coarse_mse = F.mse_loss(z_coarse_suffix[suffix_mask], z_real_suffix[suffix_mask]).item() coarse_decode_idx = (attn_mask[:, PROMPT_LEN:].sum(dim=1) > 0).nonzero(as_tuple=False).flatten() coarse_decode_idx = coarse_decode_idx[:DECODE_LOSS_BATCH] if coarse_decode_idx.numel() > 0: coarse_seq = torch.cat([z_real[:, :PROMPT_LEN, :], z_coarse_suffix], dim=1)[coarse_decode_idx] coarse_targets = input_ids[coarse_decode_idx, PROMPT_LEN:] coarse_valid = coarse_targets != 0 with torch.amp.autocast("cuda", enabled=device.type == "cuda"): coarse_logits = decoder.decode_from_latent(coarse_seq)[:, PROMPT_LEN:, :].float() coarse_probs = coarse_logits.softmax(dim=-1) coarse_target_ids = coarse_targets.clamp(0, coarse_probs.size(-1) - 1).unsqueeze(-1) coarse_target_probs = coarse_probs.gather(dim=-1, index=coarse_target_ids).squeeze(-1) coarse_target_prob = coarse_target_probs[coarse_valid].mean().item() if coarse_valid.any() else 0.0 else: coarse_target_prob = 0.0 else: coarse_mean = 0.0 coarse_std = 0.0 coarse_cosine = 0.0 coarse_decode_ce = 0.0 coarse_mse = 0.0 coarse_target_prob = 0.0 _, _, _, initial_decode_ce = generated_decode_stats( z_real, z_initial_suffix, suffix_mask, input_ids, attn_mask, decoder, device, ) _, _, _, uncal_decode_ce = generated_decode_stats( z_real, z_uncalibrated_suffix, suffix_mask, input_ids, attn_mask, decoder, device, ) metric_valid = metric_snapshot[suffix_mask] if metric_snapshot is not None else z_gen_flat.new_ones(z_gen_flat.shape) metric_mean = metric_valid.mean().item() metric_std = metric_valid.std().item() metric_min = metric_valid.min().item() metric_max = metric_valid.max().item() decode_idx = (attn_mask[:, PROMPT_LEN:].sum(dim=1) > 0).nonzero(as_tuple=False).flatten() decode_idx = decode_idx[:DECODE_LOSS_BATCH] if decode_idx.numel() > 0: z_decode_real = z_real[decode_idx] z_decode_gen = torch.cat([z_real[:, :PROMPT_LEN, :], z_projected_suffix], dim=1)[decode_idx] decode_targets = input_ids[decode_idx, PROMPT_LEN:].reshape(-1) with torch.amp.autocast("cuda", enabled=device.type == "cuda"): real_decode_logits = decoder.decode_from_latent(z_decode_real) gen_decode_logits = decoder.decode_from_latent(z_decode_gen) fused_decode_logits = fuse_aux_logits( flow_net, aux_token_head, gen_decode_logits, z_projected_suffix[decode_idx], z_cond[decode_idx], suffix_mask[decode_idx], ) real_decode_ce = F.cross_entropy( real_decode_logits[:, PROMPT_LEN:, :].reshape(-1, real_decode_logits.size(-1)), decode_targets, ignore_index=0, ).item() fused_decode_ce = F.cross_entropy( fused_decode_logits[:, PROMPT_LEN:, :].reshape(-1, fused_decode_logits.size(-1)), decode_targets, ignore_index=0, ).item() fused_suffix_logits = fused_decode_logits[:, PROMPT_LEN:, :].float() fused_targets = input_ids[decode_idx, PROMPT_LEN:] fused_valid = fused_targets != 0 if fused_valid.any(): fused_probs = fused_suffix_logits.softmax(dim=-1) fused_target_ids = fused_targets.clamp(0, fused_probs.size(-1) - 1).unsqueeze(-1) fused_decode_target_prob = ( fused_probs.gather(dim=-1, index=fused_target_ids).squeeze(-1)[fused_valid] .mean() .item() ) else: fused_decode_target_prob = 0.0 else: real_decode_ce = 0.0 fused_decode_ce = gen_decode_ce fused_decode_target_prob = 0.0 decode_ce_gap = gen_decode_ce - real_decode_ce fused_decode_ce_gap = fused_decode_ce - real_decode_ce with torch.no_grad(): sample_idx = (attn_mask[:, PROMPT_LEN:].sum(dim=1) > 0).nonzero(as_tuple=False).flatten() sample_idx = sample_idx[:n_samples] z_gen_seq = torch.cat([z_real[:, :PROMPT_LEN, :], z_projected_suffix], dim=1)[sample_idx] with torch.amp.autocast("cuda", enabled=device.type == "cuda"): raw_logits = decoder.decode_from_latent(z_gen_seq) logits = fuse_aux_logits( flow_net, aux_token_head, raw_logits, z_projected_suffix[sample_idx], z_cond[sample_idx], attn_mask[sample_idx, PROMPT_LEN:], ) oracle_logits = decoder.decode_from_latent(z_real[sample_idx]) pred_ids = logits.argmax(-1) sampled_ids = sample_token_ids(logits, tokenizer) oracle_ids = oracle_logits.argmax(-1) gen_collapse = argmax_token_collapse_stats(logits, pred_ids, tokenizer) sampled_collapse = argmax_token_collapse_stats(logits, sampled_ids, tokenizer) oracle_collapse = argmax_token_collapse_stats(oracle_logits, oracle_ids, tokenizer) sample_suffix_mask = attn_mask[sample_idx, PROMPT_LEN:] sample_target_ids = input_ids[sample_idx] gen_dist = decoder_distribution_stats( logits, tokenizer, sample_suffix_mask, oracle_logits=oracle_logits, target_ids=sample_target_ids, ) raw_gen_dist = decoder_distribution_stats( raw_logits, tokenizer, sample_suffix_mask, oracle_logits=oracle_logits, target_ids=sample_target_ids, ) oracle_dist = decoder_distribution_stats( oracle_logits, tokenizer, sample_suffix_mask, target_ids=sample_target_ids, ) print("\n-- riemannian samples -----------------------------------------") for sample_pos, batch_idx in enumerate(sample_idx.tolist()): prompt = tokenizer.decode(input_ids[batch_idx, :PROMPT_LEN], skip_special_tokens=True) target = tokenizer.decode(input_ids[batch_idx, PROMPT_LEN:], skip_special_tokens=True) sampled = tokenizer.decode(sampled_ids[sample_pos, PROMPT_LEN:], skip_special_tokens=True) argmax = tokenizer.decode(pred_ids[sample_pos, PROMPT_LEN:], skip_special_tokens=True) oracle = tokenizer.decode(oracle_ids[sample_pos, PROMPT_LEN:], skip_special_tokens=True) print(f" prompt: {prompt}") print(f" target: {target[:120]}") print(f" oracle: {oracle[:120]}") print(f" generated: {sampled[:120]}") print(f" argmax: {argmax[:120]}") print() print( " collapse argmax: " f"entropy={gen_collapse['entropy']:.2f} " f"uniq={gen_collapse['unique_ratio']:.3f} " f"maxfrac={gen_collapse['max_frac']:.3f} " f"top={', '.join(gen_collapse['top_tokens'])}" ) print( " collapse gen : " f"entropy={sampled_collapse['entropy']:.2f} " f"uniq={sampled_collapse['unique_ratio']:.3f} " f"maxfrac={sampled_collapse['max_frac']:.3f} " f"top={', '.join(sampled_collapse['top_tokens'])}" ) print( " collapse oracle: " f"entropy={oracle_collapse['entropy']:.2f} " f"uniq={oracle_collapse['unique_ratio']:.3f} " f"maxfrac={oracle_collapse['max_frac']:.3f} " f"top={', '.join(oracle_collapse['top_tokens'])}" ) print( " dist gen : " f"ent={gen_dist['entropy']:.2f} " f"top1={gen_dist['top1_prob']:.3f} " f"top5={gen_dist['top5_mass']:.3f} " f"top50={gen_dist['top50_mass']:.3f} " f"batch_top8={gen_dist['batch_top8_mass']:.3f} " f"special={gen_dist['special_mass']:.3f} " f"oracle_to_gen_kl={gen_dist['kl_to_oracle']:.3f}" ) print( " target raw : " f"top1_acc={raw_gen_dist['top1_acc']:.3f} " f"target_prob={raw_gen_dist['target_prob']:.4f}" ) print( " target gen : " f"top1_acc={gen_dist['top1_acc']:.3f} " f"target_prob={gen_dist['target_prob']:.4f} " f"oracle_top1_acc={gen_dist['oracle_top1_acc']:.3f} " f"oracle_target_prob={gen_dist['oracle_target_prob']:.4f}" ) print( " dist oracle : " f"ent={oracle_dist['entropy']:.2f} " f"top1={oracle_dist['top1_prob']:.3f} " f"top5={oracle_dist['top5_mass']:.3f} " f"top50={oracle_dist['top50_mass']:.3f} " f"batch_top8={oracle_dist['batch_top8_mass']:.3f} " f"special={oracle_dist['special_mass']:.3f} " f"target_prob={oracle_dist['target_prob']:.4f} " f"top1_acc={oracle_dist['top1_acc']:.3f}" ) print() torch.random.set_rng_state(eval_rng_state) if cuda_rng_state is not None: torch.cuda.set_rng_state_all(cuda_rng_state) latent_std_gap = abs(gen_std - real_std) raw_norm_gap = abs(uncal_norm - real_norm) collapse_uniq_penalty = max(0.0, COLLAPSE_UNIQ_TARGET - gen_collapse["unique_ratio"]) collapse_maxfrac_penalty = max(0.0, gen_collapse["max_frac"] - COLLAPSE_MAXFRAC_TARGET) val_score = ( avg_val_loss + latent_std_gap + max(0.0, 0.8 - cosine_sim) + 0.05 * max(0.0, fused_decode_ce_gap if aux_token_head is not None else decode_ce_gap) + RAW_NORM_GAP_SCORE_WEIGHT * raw_norm_gap + COLLAPSE_UNIQ_SCORE_WEIGHT * collapse_uniq_penalty + COLLAPSE_MAXFRAC_SCORE_WEIGHT * collapse_maxfrac_penalty ) print("-- val metrics ------------------------------------------------") print(f" val loss : {avg_val_loss:.4f}") print(f" real latents : mean={real_mean:.3f} std={real_std:.3f} norm={real_norm:.3f}") print(f" gen latents : mean={gen_mean:.3f} std={gen_std:.3f}") if start_mlp is not None: print(f" coarse start : mean={coarse_mean:.3f} std={coarse_std:.3f} mse={coarse_mse:.4f} cos={coarse_cosine:.4f} ce={coarse_decode_ce:.4f} p={coarse_target_prob:.4f}") if latent_projector is not None: print(f" raw gen lat : mean={raw_gen_mean:.3f} std={raw_gen_std:.3f}") print(f" projector : delta_norm={projector_delta_norm:.3f} proj_norm={projected_norm:.3f}") print(f" init latents : mean={initial_mean:.3f} std={initial_std:.3f} norm={initial_norm:.3f}") print(f" raw flow lat : mean={uncal_mean:.3f} std={uncal_std:.3f} norm={uncal_norm:.3f} norm_gap={raw_norm_gap:.3f}") print(f" metric diag : mean={metric_mean:.3f} std={metric_std:.3f} min={metric_min:.3f} max={metric_max:.3f}") print(f" cosine sim : {cosine_sim:.4f}") print(f" decoder CE : real={real_decode_ce:.4f} init={initial_decode_ce:.4f} raw={uncal_decode_ce:.4f} gen={gen_decode_ce:.4f} gap={decode_ce_gap:.4f}") if aux_token_head is not None: print( f" fused CE : beta={AUX_LOGIT_FUSION_BETA:.3f} gen={fused_decode_ce:.4f} " f"gap={fused_decode_ce_gap:.4f} p={fused_decode_target_prob:.4f}" ) if mix_decode_ce: mix_text = " ".join(f"a={alpha:.2f}:{ce:.3f}" for alpha, ce in mix_decode_ce) print(f" mix CE raw : {mix_text}") print(f" collapse pen : uniq={collapse_uniq_penalty:.4f} maxfrac={collapse_maxfrac_penalty:.4f}") print(f" ode steps : {ODE_STEPS}") print(f" val score : {val_score:.4f}") print() return avg_val_loss, val_score