liqi / code /pmps /analysis /simulate_simcot_distribution.py

Upload liqi files

0ab2566 verified 1 day ago

5.8 kB

"""
Sim-CoT Information Load Simulator
Since Sim-CoT uses rigid one-to-one segmentation (evenly dividing N_c CoT
tokens into N_l segments, each mapped to one latent token), we can simulate
its information load distribution without an actual Sim-CoT checkpoint.
This script reads CoT lengths from PMPS results and computes what the
Sim-CoT distribution would look like.
Usage:
    python analysis/simulate_simcot_distribution.py \
        --pmps_dir analysis/results/pmps \
        --num_latent 6 \
        --output_dir analysis/results/simcot_simulated
"""

import os
import json
import argparse
import numpy as np


def simulate_simcot_info_load(cot_lengths, num_latent):
    """
    Simulate Sim-CoT's rigid segmentation information load.
    Sim-CoT divides the CoT into num_latent equal segments.
    Each latent token gets floor(N_c / N_l) tokens, with the last
    latent getting the remainder.
    However, for samples where N_c < N_l, some latents get 0 tokens
    (the core information bottleneck problem).
    Returns:
        info_loads: list of (num_latent,) arrays
    """
    info_loads = []
    for N_c in cot_lengths:
        load = np.zeros(num_latent)
        if N_c <= 0:
            info_loads.append(load.tolist())
            continue

        if N_c >= num_latent:
            # Normal case: divide evenly
            base_size = N_c // num_latent
            remainder = N_c % num_latent
            for i in range(num_latent):
                # First `remainder` latents get one extra token
                load[i] = base_size + (1 if i < remainder else 0)
        else:
            # N_c < N_l: some latents get nothing (surplus latents)
            for i in range(N_c):
                load[i] = 1
            # Remaining latents get 0 → unsupervised

        info_loads.append(load.tolist())

    return info_loads


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--pmps_dir', type=str, required=True,
                        help='PMPS results directory (to read CoT lengths)')
    parser.add_argument('--num_latent', type=int, default=6)
    parser.add_argument('--output_dir', type=str, default='analysis/results/simcot_simulated')
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # Read CoT lengths from PMPS results
    cot_lengths_path = os.path.join(args.pmps_dir, 'cot_lengths.json')
    if os.path.exists(cot_lengths_path):
        with open(cot_lengths_path, 'r') as f:
            cot_lengths = json.load(f)
    else:
        # Fallback: try to read from distribution metrics
        dist_path = os.path.join(args.pmps_dir, 'distribution_metrics.json')
        with open(dist_path, 'r') as f:
            dist_metrics = json.load(f)
        # Estimate CoT length from info load sums
        cot_lengths = [int(sum(m['info_load_per_latent'])) for m in dist_metrics]

    print(f'Loaded {len(cot_lengths)} CoT lengths')
    print(f'  Mean CoT length: {np.mean(cot_lengths):.1f}')
    print(f'  Min/Max: {np.min(cot_lengths)}/{np.max(cot_lengths)}')

    # Simulate Sim-CoT info loads
    info_loads = simulate_simcot_info_load(cot_lengths, args.num_latent)

    # Compute distribution metrics
    dist_metrics = []
    for i, (loads, cot_len) in enumerate(zip(info_loads, cot_lengths)):
        loads_arr = np.array(loads)
        cv = float(np.std(loads_arr) / (np.mean(loads_arr) + 1e-8))
        dist_metrics.append({
            'sample_idx': i,
            'info_load_per_latent': loads,
            'info_load_cv': cv,
            'cot_length': cot_len,
            # Sim-CoT has no prototype concept
            'proto_usage_latent': [1.0 / args.num_latent] * args.num_latent,
            'proto_entropy_latent': float(np.log(args.num_latent)),
            'proto_usage_cot': [1.0 / args.num_latent] * args.num_latent,
            'proto_entropy_cot': float(np.log(args.num_latent)),
            'max_entropy': float(np.log(args.num_latent)),
            'avg_match_entropy': 0.0,  # rigid 1-to-1 → entropy = 0
            'match_entropies': [0.0] * cot_len,
        })

    # Also create a fake metrics.json for the collapse analysis pipeline
    # (we need cos_sim_matrix and effective_rank — use zeros as placeholders
    #  unless real Sim-CoT checkpoint data is available)
    collapse_metrics = []
    for i in range(len(info_loads)):
        N_l = args.num_latent
        collapse_metrics.append({
            'sample_idx': i,
            'cos_sim_matrix': np.eye(N_l).tolist(),  # placeholder
            'avg_off_diag_sim': 0.0,
            'std_off_diag_sim': 0.0,
            'singular_values': [1.0] * N_l,
            'singular_values_normalized': [1.0 / N_l] * N_l,
            'effective_rank': float(N_l),
            'explained_variance': [1.0 / N_l] * N_l,
            'cumulative_variance': [i / N_l for i in range(1, N_l + 1)],
        })

    # Save
    with open(os.path.join(args.output_dir, 'distribution_metrics.json'), 'w') as f:
        json.dump(dist_metrics, f, indent=2)

    with open(os.path.join(args.output_dir, 'metrics.json'), 'w') as f:
        json.dump(collapse_metrics, f, indent=2)

    # Save simulated latent hiddens (zeros placeholder)
    D = 768  # placeholder hidden dim
    fake_hiddens = np.zeros((len(info_loads), args.num_latent, D))
    np.save(os.path.join(args.output_dir, 'latent_hiddens.npy'), fake_hiddens)

    # Summary
    all_cvs = [m['info_load_cv'] for m in dist_metrics]
    print(f'\nSimulated Sim-CoT distribution:')
    print(f'  Mean Info Load CV: {np.mean(all_cvs):.4f} (higher = less uniform)')
    print(f'  Samples with surplus latents (N_c < N_l): '
          f'{sum(1 for l in cot_lengths if l < args.num_latent)}')
    print(f'\nSaved to: {args.output_dir}')


