| """ |
| Collapse Analysis Visualization (PMPS vs CODI) |
| |
| Generates publication-quality figures for latent representation collapse analysis: |
| Fig 1: Pairwise cosine similarity heatmaps (side-by-side CODI vs PMPS) |
| Fig 2: Off-diagonal cosine similarity distribution (violin plot) |
| Fig 3: Singular value spectrum comparison |
| Fig 4: Effective rank distribution comparison |
| Fig 5: t-SNE / PCA visualization of latent embeddings |
| |
| Usage: |
| python analysis/plot_collapse_analysis.py \ |
| --pmps_dir analysis/results/pmps \ |
| --codi_dir analysis/results/codi \ |
| --output_dir analysis/figures/collapse |
| """ |
|
|
| 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 sklearn.decomposition import PCA |
| from sklearn.manifold import TSNE |
|
|
| |
| |
| |
|
|
| 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, |
| 'axes.linewidth': 0.8, |
| 'axes.grid': False, |
| }) |
|
|
| COLORS = { |
| 'codi': '#E74C3C', |
| 'pmps': '#2E86C1', |
| 'simcot': '#27AE60', |
| 'cot': '#8E44AD', |
| } |
|
|
| METHOD_LABELS = { |
| 'codi': 'CODI', |
| 'pmps': 'PMPS (Ours)', |
| 'simcot': 'Sim-CoT', |
| } |
|
|
|
|
| def load_results(result_dir): |
| """Load extracted metrics from a result directory.""" |
| metrics_path = os.path.join(result_dir, 'metrics.json') |
| latent_path = os.path.join(result_dir, 'latent_hiddens.npy') |
|
|
| with open(metrics_path, 'r') as f: |
| metrics = json.load(f) |
|
|
| latent_hiddens = np.load(latent_path) |
|
|
| return metrics, latent_hiddens |
|
|
|
|
| |
| |
| |
|
|
| def plot_similarity_heatmaps(metrics_dict, output_path, methods=None): |
| """Side-by-side average pairwise cosine similarity heatmaps.""" |
| if methods is None: |
| methods = list(metrics_dict.keys()) |
|
|
| n_methods = len(methods) |
| fig, axes = plt.subplots(1, n_methods, figsize=(4.5 * n_methods, 4)) |
| if n_methods == 1: |
| axes = [axes] |
|
|
| cmap = LinearSegmentedColormap.from_list( |
| 'sim', ['#FFFFFF', '#AED6F1', '#3498DB', '#1A5276'], N=256 |
| ) |
|
|
| for idx, method in enumerate(methods): |
| metrics = metrics_dict[method] |
| sim_matrices = [np.array(m['cos_sim_matrix']) for m in metrics] |
| N_l = sim_matrices[0].shape[0] |
| avg_matrix = np.mean(sim_matrices, axis=0) |
|
|
| ax = axes[idx] |
| 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.7 else 'black' |
| ax.text(j, i, f'{avg_matrix[i, j]:.2f}', |
| ha='center', va='center', fontsize=9, color=color) |
|
|
| ax.set_title(METHOD_LABELS.get(method, method), fontsize=14, fontweight='bold') |
| ax.set_xlabel('Latent Token Index') |
| ax.set_ylabel('Latent Token Index') |
| ax.set_xticks(range(N_l)) |
| ax.set_yticks(range(N_l)) |
| ax.set_xticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)]) |
| ax.set_yticklabels([f'$l_{{{i+1}}}$' for i in range(N_l)]) |
|
|
| fig.colorbar(im, ax=axes, label='Cosine Similarity', shrink=0.8, pad=0.02) |
| plt.savefig(output_path, dpi=300) |
| plt.savefig(output_path.replace('.png', '.pdf'), dpi=300) |
| plt.close() |
| print(f' Saved: {output_path}') |
|
|
|
|
| |
| |
| |
|
|
| def plot_similarity_distribution(metrics_dict, output_path, methods=None): |
| """Violin + box plot of off-diagonal cosine similarities.""" |
| if methods is None: |
| methods = list(metrics_dict.keys()) |
|
|
| fig, ax = plt.subplots(figsize=(3 + 1.5 * len(methods), 4.5)) |
|
|
| data_list = [] |
| labels = [] |
| colors = [] |
|
|
| for method in methods: |
| metrics = metrics_dict[method] |
| off_diag_sims = [] |
| for m in metrics: |
| sim_matrix = np.array(m['cos_sim_matrix']) |
| N_l = sim_matrix.shape[0] |
| mask = ~np.eye(N_l, dtype=bool) |
| off_diag_sims.extend(sim_matrix[mask].tolist()) |
| data_list.append(off_diag_sims) |
| labels.append(METHOD_LABELS.get(method, method)) |
| colors.append(COLORS.get(method, '#888888')) |
|
|
| positions = range(1, len(methods) + 1) |
| parts = ax.violinplot(data_list, positions=positions, showmedians=False, showextrema=False) |
| for pc, color in zip(parts['bodies'], colors): |
| pc.set_facecolor(color) |
| pc.set_alpha(0.3) |
|
|
| bp = ax.boxplot(data_list, positions=positions, widths=0.15, |
| patch_artist=True, showfliers=False, |
| medianprops={'color': 'black', 'linewidth': 2}) |
| for patch, color in zip(bp['boxes'], colors): |
| patch.set_facecolor(color) |
| patch.set_alpha(0.8) |
|
|
| for i, (d, color) in enumerate(zip(data_list, colors)): |
| mean_val = np.mean(d) |
| ax.plot(i + 1, mean_val, 'D', color='white', markeredgecolor=color, |
| markersize=7, markeredgewidth=1.5, zorder=5) |
| ax.annotate(f'{mean_val:.3f}', (i + 1, mean_val), |
| textcoords='offset points', xytext=(15, 5), |
| fontsize=9, color=color, fontweight='bold') |
|
|
| ax.set_xticks(positions) |
| ax.set_xticklabels(labels, fontsize=12) |
| ax.set_ylabel('Off-Diagonal Cosine Similarity', fontsize=12) |
| ax.set_title('Latent Token Pairwise Similarity\n(Lower = Less Collapse)', fontsize=13) |
| ax.axhline(y=1.0, color='gray', linestyle='--', alpha=0.4, label='Perfect Collapse') |
| ax.legend(loc='lower right', 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_singular_values(metrics_dict, output_path, methods=None): |
| """Normalized singular value spectrum comparison.""" |
| if methods is None: |
| methods = list(metrics_dict.keys()) |
|
|
| fig, axes = plt.subplots(1, 2, figsize=(11, 4.5)) |
|
|
| |
| ax = axes[0] |
| bar_width = 0.8 / len(methods) |
| for m_idx, method in enumerate(methods): |
| metrics = metrics_dict[method] |
| all_sv_norm = [m['singular_values_normalized'] for m in metrics] |
| N_l = len(all_sv_norm[0]) |
| avg_sv = np.mean(all_sv_norm, axis=0) |
| std_sv = np.std(all_sv_norm, axis=0) |
|
|
| x = np.arange(N_l) + m_idx * bar_width |
| ax.bar(x, avg_sv, width=bar_width, yerr=std_sv, capsize=3, |
| color=COLORS.get(method, '#888'), alpha=0.8, |
| label=METHOD_LABELS.get(method, method)) |
|
|
| ax.set_xlabel('Singular Value Index') |
| ax.set_ylabel('Normalized Magnitude ($\\sigma_i / \\sum \\sigma_j$)') |
| ax.set_title('(a) Singular Value Spectrum') |
| ax.set_xticks(np.arange(N_l) + bar_width * (len(methods) - 1) / 2) |
| ax.set_xticklabels([f'$\\sigma_{{{i+1}}}$' for i in range(N_l)]) |
| ax.legend() |
|
|
| |
| ax2 = axes[1] |
| for method in methods: |
| metrics = metrics_dict[method] |
| all_cum_var = [m['cumulative_variance'] for m in metrics] |
| avg_cum = np.mean(all_cum_var, axis=0) |
| std_cum = np.std(all_cum_var, axis=0) |
| x = np.arange(1, len(avg_cum) + 1) |
| ax2.plot(x, avg_cum, 'o-', color=COLORS.get(method, '#888'), |
| linewidth=2, markersize=7, label=METHOD_LABELS.get(method, method)) |
| ax2.fill_between(x, avg_cum - std_cum, np.minimum(avg_cum + std_cum, 1.0), |
| alpha=0.15, color=COLORS.get(method, '#888')) |
|
|
| ax2.axhline(y=0.95, color='gray', linestyle='--', alpha=0.5) |
| ax2.text(0.5, 0.96, '95%', color='gray', fontsize=9) |
| ax2.set_xlabel('Number of Components') |
| ax2.set_ylabel('Cumulative Explained Variance') |
| ax2.set_title('(b) Cumulative Variance') |
| ax2.set_xticks(range(1, N_l + 1)) |
| 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_effective_rank(metrics_dict, output_path, methods=None): |
| """Effective rank comparison bar chart with individual sample scatter.""" |
| if methods is None: |
| methods = list(metrics_dict.keys()) |
|
|
| fig, ax = plt.subplots(figsize=(3 + 1.2 * len(methods), 4.5)) |
|
|
| for m_idx, method in enumerate(methods): |
| metrics = metrics_dict[method] |
| ranks = [m['effective_rank'] for m in metrics] |
| mean_rank = np.mean(ranks) |
| std_rank = np.std(ranks) |
| color = COLORS.get(method, '#888') |
|
|
| ax.bar(m_idx, mean_rank, width=0.5, yerr=std_rank, capsize=5, |
| color=color, alpha=0.7, label=METHOD_LABELS.get(method, method)) |
|
|
| jitter = np.random.normal(0, 0.05, len(ranks)) |
| ax.scatter(m_idx + jitter, ranks, color=color, alpha=0.15, s=10, zorder=3) |
|
|
| ax.text(m_idx, mean_rank + std_rank + 0.1, f'{mean_rank:.2f}', |
| ha='center', fontsize=11, fontweight='bold', color=color) |
|
|
| N_l = len(metrics_dict[methods[0]][0]['cos_sim_matrix']) |
| ax.axhline(y=N_l, color='gray', linestyle='--', alpha=0.4) |
| ax.text(len(methods) - 0.5, N_l + 0.05, f'Max = {N_l}', color='gray', fontsize=9) |
| ax.axhline(y=1.0, color='red', linestyle='--', alpha=0.3) |
| ax.text(len(methods) - 0.5, 1.05, 'Full Collapse', color='red', fontsize=9) |
|
|
| ax.set_xticks(range(len(methods))) |
| ax.set_xticklabels([METHOD_LABELS.get(m, m) for m in methods], fontsize=12) |
| ax.set_ylabel('Effective Rank', fontsize=12) |
| ax.set_title('Effective Rank of Latent Embeddings\n(Higher = More Diverse)', fontsize=13) |
| ax.set_ylim(0, N_l + 1) |
|
|
| 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_tsne_visualization(hiddens_dict, output_path, methods=None, |
| max_samples=100, perplexity=30): |
| """t-SNE visualization of latent embeddings colored by latent position.""" |
| if methods is None: |
| methods = list(hiddens_dict.keys()) |
|
|
| n_methods = len(methods) |
| fig, axes = plt.subplots(1, n_methods, figsize=(5.5 * n_methods, 5)) |
| if n_methods == 1: |
| axes = [axes] |
|
|
| all_embeddings = [] |
| all_labels = [] |
| all_method_labels = [] |
|
|
| for method in methods: |
| hiddens = hiddens_dict[method][:max_samples] |
| S, N_l, D = hiddens.shape |
| flat = hiddens.reshape(-1, D) |
| positions = np.tile(np.arange(N_l), S) |
| method_ids = np.full(len(positions), methods.index(method)) |
| all_embeddings.append(flat) |
| all_labels.append(positions) |
| all_method_labels.append(method_ids) |
|
|
| combined = np.concatenate(all_embeddings, axis=0) |
| combined_labels = np.concatenate(all_labels, axis=0) |
| combined_methods = np.concatenate(all_method_labels, axis=0) |
|
|
| |
| pca = PCA(n_components=min(50, combined.shape[1])) |
| combined_pca = pca.fit_transform(combined) |
|
|
| tsne = TSNE(n_components=2, perplexity=perplexity, random_state=42, max_iter=1000) |
| combined_2d = tsne.fit_transform(combined_pca) |
|
|
| N_l = hiddens_dict[methods[0]].shape[1] |
| cmap = plt.cm.get_cmap('tab10', N_l) |
|
|
| offset = 0 |
| for m_idx, method in enumerate(methods): |
| n_points = hiddens_dict[method][:max_samples].shape[0] * N_l |
| points = combined_2d[offset:offset + n_points] |
| labels = combined_labels[offset:offset + n_points] |
| offset += n_points |
|
|
| ax = axes[m_idx] |
| for lat_idx in range(N_l): |
| mask = labels == lat_idx |
| ax.scatter(points[mask, 0], points[mask, 1], |
| c=[cmap(lat_idx)], s=15, alpha=0.5, |
| label=f'$l_{{{lat_idx+1}}}$') |
|
|
| ax.set_title(METHOD_LABELS.get(method, method), fontsize=14, fontweight='bold') |
| ax.set_xlabel('t-SNE 1') |
| ax.set_ylabel('t-SNE 2') |
| if m_idx == n_methods - 1: |
| ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left', |
| fontsize=9, markerscale=1.5) |
|
|
| 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_pca_visualization(hiddens_dict, output_path, methods=None, max_samples=100): |
| """PCA visualization — same idea as t-SNE but deterministic and faster.""" |
| if methods is None: |
| methods = list(hiddens_dict.keys()) |
|
|
| n_methods = len(methods) |
| fig, axes = plt.subplots(1, n_methods, figsize=(5.5 * n_methods, 5)) |
| if n_methods == 1: |
| axes = [axes] |
|
|
| all_embeddings = [] |
| for method in methods: |
| hiddens = hiddens_dict[method][:max_samples] |
| flat = hiddens.reshape(-1, hiddens.shape[-1]) |
| all_embeddings.append(flat) |
|
|
| combined = np.concatenate(all_embeddings, axis=0) |
| pca = PCA(n_components=2) |
| combined_2d = pca.fit_transform(combined) |
|
|
| N_l = hiddens_dict[methods[0]].shape[1] |
| cmap = plt.cm.get_cmap('tab10', N_l) |
|
|
| offset = 0 |
| for m_idx, method in enumerate(methods): |
| S = min(max_samples, hiddens_dict[method].shape[0]) |
| n_points = S * N_l |
| points = combined_2d[offset:offset + n_points] |
| offset += n_points |
|
|
| labels = np.tile(np.arange(N_l), S) |
|
|
| ax = axes[m_idx] |
| for lat_idx in range(N_l): |
| mask = labels == lat_idx |
| ax.scatter(points[mask, 0], points[mask, 1], |
| c=[cmap(lat_idx)], s=15, alpha=0.5, |
| label=f'$l_{{{lat_idx+1}}}$') |
|
|
| ax.set_title(METHOD_LABELS.get(method, method), fontsize=14, fontweight='bold') |
| ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})') |
| ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})') |
| if m_idx == n_methods - 1: |
| ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left', |
| fontsize=9, markerscale=1.5) |
|
|
| 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/figures/collapse') |
| args = parser.parse_args() |
|
|
| os.makedirs(args.output_dir, exist_ok=True) |
|
|
| methods = ['codi', 'pmps'] |
| dirs = {'codi': args.codi_dir, 'pmps': args.pmps_dir} |
| if args.simcot_dir: |
| methods.append('simcot') |
| dirs['simcot'] = args.simcot_dir |
|
|
| print('Loading results...') |
| metrics_dict = {} |
| hiddens_dict = {} |
| for method in methods: |
| m, h = load_results(dirs[method]) |
| metrics_dict[method] = m |
| hiddens_dict[method] = h |
| print(f' {method}: {len(m)} samples, latent shape {h.shape}') |
|
|
| print('\nGenerating figures...') |
|
|
| plot_similarity_heatmaps( |
| metrics_dict, |
| os.path.join(args.output_dir, 'fig_similarity_heatmap.png'), |
| methods=methods, |
| ) |
|
|
| plot_similarity_distribution( |
| metrics_dict, |
| os.path.join(args.output_dir, 'fig_similarity_distribution.png'), |
| methods=methods, |
| ) |
|
|
| plot_singular_values( |
| metrics_dict, |
| os.path.join(args.output_dir, 'fig_singular_values.png'), |
| methods=methods, |
| ) |
|
|
| plot_effective_rank( |
| metrics_dict, |
| os.path.join(args.output_dir, 'fig_effective_rank.png'), |
| methods=methods, |
| ) |
|
|
| plot_tsne_visualization( |
| hiddens_dict, |
| os.path.join(args.output_dir, 'fig_tsne.png'), |
| methods=methods, |
| ) |
|
|
| plot_pca_visualization( |
| hiddens_dict, |
| os.path.join(args.output_dir, 'fig_pca.png'), |
| methods=methods, |
| ) |
|
|
| print(f'\nAll figures saved to: {args.output_dir}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
