| """ |
| Paper-Quality Combined Analysis Figures for PMPS |
| |
| Generates the main analysis figures for the AAAI paper: |
| Fig A: Three-way comparison (CODI vs Sim-CoT vs PMPS) — the hero figure |
| Fig B: Matching weight case study with CoT annotation |
| Fig C: Training dynamics (collapse metrics over training steps) |
| |
| Usage: |
| python analysis/plot_paper_figures.py \ |
| --pmps_dir analysis/results/pmps \ |
| --codi_dir analysis/results/codi \ |
| --simcot_dir analysis/results/simcot \ |
| --output_dir analysis/paper_figures |
| """ |
|
|
| import os |
| import json |
| import argparse |
| import numpy as np |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| import matplotlib.gridspec as gridspec |
| from matplotlib.colors import LinearSegmentedColormap |
| from matplotlib.patches import FancyBboxPatch |
|
|
| plt.rcParams.update({ |
| 'font.family': 'serif', |
| 'font.serif': ['Times New Roman', 'DejaVu Serif'], |
| 'font.size': 10, |
| 'axes.titlesize': 11, |
| 'axes.labelsize': 10, |
| 'xtick.labelsize': 9, |
| 'ytick.labelsize': 9, |
| 'legend.fontsize': 9, |
| 'figure.dpi': 300, |
| 'savefig.bbox': 'tight', |
| 'savefig.pad_inches': 0.05, |
| 'axes.linewidth': 0.8, |
| }) |
|
|
| COLORS = { |
| 'codi': '#E74C3C', |
| 'pmps': '#2E86C1', |
| 'simcot': '#27AE60', |
| } |
| METHOD_LABELS = { |
| 'codi': 'CODI', |
| 'pmps': 'PMPS (Ours)', |
| 'simcot': 'Sim-CoT', |
| } |
|
|
|
|
| def load_metrics(result_dir): |
| with open(os.path.join(result_dir, 'metrics.json'), 'r') as f: |
| return json.load(f) |
|
|
|
|
| def load_dist_metrics(result_dir): |
| path = os.path.join(result_dir, 'distribution_metrics.json') |
| if os.path.exists(path): |
| with open(path, 'r') as f: |
| return json.load(f) |
| return None |
|
|
|
|
| def load_latent_hiddens(result_dir): |
| return np.load(os.path.join(result_dir, 'latent_hiddens.npy')) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
| def plot_hero_figure(metrics_dict, dist_metrics_dict, output_path, |
| methods=['codi', 'simcot', 'pmps']): |
| """The main analysis figure for the paper.""" |
|
|
| available_methods = [m for m in methods if m in metrics_dict] |
|
|
| fig = plt.figure(figsize=(14, 4.8)) |
| gs = gridspec.GridSpec(1, 3, figure=fig, wspace=0.35, |
| width_ratios=[2.5, 1, 1.5]) |
|
|
| |
| gs_heatmaps = gs[0].subgridspec(1, len(available_methods), wspace=0.15) |
|
|
| cmap = LinearSegmentedColormap.from_list( |
| 'sim', ['#FFFFFF', '#AED6F1', '#3498DB', '#1A5276'], N=256 |
| ) |
|
|
| for m_idx, method in enumerate(available_methods): |
| ax = fig.add_subplot(gs_heatmaps[m_idx]) |
| metrics = metrics_dict[method] |
| sim_matrices = [np.array(m['cos_sim_matrix']) for m in metrics] |
| avg_matrix = np.mean(sim_matrices, axis=0) |
| N_l = avg_matrix.shape[0] |
|
|
| im = ax.imshow(avg_matrix, cmap=cmap, vmin=0.0, vmax=1.0, aspect='equal') |
|
|
| for i in range(N_l): |
| for j in range(N_l): |
| color = 'white' if avg_matrix[i, j] > 0.65 else 'black' |
| ax.text(j, i, f'{avg_matrix[i, j]:.2f}', |
| ha='center', va='center', fontsize=7, color=color) |
|
|
| ax.set_title(METHOD_LABELS.get(method, method), fontsize=11, fontweight='bold', |
| color=COLORS.get(method, 'black')) |
| ax.set_xticks(range(N_l)) |
| ax.set_yticks(range(N_l)) |
| ax.set_xticklabels([f'{i+1}' for i in range(N_l)], fontsize=8) |
| if m_idx == 0: |
| ax.set_yticklabels([f'{i+1}' for i in range(N_l)], fontsize=8) |
| ax.set_ylabel('Latent Index', fontsize=10) |
| else: |
| ax.set_yticklabels([]) |
| ax.set_xlabel('Latent Index', fontsize=10) |
|
|
| avg_offdiag = np.mean([m['avg_off_diag_sim'] for m in metrics]) |
| ax.text(0.5, -0.22, f'Avg off-diag: {avg_offdiag:.3f}', |
| transform=ax.transAxes, ha='center', fontsize=8, |
| color=COLORS.get(method, 'black'), fontweight='bold') |
|
|
| |
| ax_rank = fig.add_subplot(gs[1]) |
|
|
| for m_idx, method in enumerate(available_methods): |
| ranks = [m['effective_rank'] for m in metrics_dict[method]] |
| mean_r = np.mean(ranks) |
| std_r = np.std(ranks) |
| color = COLORS.get(method, '#888') |
|
|
| ax_rank.bar(m_idx, mean_r, width=0.6, yerr=std_r, capsize=4, |
| color=color, alpha=0.75) |
| ax_rank.text(m_idx, mean_r + std_r + 0.08, f'{mean_r:.2f}', |
| ha='center', fontsize=9, fontweight='bold', color=color) |
|
|
| N_l = len(metrics_dict[available_methods[0]][0]['cos_sim_matrix']) |
| ax_rank.axhline(y=N_l, color='gray', linestyle='--', alpha=0.4) |
| ax_rank.text(len(available_methods) - 0.3, N_l + 0.05, f'Max={N_l}', |
| color='gray', fontsize=8) |
| ax_rank.axhline(y=1.0, color='#E74C3C', linestyle=':', alpha=0.4) |
| ax_rank.text(len(available_methods) - 0.3, 1.1, 'Collapse', color='#E74C3C', fontsize=8) |
|
|
| ax_rank.set_xticks(range(len(available_methods))) |
| ax_rank.set_xticklabels([METHOD_LABELS.get(m, m) for m in available_methods], |
| fontsize=9, rotation=15) |
| ax_rank.set_ylabel('Effective Rank') |
| ax_rank.set_title('(b) Effective Rank', fontsize=11) |
| ax_rank.set_ylim(0, N_l + 1.5) |
|
|
| |
| ax_load = fig.add_subplot(gs[2]) |
|
|
| load_methods = [m for m in available_methods if m in dist_metrics_dict and dist_metrics_dict[m]] |
|
|
| if load_methods: |
| bar_width = 0.7 / len(load_methods) |
| for lm_idx, method in enumerate(load_methods): |
| dist_m = dist_metrics_dict[method] |
| loads = np.array([m['info_load_per_latent'] for m in dist_m]) |
| mean_loads = loads.mean(axis=0) |
| std_loads = loads.std(axis=0) |
| N_l = len(mean_loads) |
| x = np.arange(N_l) + lm_idx * bar_width |
| color = COLORS.get(method, '#888') |
| ax_load.bar(x, mean_loads, bar_width, yerr=std_loads, capsize=2, |
| color=color, alpha=0.75, |
| label=METHOD_LABELS.get(method, method)) |
|
|
| ideal = mean_loads.sum() / N_l |
| ax_load.axhline(y=ideal, color='gray', linestyle='--', alpha=0.5, |
| label=f'Ideal = {ideal:.1f}') |
| ax_load.set_xticks(np.arange(N_l) + bar_width * (len(load_methods) - 1) / 2) |
| ax_load.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)], fontsize=9) |
| ax_load.set_xlabel('Latent Index') |
| ax_load.set_ylabel('Information Load') |
| ax_load.set_title('(c) Info Distribution', fontsize=11) |
| ax_load.legend(fontsize=8, loc='upper right') |
| else: |
| ax_load.text(0.5, 0.5, 'No distribution\ndata available', |
| transform=ax_load.transAxes, ha='center', va='center', |
| fontsize=11, color='gray') |
| ax_load.set_title('(c) Info Distribution', fontsize=11) |
|
|
| |
| fig.text(0.18, 1.02, '(a) Latent Pairwise Cosine Similarity', |
| fontsize=11, ha='center', transform=fig.transFigure) |
|
|
| plt.savefig(output_path, dpi=300) |
| plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) |
| plt.close() |
| print(f' Saved: {output_path}') |
|
|
|
|
| |
| |
| |
|
|
| def plot_matching_case_study(pmps_internals_path, output_path, num_cases=2): |
| """Detailed matching weight visualization for individual samples.""" |
| with open(pmps_internals_path, 'r') as f: |
| pmps_internals = json.load(f) |
|
|
| n_cases = min(num_cases, len(pmps_internals)) |
| fig, axes = plt.subplots(1, n_cases, figsize=(5.5 * n_cases, 5)) |
| if n_cases == 1: |
| axes = [axes] |
|
|
| cmap = LinearSegmentedColormap.from_list( |
| 'weight', ['#FDFEFE', '#D5F5E3', '#27AE60', '#1E8449', '#145A32'], N=256 |
| ) |
|
|
| for idx in range(n_cases): |
| sample = pmps_internals[idx] |
| A_c2l = np.array(sample['A_c2l']) |
| N_c, N_l = A_c2l.shape |
| info_load = np.array(sample['info_load']) |
|
|
| ax = axes[idx] |
| im = ax.imshow(A_c2l, cmap=cmap, aspect='auto', vmin=0, interpolation='nearest') |
|
|
| ax.set_xlabel('Latent Token') |
| ax.set_ylabel('CoT Token Position') |
| ax.set_title(f'Sample #{sample["sample_idx"]}', fontsize=12, fontweight='bold') |
| ax.set_xticks(range(N_l)) |
| ax.set_xticklabels([f'$l_{{{i+1}}}$\n({info_load[i]:.1f})' for i in range(N_l)], |
| fontsize=9) |
|
|
| if N_c > 40: |
| tick_step = max(1, N_c // 10) |
| ax.set_yticks(range(0, N_c, tick_step)) |
|
|
| fig.colorbar(im, ax=axes, label='$A_{c2l}$ Weight', shrink=0.8, pad=0.03) |
| fig.suptitle('PMPS Matching Weights: CoT Token → Latent Token\n' |
| '(Numbers below latent labels show information load)', |
| fontsize=12, y=1.05) |
|
|
| plt.tight_layout() |
| plt.savefig(output_path, dpi=300) |
| plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) |
| plt.close() |
| print(f' Saved: {output_path}') |
|
|
|
|
| |
| |
| |
|
|
| def plot_sv_compact(metrics_dict, output_path, methods=None): |
| """Compact singular value comparison for supplementary material.""" |
| if methods is None: |
| methods = list(metrics_dict.keys()) |
|
|
| fig, axes = plt.subplots(1, 2, figsize=(10, 4)) |
|
|
| |
| ax = axes[0] |
| for method in methods: |
| sv_norms = np.array([m['singular_values_normalized'] for m in metrics_dict[method]]) |
| mean_sv = sv_norms.mean(axis=0) |
| std_sv = sv_norms.std(axis=0) |
| N_l = len(mean_sv) |
| x = np.arange(1, N_l + 1) |
| color = COLORS.get(method, '#888') |
| ax.plot(x, mean_sv, 'o-', color=color, linewidth=2, markersize=6, |
| label=METHOD_LABELS.get(method, method)) |
| ax.fill_between(x, mean_sv - std_sv, mean_sv + std_sv, |
| alpha=0.15, color=color) |
|
|
| ax.set_xlabel('Singular Value Index') |
| ax.set_ylabel('Normalized Magnitude') |
| ax.set_title('(a) Singular Value Spectrum') |
| ax.legend() |
| ax.set_xticks(range(1, N_l + 1)) |
|
|
| |
| ax2 = axes[1] |
| for method in methods: |
| ranks = [m['effective_rank'] for m in metrics_dict[method]] |
| color = COLORS.get(method, '#888') |
| ax2.hist(ranks, bins=20, color=color, alpha=0.4, edgecolor=color, |
| linewidth=1.2, label=METHOD_LABELS.get(method, method), density=True) |
| ax2.axvline(x=np.mean(ranks), color=color, linestyle='--', linewidth=2) |
|
|
| ax2.set_xlabel('Effective Rank') |
| ax2.set_ylabel('Density') |
| ax2.set_title('(b) Effective Rank Distribution') |
| ax2.legend() |
|
|
| plt.tight_layout() |
| plt.savefig(output_path, dpi=300) |
| plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) |
| plt.close() |
| print(f' Saved: {output_path}') |
|
|
|
|
| |
| |
| |
|
|
| def plot_prototype_analysis(pmps_dist_metrics, output_path): |
| """Prototype utilization and entropy analysis.""" |
| fig, axes = plt.subplots(1, 3, figsize=(14, 4)) |
|
|
| |
| ax = axes[0] |
| all_usage_lat = np.array([m['proto_usage_latent'] for m in pmps_dist_metrics]) |
| avg_lat = all_usage_lat.mean(axis=0) |
| K = len(avg_lat) |
| ax.bar(range(K), avg_lat, color=COLORS['pmps'], alpha=0.7, edgecolor='white') |
| ax.axhline(y=1.0/K, color='red', linestyle='--', alpha=0.6, label='Uniform') |
| ax.set_xlabel('Prototype Index') |
| ax.set_ylabel('Avg Weight') |
| ax.set_title('(a) Proto Usage (Latent)') |
| ax.legend(fontsize=8) |
|
|
| |
| ax2 = axes[1] |
| all_usage_cot = np.array([m['proto_usage_cot'] for m in pmps_dist_metrics]) |
| avg_cot = all_usage_cot.mean(axis=0) |
| ax2.bar(range(K), avg_cot, color=COLORS['simcot'], alpha=0.7, edgecolor='white') |
| ax2.axhline(y=1.0/K, color='red', linestyle='--', alpha=0.6, label='Uniform') |
| ax2.set_xlabel('Prototype Index') |
| ax2.set_ylabel('Avg Weight') |
| ax2.set_title('(b) Proto Usage (CoT)') |
| ax2.legend(fontsize=8) |
|
|
| |
| ax3 = axes[2] |
| lat_ents = [m['proto_entropy_latent'] for m in pmps_dist_metrics] |
| cot_ents = [m['proto_entropy_cot'] for m in pmps_dist_metrics] |
| max_ent = pmps_dist_metrics[0]['max_entropy'] |
|
|
| ax3.scatter(lat_ents, cot_ents, alpha=0.3, s=15, color=COLORS['pmps']) |
| ax3.plot([0, max_ent], [0, max_ent], 'k--', alpha=0.3, label='y=x') |
| ax3.axhline(y=max_ent, color='red', linestyle=':', alpha=0.3) |
| ax3.axvline(x=max_ent, color='red', linestyle=':', alpha=0.3) |
| ax3.set_xlabel('Latent Proto Entropy') |
| ax3.set_ylabel('CoT Proto Entropy') |
| ax3.set_title(f'(c) Entropy (max={max_ent:.2f})') |
| ax3.legend(fontsize=8) |
| ax3.set_aspect('equal') |
|
|
| plt.tight_layout() |
| plt.savefig(output_path, dpi=300) |
| plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) |
| plt.close() |
| print(f' Saved: {output_path}') |
|
|
|
|
| |
| |
| |
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--pmps_dir', type=str, required=True) |
| parser.add_argument('--codi_dir', type=str, required=True) |
| parser.add_argument('--simcot_dir', type=str, default=None) |
| parser.add_argument('--output_dir', type=str, default='analysis/paper_figures') |
| args = parser.parse_args() |
|
|
| os.makedirs(args.output_dir, exist_ok=True) |
|
|
| |
| print('Loading results...') |
| methods = ['codi', 'pmps'] |
| dirs = {'codi': args.codi_dir, 'pmps': args.pmps_dir} |
| if args.simcot_dir: |
| methods.append('simcot') |
| dirs['simcot'] = args.simcot_dir |
|
|
| |
| method_order = ['codi'] |
| if 'simcot' in methods: |
| method_order.append('simcot') |
| method_order.append('pmps') |
|
|
| metrics_dict = {} |
| dist_metrics_dict = {} |
|
|
| for method in method_order: |
| metrics_dict[method] = load_metrics(dirs[method]) |
| dist_metrics_dict[method] = load_dist_metrics(dirs[method]) |
| print(f' {method}: {len(metrics_dict[method])} samples') |
|
|
| |
| print('\nGenerating paper figures...') |
|
|
| |
| plot_hero_figure( |
| metrics_dict, dist_metrics_dict, |
| os.path.join(args.output_dir, 'fig_analysis_main.png'), |
| methods=method_order, |
| ) |
|
|
| |
| plot_sv_compact( |
| metrics_dict, |
| os.path.join(args.output_dir, 'fig_singular_values.png'), |
| methods=method_order, |
| ) |
|
|
| |
| pmps_pmps_path = os.path.join(args.pmps_dir, 'pmps_internals.json') |
| if os.path.exists(pmps_pmps_path): |
| plot_matching_case_study( |
| pmps_pmps_path, |
| os.path.join(args.output_dir, 'fig_matching_case_study.png'), |
| num_cases=3, |
| ) |
|
|
| |
| if dist_metrics_dict.get('pmps'): |
| plot_prototype_analysis( |
| dist_metrics_dict['pmps'], |
| os.path.join(args.output_dir, 'fig_prototype_analysis.png'), |
| ) |
|
|
| print(f'\nAll paper figures saved to: {args.output_dir}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