if __name__ == '__main__':
    main()

	"""
	Sim-CoT Information Load Simulator

	Since Sim-CoT uses rigid one-to-one segmentation (evenly dividing N_c CoT
	tokens into N_l segments, each mapped to one latent token), we can simulate
	its information load distribution without an actual Sim-CoT checkpoint.

	This script reads CoT lengths from PMPS results and computes what the
	Sim-CoT distribution would look like.

	Usage:
	python analysis/simulate_simcot_distribution.py \
	--pmps_dir analysis/results/pmps \
	--num_latent 6 \
	--output_dir analysis/results/simcot_simulated
	"""

	import os
	import json
	import argparse
	import numpy as np


	def simulate_simcot_info_load(cot_lengths, num_latent):
	"""
	Simulate Sim-CoT's rigid segmentation information load.

	Sim-CoT divides the CoT into num_latent equal segments.
	Each latent token gets floor(N_c / N_l) tokens, with the last
	latent getting the remainder.

	However, for samples where N_c < N_l, some latents get 0 tokens
	(the core information bottleneck problem).

	Returns:
	info_loads: list of (num_latent,) arrays
	"""
	info_loads = []
	for N_c in cot_lengths:
	load = np.zeros(num_latent)
	if N_c <= 0:
	info_loads.append(load.tolist())
	continue

	if N_c >= num_latent:
	# Normal case: divide evenly
	base_size = N_c // num_latent
	remainder = N_c % num_latent
	for i in range(num_latent):
	# First `remainder` latents get one extra token
	load[i] = base_size + (1 if i < remainder else 0)
	else:
	# N_c < N_l: some latents get nothing (surplus latents)
	for i in range(N_c):
	load[i] = 1
	# Remaining latents get 0 → unsupervised

	info_loads.append(load.tolist())

	return info_loads


	def main():
	parser = argparse.ArgumentParser()
	parser.add_argument('--pmps_dir', type=str, required=True,
	help='PMPS results directory (to read CoT lengths)')
	parser.add_argument('--num_latent', type=int, default=6)
	parser.add_argument('--output_dir', type=str, default='analysis/results/simcot_simulated')
	args = parser.parse_args()

	os.makedirs(args.output_dir, exist_ok=True)

	# Read CoT lengths from PMPS results
	cot_lengths_path = os.path.join(args.pmps_dir, 'cot_lengths.json')
	if os.path.exists(cot_lengths_path):
	with open(cot_lengths_path, 'r') as f:
	cot_lengths = json.load(f)
	else:
	# Fallback: try to read from distribution metrics
	dist_path = os.path.join(args.pmps_dir, 'distribution_metrics.json')
	with open(dist_path, 'r') as f:
	dist_metrics = json.load(f)
	# Estimate CoT length from info load sums
	cot_lengths = [int(sum(m['info_load_per_latent'])) for m in dist_metrics]

	print(f'Loaded {len(cot_lengths)} CoT lengths')
	print(f' Mean CoT length: {np.mean(cot_lengths):.1f}')
	print(f' Min/Max: {np.min(cot_lengths)}/{np.max(cot_lengths)}')

	# Simulate Sim-CoT info loads
	info_loads = simulate_simcot_info_load(cot_lengths, args.num_latent)

	# Compute distribution metrics
	dist_metrics = []
	for i, (loads, cot_len) in enumerate(zip(info_loads, cot_lengths)):
	loads_arr = np.array(loads)
	cv = float(np.std(loads_arr) / (np.mean(loads_arr) + 1e-8))
	dist_metrics.append({
	'sample_idx': i,
	'info_load_per_latent': loads,
	'info_load_cv': cv,
	'cot_length': cot_len,
	# Sim-CoT has no prototype concept
	'proto_usage_latent': [1.0 / args.num_latent] * args.num_latent,
	'proto_entropy_latent': float(np.log(args.num_latent)),
	'proto_usage_cot': [1.0 / args.num_latent] * args.num_latent,
	'proto_entropy_cot': float(np.log(args.num_latent)),
	'max_entropy': float(np.log(args.num_latent)),
	'avg_match_entropy': 0.0, # rigid 1-to-1 → entropy = 0
	'match_entropies': [0.0] * cot_len,
	})

	# Also create a fake metrics.json for the collapse analysis pipeline
	# (we need cos_sim_matrix and effective_rank — use zeros as placeholders
	# unless real Sim-CoT checkpoint data is available)
	collapse_metrics = []
	for i in range(len(info_loads)):
	N_l = args.num_latent
	collapse_metrics.append({
	'sample_idx': i,
	'cos_sim_matrix': np.eye(N_l).tolist(), # placeholder
	'avg_off_diag_sim': 0.0,
	'std_off_diag_sim': 0.0,
	'singular_values': [1.0] * N_l,
	'singular_values_normalized': [1.0 / N_l] * N_l,
	'effective_rank': float(N_l),
	'explained_variance': [1.0 / N_l] * N_l,
	'cumulative_variance': [i / N_l for i in range(1, N_l + 1)],
	})

	# Save
	with open(os.path.join(args.output_dir, 'distribution_metrics.json'), 'w') as f:
	json.dump(dist_metrics, f, indent=2)

	with open(os.path.join(args.output_dir, 'metrics.json'), 'w') as f:
	json.dump(collapse_metrics, f, indent=2)

	# Save simulated latent hiddens (zeros placeholder)
	D = 768 # placeholder hidden dim
	fake_hiddens = np.zeros((len(info_loads), args.num_latent, D))
	np.save(os.path.join(args.output_dir, 'latent_hiddens.npy'), fake_hiddens)

	# Summary
	all_cvs = [m['info_load_cv'] for m in dist_metrics]
	print(f'\nSimulated Sim-CoT distribution:')
	print(f' Mean Info Load CV: {np.mean(all_cvs):.4f} (higher = less uniform)')
	print(f' Samples with surplus latents (N_c < N_l): '
	f'{sum(1 for l in cot_lengths if l < args.num_latent)}')
	print(f'\nSaved to: {args.output_dir}')


	if __name__ == '__main__':
	main()