| """ |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| result = { |
| 'latent_hiddens': latent_hiddens.squeeze(0).cpu().float().numpy(), |
| } |
|
|
| |
| 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 |
|
|
| |
| 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) |
| cot_h = cot_hidden.unsqueeze(0).to(device) |
|
|
| B, N_l, D = lat_h.shape |
| _, N_c, _ = cot_h.shape |
| K = pmps.nmb_prototypes |
|
|
| lat_scores = pmps._project_and_score(lat_h) |
| cot_scores = pmps._project_and_score(cot_h) |
|
|
| 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) |
| A_l2c = F.softmax(match_logits.transpose(1, 2), dim=-1) |
|
|
| 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(), |
| 'Q_cot': Q_cot.squeeze(0).cpu().float().numpy(), |
| 'A_c2l': A_c2l.squeeze(0).cpu().float().numpy(), |
| 'A_l2c': A_l2c.squeeze(0).cpu().float().numpy(), |
| 'p_lat': p_lat.squeeze(0).cpu().float().numpy(), |
| 'p_cot': p_cot.squeeze(0).cpu().float().numpy(), |
| 'prototype_weights': pmps.prototypes.weight.data.cpu().float().numpy(), |
| } |
|
|
| 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 |
|
|
| |
| norms = np.linalg.norm(latent_hiddens, axis=1, keepdims=True) |
| normed = latent_hiddens / (norms + 1e-8) |
| cos_sim_matrix = normed @ normed.T |
|
|
| 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)) |
|
|
| |
| 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)) |
|
|
| |
| 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'] |
| Q_lat = pmps_internals['Q_lat'] |
| Q_cot = pmps_internals['Q_cot'] |
| N_c, N_l = A_c2l.shape |
| K = Q_lat.shape[1] |
|
|
| |
| info_load = A_c2l.sum(axis=0) |
| info_load_cv = float(np.std(info_load) / (np.mean(info_load) + 1e-8)) |
|
|
| |
| proto_usage_lat = Q_lat.mean(axis=0) |
| 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)) |
|
|
| |
| 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)) |
| )) |
|
|
| |
| 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'], |
| }) |
|
|
| |
| print(f"\nSaving results to {args.output_dir}...") |
|
|
| |
| all_latent_np = np.stack(all_latent_hiddens, axis=0) |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
