| """ |
| Information Distribution Analysis Visualization (PMPS vs Sim-CoT) |
| |
| Generates publication-quality figures for information distribution uniformity: |
| Fig 1: Information load per latent token (bar chart comparison) |
| Fig 2: Matching weight heatmap (PMPS, case studies) |
| Fig 3: Prototype utilization histogram |
| Fig 4: Latent-CoT coverage heatmap |
| |
| Usage: |
| python analysis/plot_distribution_analysis.py \ |
| --pmps_dir analysis/results/pmps \ |
| --simcot_dir analysis/results/simcot \ |
| --output_dir analysis/figures/distribution |
| """ |
|
|
| import os |
| import json |
| import argparse |
| import numpy as np |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| from matplotlib.colors import LinearSegmentedColormap |
|
|
| plt.rcParams.update({ |
| 'font.family': 'serif', |
| 'font.serif': ['Times New Roman', 'DejaVu Serif'], |
| 'font.size': 11, |
| 'axes.titlesize': 13, |
| 'axes.labelsize': 12, |
| 'xtick.labelsize': 10, |
| 'ytick.labelsize': 10, |
| 'legend.fontsize': 10, |
| 'figure.dpi': 300, |
| 'savefig.bbox': 'tight', |
| 'savefig.pad_inches': 0.05, |
| }) |
|
|
| COLORS = { |
| 'codi': '#E74C3C', |
| 'pmps': '#2E86C1', |
| 'simcot': '#27AE60', |
| } |
|
|
| METHOD_LABELS = { |
| 'codi': 'CODI', |
| 'pmps': 'PMPS (Ours)', |
| 'simcot': 'Sim-CoT', |
| } |
|
|
|
|
| |
| |
| |
|
|
| def plot_info_load_comparison(pmps_dist_metrics, simcot_info_loads, output_path, |
| num_latent=6): |
| """ |
| Compare information load per latent token between PMPS and Sim-CoT. |
| |
| PMPS: uses matching weight sums from A_c2l |
| Sim-CoT: uses uniform segmentation (N_c / N_l per latent, with remainder |
| going to the last latent) |
| """ |
| fig, axes = plt.subplots(1, 2, figsize=(11, 4.5)) |
|
|
| |
| ax = axes[0] |
|
|
| |
| pmps_loads = np.array([m['info_load_per_latent'] for m in pmps_dist_metrics]) |
| pmps_mean = pmps_loads.mean(axis=0) |
| pmps_std = pmps_loads.std(axis=0) |
| N_l = len(pmps_mean) |
|
|
| |
| if simcot_info_loads is not None: |
| simcot_loads = np.array(simcot_info_loads) |
| simcot_mean = simcot_loads.mean(axis=0) |
| simcot_std = simcot_loads.std(axis=0) |
| else: |
| simcot_mean = None |
|
|
| x = np.arange(N_l) |
| w = 0.35 |
|
|
| bars_pmps = ax.bar(x - w/2, pmps_mean, w, yerr=pmps_std, capsize=4, |
| color=COLORS['pmps'], alpha=0.8, label=METHOD_LABELS['pmps']) |
|
|
| if simcot_mean is not None: |
| bars_simcot = ax.bar(x + w/2, simcot_mean, w, yerr=simcot_std, capsize=4, |
| color=COLORS['simcot'], alpha=0.8, label=METHOD_LABELS['simcot']) |
|
|
| ideal_load = pmps_mean.sum() / N_l |
| ax.axhline(y=ideal_load, color='gray', linestyle='--', alpha=0.6, |
| label=f'Ideal (uniform = {ideal_load:.1f})') |
|
|
| ax.set_xlabel('Latent Token Index') |
| ax.set_ylabel('Information Load\n($\\sum_j A_{c2l}(j, i)$)') |
| ax.set_title('(a) Information Load per Latent', fontsize=13) |
| ax.set_xticks(x) |
| ax.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)]) |
| ax.legend(fontsize=9) |
|
|
| |
| ax2 = axes[1] |
|
|
| pmps_cvs = [m['info_load_cv'] for m in pmps_dist_metrics] |
|
|
| categories = [METHOD_LABELS['pmps']] |
| means = [np.mean(pmps_cvs)] |
| stds = [np.std(pmps_cvs)] |
| bar_colors = [COLORS['pmps']] |
| scatter_data = [pmps_cvs] |
|
|
| if simcot_info_loads is not None: |
| simcot_cvs = [np.std(loads) / (np.mean(loads) + 1e-8) for loads in simcot_info_loads] |
| categories.append(METHOD_LABELS['simcot']) |
| means.append(np.mean(simcot_cvs)) |
| stds.append(np.std(simcot_cvs)) |
| bar_colors.append(COLORS['simcot']) |
| scatter_data.append(simcot_cvs) |
|
|
| bars = ax2.bar(range(len(categories)), means, yerr=stds, capsize=5, |
| color=bar_colors, alpha=0.7, width=0.5) |
|
|
| for i, (d, color) in enumerate(zip(scatter_data, bar_colors)): |
| jitter = np.random.normal(0, 0.03, len(d)) |
| ax2.scatter(i + jitter, d, color=color, alpha=0.15, s=10, zorder=3) |
| ax2.text(i, means[i] + stds[i] + 0.02, f'{means[i]:.3f}', |
| ha='center', fontsize=11, fontweight='bold', color=color) |
|
|
| ax2.set_xticks(range(len(categories))) |
| ax2.set_xticklabels(categories, fontsize=12) |
| ax2.set_ylabel('Coefficient of Variation\n(Lower = More Uniform)', fontsize=11) |
| ax2.set_title('(b) Load Uniformity (CV)', fontsize=13) |
|
|
| 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_matching_weight_heatmaps(pmps_internals, output_path, num_cases=3): |
| """ |
| Heatmap of A_c2l for individual samples. |
| Shows which CoT tokens are assigned to which latent tokens. |
| """ |
| n_cases = min(num_cases, len(pmps_internals)) |
| fig, axes = plt.subplots(1, n_cases, figsize=(5 * n_cases, 5)) |
| if n_cases == 1: |
| axes = [axes] |
|
|
| cmap = LinearSegmentedColormap.from_list( |
| 'weight', ['#FFFFFF', '#E8F8F5', '#1ABC9C', '#117864'], 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 |
|
|
| ax = axes[idx] |
| im = ax.imshow(A_c2l, cmap=cmap, aspect='auto', vmin=0) |
|
|
| ax.set_xlabel('Latent Token Index') |
| 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}}}$' for i in range(N_l)]) |
|
|
| if N_c > 30: |
| tick_step = max(1, N_c // 10) |
| ax.set_yticks(range(0, N_c, tick_step)) |
|
|
| fig.colorbar(im, ax=axes, label='Matching Weight $A_{c2l}$', shrink=0.8, pad=0.02) |
| fig.suptitle('CoT → Latent Matching Weights (PMPS)', fontsize=14, y=1.02) |
|
|
| 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_utilization(pmps_dist_metrics, output_path): |
| """ |
| Prototype utilization histogram + entropy comparison. |
| """ |
| fig, axes = plt.subplots(1, 2, figsize=(11, 4.5)) |
|
|
| |
| ax = axes[0] |
| all_proto_usage = np.array([m['proto_usage_latent'] for m in pmps_dist_metrics]) |
| avg_usage = all_proto_usage.mean(axis=0) |
| std_usage = all_proto_usage.std(axis=0) |
| K = len(avg_usage) |
| uniform_level = 1.0 / K |
|
|
| ax.bar(range(K), avg_usage, yerr=std_usage, capsize=2, |
| color=COLORS['pmps'], alpha=0.7, edgecolor='white', linewidth=0.5) |
| ax.axhline(y=uniform_level, color='red', linestyle='--', alpha=0.6, |
| label=f'Uniform = {uniform_level:.4f}') |
| ax.set_xlabel('Prototype Index') |
| ax.set_ylabel('Average Assignment Weight') |
| ax.set_title('(a) Prototype Utilization (Latent Side)', fontsize=13) |
| ax.legend(fontsize=9) |
|
|
| if K > 16: |
| ax.set_xticks(range(0, K, max(1, K // 8))) |
|
|
| |
| ax2 = axes[1] |
| proto_ents = [m['proto_entropy_latent'] for m in pmps_dist_metrics] |
| max_ent = pmps_dist_metrics[0]['max_entropy'] |
|
|
| ax2.hist(proto_ents, bins=30, color=COLORS['pmps'], alpha=0.7, |
| edgecolor='white', linewidth=0.5, density=True) |
| ax2.axvline(x=max_ent, color='red', linestyle='--', linewidth=2, alpha=0.7, |
| label=f'Max Entropy = $\\log K$ = {max_ent:.2f}') |
| ax2.axvline(x=np.mean(proto_ents), color=COLORS['pmps'], linestyle='-', |
| linewidth=2, label=f'Mean = {np.mean(proto_ents):.2f}') |
|
|
| ax2.set_xlabel('Prototype Assignment Entropy') |
| ax2.set_ylabel('Density') |
| ax2.set_title('(b) Prototype Entropy Distribution', fontsize=13) |
| ax2.legend(fontsize=9) |
|
|
| 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_coverage_comparison(pmps_dist_metrics, simcot_info_loads, output_path): |
| """ |
| Compare how CoT information is distributed across latents. |
| |
| For PMPS: use matching weight entropy per CoT token (higher = distributed more evenly) |
| For Sim-CoT: compute rigid-segment coverage |
| """ |
| fig, ax = plt.subplots(figsize=(7, 4.5)) |
|
|
| |
| pmps_match_ents = [] |
| for m in pmps_dist_metrics: |
| pmps_match_ents.extend(m['match_entropies']) |
|
|
| data_list = [pmps_match_ents] |
| labels = [METHOD_LABELS['pmps']] |
| colors = [COLORS['pmps']] |
|
|
| |
| |
| if simcot_info_loads is not None: |
| simcot_ents = [0.0] * len(pmps_match_ents) |
| data_list.append(simcot_ents) |
| labels.append(METHOD_LABELS['simcot']) |
| colors.append(COLORS['simcot']) |
|
|
| parts = ax.violinplot(data_list, positions=range(len(labels)), |
| showmedians=True, showextrema=False) |
| for pc, color in zip(parts['bodies'], colors): |
| pc.set_facecolor(color) |
| pc.set_alpha(0.4) |
| parts['cmedians'].set_color('black') |
| parts['cmedians'].set_linewidth(2) |
|
|
| for i, (d, color) in enumerate(zip(data_list, colors)): |
| mean_val = np.mean(d) |
| ax.plot(i, mean_val, 'D', color='white', markeredgecolor=color, |
| markersize=8, markeredgewidth=2, zorder=5) |
| ax.annotate(f'{mean_val:.3f}', (i, mean_val), |
| textcoords='offset points', xytext=(18, 5), |
| fontsize=10, color=color, fontweight='bold') |
|
|
| ax.set_xticks(range(len(labels))) |
| ax.set_xticklabels(labels, fontsize=12) |
| ax.set_ylabel('Per-CoT-Token Matching Entropy\n(Higher = More Distributed)', fontsize=11) |
| ax.set_title('Matching Weight Distribution\n(How Diversely Each CoT Token Maps to Latents)', fontsize=13) |
|
|
| 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('--simcot_dir', type=str, default=None) |
| parser.add_argument('--output_dir', type=str, default='analysis/figures/distribution') |
| args = parser.parse_args() |
|
|
| os.makedirs(args.output_dir, exist_ok=True) |
|
|
| |
| print('Loading PMPS distribution metrics...') |
| pmps_dist_path = os.path.join(args.pmps_dir, 'distribution_metrics.json') |
| pmps_pmps_path = os.path.join(args.pmps_dir, 'pmps_internals.json') |
|
|
| with open(pmps_dist_path, 'r') as f: |
| pmps_dist_metrics = json.load(f) |
| print(f' PMPS: {len(pmps_dist_metrics)} samples with distribution metrics') |
|
|
| pmps_internals = None |
| if os.path.exists(pmps_pmps_path): |
| with open(pmps_pmps_path, 'r') as f: |
| pmps_internals = json.load(f) |
| print(f' PMPS: {len(pmps_internals)} samples with PMPS internals') |
|
|
| |
| simcot_info_loads = None |
| if args.simcot_dir: |
| simcot_dist_path = os.path.join(args.simcot_dir, 'distribution_metrics.json') |
| if os.path.exists(simcot_dist_path): |
| with open(simcot_dist_path, 'r') as f: |
| simcot_dist = json.load(f) |
| simcot_info_loads = [m['info_load_per_latent'] for m in simcot_dist] |
| print(f' Sim-CoT: {len(simcot_info_loads)} samples') |
| else: |
| print(' Sim-CoT: no distribution_metrics.json found, will simulate rigid segmentation') |
|
|
| print('\nGenerating figures...') |
|
|
| plot_info_load_comparison( |
| pmps_dist_metrics, |
| simcot_info_loads, |
| os.path.join(args.output_dir, 'fig_info_load.png'), |
| ) |
|
|
| if pmps_internals: |
| plot_matching_weight_heatmaps( |
| pmps_internals, |
| os.path.join(args.output_dir, 'fig_matching_weights.png'), |
| num_cases=3, |
| ) |
|
|
| plot_prototype_utilization( |
| pmps_dist_metrics, |
| os.path.join(args.output_dir, 'fig_prototype_utilization.png'), |
| ) |
|
|
| plot_coverage_comparison( |
| pmps_dist_metrics, |
| simcot_info_loads, |
| os.path.join(args.output_dir, 'fig_coverage_comparison.png'), |
| ) |
|
|
| print(f'\nAll figures saved to: {args.output_dir}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
