""" Latent Hidden State Extraction Script for PMPS Analysis Experiments This script loads a trained CODI/PMPS checkpoint and extracts: 1. Latent hidden states at each iteration step (for collapse analysis) 2. CoT hidden states from teacher forward pass (for alignment analysis) 3. Prototype assignments and matching weights (PMPS only, for distribution analysis) Usage: python analysis/extract_latent_states.py \ --config_file configs/test_pmps.yaml \ --output_dir analysis/results/pmps \ --num_samples 200 """ import os import sys import json import argparse import yaml import torch import torch.nn.functional as F import numpy as np from tqdm import tqdm from dataclasses import asdict sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import transformers from peft import LoraConfig, TaskType from src.model import ( CODI_CT, ModelArguments, DataArguments, TrainingArguments, load_safetensors_state_dict, ) def read_json(file_path): with open(file_path, 'r', encoding='utf-8') as f: return json.load(f) def load_model_from_config(config_file, ckpt_dir=None): """Load model and tokenizer from a YAML config file.""" with open(config_file, 'r') as f: config = yaml.safe_load(f) cmd_args = [] for key, value in config.items(): if value is None: continue if isinstance(value, bool): cmd_args.extend([f"--{key}", "True" if value else "False"]) else: cmd_args.extend([f"--{key}", str(value)]) parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=cmd_args) if ckpt_dir: model_args.ckpt_dir = ckpt_dir task_type = TaskType.CAUSAL_LM if any(name in model_args.model_name_or_path.lower() for name in ["llama", "mistral", "falcon", "qwen"]): target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"] elif "gpt2" in model_args.model_name_or_path.lower(): target_modules = ["c_attn", "c_proj", "c_fc"] else: raise ValueError(f"Unsupported model: {model_args.model_name_or_path}") lora_config = LoraConfig( task_type=task_type, inference_mode=False, r=model_args.lora_r, lora_alpha=model_args.lora_alpha, lora_dropout=0.1, target_modules=target_modules, init_lora_weights=True, ) training_args.train = False model_args.train = False model = CODI_CT(model_args, training_args, lora_config) state_dict = load_safetensors_state_dict(model_args.ckpt_dir) model.load_state_dict(state_dict, strict=False) model.codi.tie_weights() tokenizer = transformers.AutoTokenizer.from_pretrained( model_args.model_name_or_path, token=model_args.token, model_max_length=training_args.model_max_length, padding_side="left", use_fast=False, ) if tokenizer.pad_token_id is None: tokenizer.add_special_tokens({'pad_token': '[PAD]'}) tokenizer.pad_token_id = model.pad_token_id model.eval() model.to('cuda') model.to(torch.bfloat16) return model, tokenizer, model_args, data_args, training_args def load_test_data(data_name, data_path=None): """Load test dataset.""" if data_name == "gsm8k": return read_json(data_path or '/data/oss_bucket_0/liqi/gsk8k-aug/gsm_test.json') elif data_name == "math": return read_json(data_path or '/data/oss_bucket_0/liqi/dataset/MATH/test.json') else: raise NotImplementedError(f"Dataset {data_name} not supported") @torch.no_grad() def extract_latent_and_cot_states( model, tokenizer, question_text, cot_text, answer_text, training_args, return_pmps_internals=False, ): """ Extract latent hidden states and (optionally) CoT hidden states + PMPS internals. Returns a dict with: - latent_hiddens: (num_latent, hidden_dim) numpy array - cot_hiddens: (num_cot_tokens, hidden_dim) numpy array (if cot_text provided) - pmps_internals: dict with Q_lat, Q_cot, A_c2l, A_l2c, prototype_weights (only if return_pmps_internals=True and model has pmps_module) """ device = next(model.parameters()).device embed_layer = model.get_embd(model.codi, model.model_name) # Tokenize question q_tokens = tokenizer.encode(question_text, add_special_tokens=True) if training_args.remove_eos: q_tokens = q_tokens + [model.bot_id] else: q_tokens = q_tokens + [tokenizer.eos_token_id, model.bot_id] q_ids = torch.tensor([q_tokens], dtype=torch.long, device=device) q_mask = torch.ones_like(q_ids) # === Extract latent hidden states === outputs = model.codi( input_ids=q_ids, use_cache=True, output_hidden_states=True, attention_mask=q_mask, ) past_key_values = outputs.past_key_values if training_args.enable_top_k: logits_last = outputs.logits[:, -1, :] if training_args.latent_top_k == 0: latent_embd = model._compute_all_weighted_embedding(logits_last, embed_layer) else: latent_embd = model._compute_topk_weighted_embedding(logits_last, embed_layer) else: latent_embd = outputs.hidden_states[-1][:, -1, :].unsqueeze(1) if training_args.use_prj: latent_embd = model.prj(latent_embd) latent_hiddens_list = [] num_latent = training_args.num_latent for i in range(num_latent): outputs = model.codi( inputs_embeds=latent_embd, use_cache=True, output_hidden_states=True, past_key_values=past_key_values, ) past_key_values = outputs.past_key_values if training_args.enable_top_k: logits_last = outputs.logits[:, -1, :] if training_args.latent_top_k == 0: next_latent = model._compute_all_weighted_embedding(logits_last, embed_layer) else: next_latent = model._compute_topk_weighted_embedding(logits_last, embed_layer) else: next_latent = outputs.hidden_states[-1][:, -1, :].unsqueeze(1) if training_args.enable_top_k: latent_hiddens_list.append(next_latent.squeeze(1)) else: latent_hidden = model._extract_hidden_states_at_layer( outputs.hidden_states, model.pmps_layer if hasattr(model, 'pmps_layer') else 'last' )[:, -1, :] latent_hiddens_list.append(latent_hidden) if training_args.use_prj: latent_embd = model.prj(next_latent) else: latent_embd = next_latent latent_hiddens = torch.stack(latent_hiddens_list, dim=1) # (1, num_latent, D) result = { 'latent_hiddens': latent_hiddens.squeeze(0).cpu().float().numpy(), } # === Extract CoT hidden states (teacher forward pass) === if cot_text is not None and answer_text is not None: q_raw = tokenizer.encode(question_text, add_special_tokens=True) cot_raw = tokenizer.encode(cot_text, add_special_tokens=False) ans_raw = tokenizer.encode(answer_text, add_special_tokens=False) if len(q_raw) > 0 and q_raw[0] == tokenizer.bos_token_id: pass if len(cot_raw) > 0 and cot_raw[0] == tokenizer.bos_token_id: cot_raw = cot_raw[1:] if len(ans_raw) > 0 and ans_raw[0] == tokenizer.bos_token_id: ans_raw = ans_raw[1:] if not training_args.remove_eos: q_raw = q_raw + [tokenizer.eos_token_id] cot_raw = cot_raw + [tokenizer.eos_token_id] ans_raw = ans_raw + [tokenizer.eos_token_id] ref_ids = q_raw + cot_raw + ans_raw cot_start = len(q_raw) cot_end = len(q_raw) + len(cot_raw) ref_input = torch.tensor([ref_ids], dtype=torch.long, device=device) ref_mask = torch.ones_like(ref_input) ref_outputs = model.codi( input_ids=ref_input, output_hidden_states=True, attention_mask=ref_mask, ) pmps_layer = model.pmps_layer if hasattr(model, 'pmps_layer') else 'last' ref_hidden = model._extract_hidden_states_at_layer( ref_outputs.hidden_states, pmps_layer ) if training_args.enable_top_k: ref_logits = ref_outputs.logits cot_logits = ref_logits[0, cot_start - 1: cot_end - 1, :] if training_args.latent_top_k == 0: cot_hidden = model._compute_all_weighted_embedding(cot_logits, embed_layer).squeeze(0) else: cot_hidden = model._compute_topk_weighted_embedding(cot_logits, embed_layer).squeeze(0) else: cot_hidden = ref_hidden[0, cot_start - 1: cot_end - 1, :] result['cot_hiddens'] = cot_hidden.cpu().float().numpy() result['cot_start'] = cot_start result['cot_end'] = cot_end result['num_cot_tokens'] = cot_end - cot_start # === Extract PMPS internals === if return_pmps_internals and hasattr(model, 'pmps_module') and model.pmps_loss_factor > 0: pmps = model.pmps_module pmps.normalize_prototypes() lat_h = latent_hiddens.to(device) # (1, N_l, D) cot_h = cot_hidden.unsqueeze(0).to(device) # (1, N_c, D) B, N_l, D = lat_h.shape _, N_c, _ = cot_h.shape K = pmps.nmb_prototypes lat_scores = pmps._project_and_score(lat_h) # (N_l, K) cot_scores = pmps._project_and_score(cot_h) # (N_c, K) Q_lat = pmps.sinkhorn(lat_scores).reshape(1, N_l, K) Q_cot = pmps.sinkhorn(cot_scores).reshape(1, N_c, K) match_logits = torch.bmm(Q_cot, Q_lat.transpose(1, 2)) / pmps.match_temperature A_c2l = F.softmax(match_logits, dim=-1) # (1, N_c, N_l) A_l2c = F.softmax(match_logits.transpose(1, 2), dim=-1) # (1, N_l, N_c) p_lat = F.softmax(lat_scores.reshape(1, N_l, K) / pmps.temperature, dim=-1) p_cot = F.softmax(cot_scores.reshape(1, N_c, K) / pmps.temperature, dim=-1) result['pmps_internals'] = { 'Q_lat': Q_lat.squeeze(0).cpu().float().numpy(), # (N_l, K) 'Q_cot': Q_cot.squeeze(0).cpu().float().numpy(), # (N_c, K) 'A_c2l': A_c2l.squeeze(0).cpu().float().numpy(), # (N_c, N_l) 'A_l2c': A_l2c.squeeze(0).cpu().float().numpy(), # (N_l, N_c) 'p_lat': p_lat.squeeze(0).cpu().float().numpy(), # (N_l, K) 'p_cot': p_cot.squeeze(0).cpu().float().numpy(), # (N_c, K) 'prototype_weights': pmps.prototypes.weight.data.cpu().float().numpy(), # (K, d) } return result def compute_metrics(latent_hiddens): """ Compute collapse-related metrics from latent hidden states. Args: latent_hiddens: (N_l, D) numpy array Returns: dict with metrics """ N_l, D = latent_hiddens.shape # 1. Pairwise cosine similarity norms = np.linalg.norm(latent_hiddens, axis=1, keepdims=True) normed = latent_hiddens / (norms + 1e-8) cos_sim_matrix = normed @ normed.T # (N_l, N_l) mask = ~np.eye(N_l, dtype=bool) off_diag_sims = cos_sim_matrix[mask] avg_off_diag_sim = float(np.mean(off_diag_sims)) std_off_diag_sim = float(np.std(off_diag_sims)) # 2. Singular value analysis / Effective rank centered = latent_hiddens - latent_hiddens.mean(axis=0, keepdims=True) U, S, Vh = np.linalg.svd(centered, full_matrices=False) S_norm = S / (S.sum() + 1e-8) S_norm_clipped = np.clip(S_norm, 1e-12, None) entropy = -np.sum(S_norm_clipped * np.log(S_norm_clipped)) effective_rank = float(np.exp(entropy)) # 3. Cumulative explained variance explained_var = (S ** 2) / ((S ** 2).sum() + 1e-8) cumulative_var = np.cumsum(explained_var) return { 'cos_sim_matrix': cos_sim_matrix.tolist(), 'avg_off_diag_sim': avg_off_diag_sim, 'std_off_diag_sim': std_off_diag_sim, 'singular_values': S.tolist(), 'singular_values_normalized': S_norm.tolist(), 'effective_rank': effective_rank, 'explained_variance': explained_var.tolist(), 'cumulative_variance': cumulative_var.tolist(), } def compute_distribution_metrics(pmps_internals): """ Compute information distribution metrics from PMPS internals. Args: pmps_internals: dict with Q_lat, Q_cot, A_c2l, A_l2c, etc. Returns: dict with distribution metrics """ A_c2l = pmps_internals['A_c2l'] # (N_c, N_l) Q_lat = pmps_internals['Q_lat'] # (N_l, K) Q_cot = pmps_internals['Q_cot'] # (N_c, K) N_c, N_l = A_c2l.shape K = Q_lat.shape[1] # 1. Information load per latent token info_load = A_c2l.sum(axis=0) # (N_l,) — total CoT weight assigned to each latent info_load_cv = float(np.std(info_load) / (np.mean(info_load) + 1e-8)) # 2. Prototype utilization (from Q_lat) proto_usage_lat = Q_lat.mean(axis=0) # (K,) average assignment to each prototype proto_usage_lat_norm = proto_usage_lat / (proto_usage_lat.sum() + 1e-8) proto_entropy_lat = -float(np.sum( proto_usage_lat_norm * np.log(np.clip(proto_usage_lat_norm, 1e-12, None)) )) max_entropy = float(np.log(K)) # 3. Prototype utilization (from Q_cot) proto_usage_cot = Q_cot.mean(axis=0) proto_usage_cot_norm = proto_usage_cot / (proto_usage_cot.sum() + 1e-8) proto_entropy_cot = -float(np.sum( proto_usage_cot_norm * np.log(np.clip(proto_usage_cot_norm, 1e-12, None)) )) # 4. Matching weight entropy (per CoT token) match_entropies = [] for j in range(N_c): w = A_c2l[j] w_clipped = np.clip(w, 1e-12, None) ent = -float(np.sum(w * np.log(w_clipped))) match_entropies.append(ent) avg_match_entropy = float(np.mean(match_entropies)) return { 'info_load_per_latent': info_load.tolist(), 'info_load_cv': info_load_cv, 'proto_usage_latent': proto_usage_lat.tolist(), 'proto_entropy_latent': proto_entropy_lat, 'proto_usage_cot': proto_usage_cot.tolist(), 'proto_entropy_cot': proto_entropy_cot, 'max_entropy': max_entropy, 'avg_match_entropy': avg_match_entropy, 'match_entropies': match_entropies, } def main(): parser = argparse.ArgumentParser() parser.add_argument('--config_file', type=str, required=True) parser.add_argument('--ckpt_dir', type=str, default=None) parser.add_argument('--output_dir', type=str, default='analysis/results') parser.add_argument('--data_path', type=str, default=None) parser.add_argument('--num_samples', type=int, default=200) parser.add_argument('--extract_pmps', action='store_true', help='Extract PMPS internals (matching weights, prototype assignments)') args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) print("Loading model...") model, tokenizer, model_args, data_args, training_args = load_model_from_config( args.config_file, args.ckpt_dir ) print("Loading test data...") data_name = "gsm8k" if "math" in data_args.data_name: data_name = "math" test_data = load_test_data(data_name, args.data_path) all_latent_hiddens = [] all_cot_hiddens = [] all_metrics = [] all_distribution_metrics = [] all_pmps_internals = [] num_samples = min(args.num_samples, len(test_data)) print(f"Extracting latent states for {num_samples} samples...") for idx in tqdm(range(num_samples)): sample = test_data[idx] if "question" in sample: question = sample["question"].strip() elif "problem" in sample: question = sample["problem"].strip() else: continue if "answer" in sample: answer_raw = str(sample["answer"]).strip() elif "target" in sample: answer_raw = str(sample["target"]).strip() else: answer_raw = None answer_text = f"The answer is: {answer_raw}" if answer_raw else None cot_text = None if "cot" in sample: cot_text = sample["cot"].strip() elif "solution" in sample: cot_text = sample["solution"].strip() elif "steps" in sample: if isinstance(sample["steps"], list): cot_text = " ".join(sample["steps"]).strip() else: cot_text = str(sample["steps"]).strip() try: result = extract_latent_and_cot_states( model, tokenizer, question, cot_text, answer_text, training_args, return_pmps_internals=args.extract_pmps, ) except Exception as e: print(f" Sample {idx} failed: {e}") continue latent_h = result['latent_hiddens'] all_latent_hiddens.append(latent_h) metrics = compute_metrics(latent_h) metrics['sample_idx'] = idx all_metrics.append(metrics) if 'cot_hiddens' in result: all_cot_hiddens.append(result['cot_hiddens']) if 'pmps_internals' in result: dist_metrics = compute_distribution_metrics(result['pmps_internals']) dist_metrics['sample_idx'] = idx all_distribution_metrics.append(dist_metrics) all_pmps_internals.append({ 'sample_idx': idx, 'A_c2l': result['pmps_internals']['A_c2l'].tolist(), 'Q_lat': result['pmps_internals']['Q_lat'].tolist(), 'Q_cot': result['pmps_internals']['Q_cot'].tolist(), 'info_load': dist_metrics['info_load_per_latent'], }) # === Save results === print(f"\nSaving results to {args.output_dir}...") # Save latent hiddens as numpy all_latent_np = np.stack(all_latent_hiddens, axis=0) # (num_samples, N_l, D) np.save(os.path.join(args.output_dir, 'latent_hiddens.npy'), all_latent_np) if all_cot_hiddens: max_cot_len = max(h.shape[0] for h in all_cot_hiddens) D = all_cot_hiddens[0].shape[1] cot_padded = np.zeros((len(all_cot_hiddens), max_cot_len, D)) cot_lengths = [] for i, h in enumerate(all_cot_hiddens): cot_padded[i, :h.shape[0]] = h cot_lengths.append(h.shape[0]) np.save(os.path.join(args.output_dir, 'cot_hiddens.npy'), cot_padded) with open(os.path.join(args.output_dir, 'cot_lengths.json'), 'w') as f: json.dump(cot_lengths, f) # Save metrics with open(os.path.join(args.output_dir, 'metrics.json'), 'w') as f: json.dump(all_metrics, f, indent=2) if all_distribution_metrics: with open(os.path.join(args.output_dir, 'distribution_metrics.json'), 'w') as f: json.dump(all_distribution_metrics, f, indent=2) if all_pmps_internals: with open(os.path.join(args.output_dir, 'pmps_internals.json'), 'w') as f: json.dump(all_pmps_internals[:50], f, indent=2) # === Print summary === avg_sims = [m['avg_off_diag_sim'] for m in all_metrics] eff_ranks = [m['effective_rank'] for m in all_metrics] print(f"\n{'='*60}") print(f"SUMMARY ({len(all_metrics)} samples)") print(f"{'='*60}") print(f"Avg Off-Diagonal Cosine Similarity: {np.mean(avg_sims):.4f} ± {np.std(avg_sims):.4f}") print(f"Avg Effective Rank: {np.mean(eff_ranks):.4f} ± {np.std(eff_ranks):.4f}") if all_distribution_metrics: cvs = [m['info_load_cv'] for m in all_distribution_metrics] proto_ents = [m['proto_entropy_latent'] for m in all_distribution_metrics] max_ent = all_distribution_metrics[0]['max_entropy'] print(f"Info Load CV (lower=more uniform): {np.mean(cvs):.4f} ± {np.std(cvs):.4f}") print(f"Proto Entropy (max={max_ent:.4f}): {np.mean(proto_ents):.4f} ± {np.std(proto_ents):.4f}") print(f"\nResults saved to: {args.output_dir}") if __name__ == '__main__': main()