--- license: apache-2.0 base_model: Qwen/Qwen2-0.5B tags: - experimental - inference - feature-selection - bayesian - compression - lossless --- # Adaptive Sparse Feature Selection at Inference Time (Qwen2-0.5B) **Status: experimental / work in progress.** This is a research test, not a production release, not a compression method, and not a claim of a new architecture. Numbers below are placeholders — real plots and stats will be added once benchmark runs are complete. ## What this actually is During autoregressive generation, this repo tracks a small set of statistical features (mean, std, quantiles, rolling window stats, autocorrelation — 64 features per layer) computed from the **input activations** hitting each attention/FFN weight matrix in Qwen2-0.5B. A lightweight Bayesian selector then flags which of those features deviate meaningfully from their running distribution at each generation step, instead of treating every feature as equally relevant every time. The output metric is simple: **what fraction of tracked features get flagged as informative per step**, averaged over a generation. That's it. It's an exploration of whether activation statistics carry sparse, structured signal during inference — not a finished result and not a benchmark win yet. ## What this is *not* To be upfront about scope, since it's easy to over-read either script: - **Not a compression method.** There is a second script in this repo (`storage_reconstruction_test.py`) that splits a weight tensor into a scalar mean and a residual tensor stored in two separate files, then reconstructs the original by adding them back together. This is a **storage/loading split test** — reconstruction is mathematically exact by construction (`mean + (original - mean) = original`), so the R²=1.0 you'll see is expected and is not a compression result. No compression ratio is claimed anywhere in this repo. - **Not a new computation method.** The Bayesian selector changes *what gets measured and tracked* during inference, not *how the forward pass computes logits*. The underlying Qwen2-0.5B forward pass is untouched. - **Not validated against a baseline yet.** There's no side-by-side comparison here (yet) showing that the selected feature subset actually predicts anything useful about output quality, speed, or attention patterns. Right now this is instrumentation, not a proven technique. If any of that changes as testing continues, this README will be updated to reflect it — the goal is to keep the claims here matched to what's actually been measured. ## A third script: entropy-gated lossless compression Separate from the two experiments above, this repo also includes `entropy_gated_compression.py`, which applies **standard lossless compression to raw model weight bytes**, using a per-packet entropy check to decide whether compression is worth attempting. ### Measured result (single run, self-reported) ``` Model: Qwen2-0.5B, model.safetensors Original: 988,807,168 bytes (943 MB) Compressed: 793,772,032 bytes (757 MB) Space saved: 19.7% Compression ratio: 1.246x ``` This is **lossless** — decompression was verified to reproduce every tensor bit-for-bit (SHA256 match on 100% of checked tensors, see "Integrity verification" below). There is no approximation, no quantization, and no information discarded. It's exact round-trip storage in less space, not a compressed *representation* of the model in the sense pruning or quantization would give you. ### How the entropy gate works For every 8 MB packet of raw tensor bytes, the packet is treated as a stream of 256 possible byte values (0–255) and a frequency histogram is built. From that histogram: ``` p(v) = (count(v) + 1) / (total_bytes + 256) # Laplace-smoothed frequency H = - Sum p(v) * log2 p(v) # Shannon entropy, bits/byte ``` The `+1` in the numerator (and `+256` in the denominator to keep it a valid distribution) is **Laplace smoothing** — a standard fix so that byte values with zero observed occurrences don't produce `log2(0)`. This keeps every probability strictly positive. To be precise about terminology: this smoothing is a Bayesian-flavored technique (it's equivalent to placing a uniform Dirichlet prior over the byte distribution and taking a posterior mean), but the compression *decision* itself is a simple **entropy threshold**, not a Bayesian inference procedure over model parameters. If `H >= 7.90` bits/byte, the packet is treated as high-entropy (close to indistinguishable from random noise) and stored raw — attempting `zlib` on it would cost CPU time for little or no space savings. If `H < 7.90`, the packet goes through `zlib` at compression level 1, and the result is kept only if it's actually smaller than the raw packet. ### Why weight tensors compress at all Raw `float16`/`bfloat16` weights are not uniformly random at the byte level, even though they look noise-like in aggregate. Mantissa bits tend to be close to high-entropy, but: - Normalization layer weights (`input_layernorm`, `post_attention_layernorm`) and biases occupy a narrower value range than general weight matrices, which tends to concentrate their byte-value histograms and lower measured entropy. - Repeated or near-repeated byte patterns across a large tensor give `zlib`'s LZ77-style matching something to exploit, independent of the entropy gate. A 19.7% reduction is consistent with the range generally reported for applying generic lossless compressors (gzip/zlib-class) directly to raw transformer weight bytes — typically single-digit to ~25%, varying by layer type and model. This is **not a novel compression ratio** by research standards; it's roughly what's expected from entropy-aware lossless compression on this kind of data. The contribution here, if there is one, is the packet-level entropy gate deciding *where* to spend compression effort rather than applying `zlib` uniformly — this mainly saves CPU time on incompressible packets rather than improving the ratio itself. ### Integrity verification Every packet is SHA256-hashed before compression and checked again after decompression. The proof-of-correctness run showed **100.0000% match across all checked tensors** — every layer's weight, bias, and normalization tensor was recovered bit-identical to the source file. This is the expected outcome for a correctly implemented lossless codec with a working verification step, not a measured quality metric — it confirms the round-trip is correct, the same way a working checksum is expected to match rather than being a research result in itself. ## Why this might be interesting anyway Most work on transformer internals looks at weights (pruning, quantization, low-rank decomposition). This script instead asks: at inference time, does the *activation* stream flowing through each layer have a small, identifiable subset of statistics that matter more than the rest at any given step? If that subset is small and stable, it's a hint (not proof) that there's structure worth digging into — for interpretability, for adaptive compute, or just as a diagnostic tool for understanding what a layer is "paying attention to" numerically. That's the honest pitch. No claims beyond it yet. ## Files | File | What it does | |---|---| | `terminal_chat_bayesian.py` | Main experiment. Loads Qwen2-0.5B, hooks every attention/FFN weight's input activations, runs the Bayesian feature selector during generation, prints the fraction of flagged features per response. Requires `bayes_analysis.safetensors` (see below). | | `storage_reconstruction_test.py` | Secondary test. Splits weight tensors into `(mean_scalar, residual_tensor)` across a JSON + safetensors file, reconstructs on load. Included for transparency — this is a loading mechanics test, not a result. | | `entropy_gated_compression.py` | Lossless compression of raw safetensors weight bytes using an entropy-gated packet scheme. Measured 19.7% size reduction on Qwen2-0.5B with verified bit-exact decompression. See "A third script" section above for details. | ## Requirements ```bash pip install torch transformers safetensors numpy ``` CUDA GPU required for `terminal_chat_bayesian.py` (checks `torch.cuda.is_available()` and will exit if not found). `storage_reconstruction_test.py` runs on CPU. ## How to run ### 1. Bayesian feature selector chat (main experiment) You need a `bayes_analysis.safetensors` file in the working directory containing precomputed per-layer feature tensors (keys ending in `__feat`). This file is produced by a separate analysis pass over the model's weights — generate it before running this script, or use the one provided in this repo's Files tab if included. ```bash python terminal_chat_bayesian.py ``` In the chat session: - Type normally to talk to the model - `/stats` — shows how many features were flagged vs. total possible in the last response - `/bayes` — shows the top 10 layers by number of currently-flagged features - `/clear` — resets conversation history - `/exit` — quit ### 2. Storage/reconstruction test (secondary, not a compression result) Requires `bayesian_features.json` and `layer_residuals.safetensors` in `/content/` (paths are hardcoded for Colab — edit `json_path` / `safetensors_path` in `prepare_fast_hybrid_model()` if running elsewhere). ```bash python storage_reconstruction_test.py ``` This will strip attention/FFN weights from the loaded model and reconstruct them from the two files, then start a basic chat loop. Reconstruction is exact by construction — see the "What this is not" section above for why. ### 3. Entropy-gated lossless compression ```bash python entropy_gated_compression.py ``` By default this compresses `/content/Qwen2-0.5B` (edit `source_model_dir` at the bottom of the script for a different path) into a bundle directory, then immediately decompresses it back to verify correctness. Output: ``` Compression finished. Payload: qwen2_0_5b_bayes_zlib_bundle/payload.bin Manifest: qwen2_0_5b_bayes_zlib_bundle/manifest.json Packet: 8.0 MB Restored to: qwen2_0_5b_restored ``` To use the compression/decompression functions directly instead of running the full round-trip: ```python from entropy_gated_compression import compress_qwen2_safetensors_fast, decompress_qwen2_safetensors_fast compress_qwen2_safetensors_fast( model_dir="/path/to/Qwen2-0.5B", output_bundle_dir="my_bundle", packet_mb=8, ) decompress_qwen2_safetensors_fast( bundle_dir="my_bundle", restored_model_dir="my_restored_model", ) ``` ## Code ### `terminal_chat_bayesian.py` ```python import torch import numpy as np from safetensors.torch import load_file from transformers import AutoTokenizer, AutoModelForCausalLM import time import os import sys MODEL_NAME = "Qwen/Qwen2-0.5B" MAX_NEW_TOKENS = 200 TEMPERATURE = 0.7 ANALYSIS_FILE = "bayes_analysis.safetensors" SYSTEM_PROMPT = "You are a helpful assistant." NUM_FEATURES = 64 BAYES_EVERY_N = 8 # compute bayes stats every N tokens instead of every token BAYES_ENABLED = True # can be fully disabled with this flag def _row_features_torch(x: torch.Tensor, n_features: int = NUM_FEATURES) -> torch.Tensor: x = x.float() L = x.shape[0] mean = x.mean() std = x.std(unbiased=False) abs_x = x.abs() feats = torch.zeros(n_features, dtype=torch.float32, device=x.device) feats[0] = mean feats[1] = std feats[2] = x.max() feats[3] = x.min() q = torch.quantile(x, torch.tensor([0.25, 0.5, 0.75, 0.05, 0.10, 0.90, 0.95], device=x.device)) feats[4], feats[5], feats[6] = q[0], q[1], q[2] feats[16], feats[17], feats[18], feats[19] = q[3], q[4], q[5], q[6] feats[7] = (x > mean + std).sum() feats[8] = (x < mean - std).sum() feats[9] = abs_x.mean() feats[10] = abs_x.median() w = 8 if L >= w: wins = x.unfold(0, w, 1) feats[11] = wins.mean(dim=1).mean() feats[12] = wins.std(dim=1, unbiased=False).mean() feats[13] = wins.max(dim=1).values.mean() feats[14] = wins.min(dim=1).values.mean() feats[15] = x.diff().abs().mean() else: feats[11], feats[12], feats[13], feats[14], feats[15] = mean, std, x.max(), x.min(), 0.0 if L > 1 and std > 1e-12: a, b = x[:-1], x[1:] a_c, b_c = a - a.mean(), b - b.mean() denom = torch.sqrt((a_c * a_c).sum() * (b_c * b_c).sum()) feats[20] = (a_c * b_c).sum() / denom if denom > 1e-12 else 0.0 else: feats[20] = 0.0 return feats[:n_features] class BayesData: def __init__(self, path: str = ANALYSIS_FILE): if not os.path.exists(path): print(f"[error] {path} not found. Run the analysis pass first to generate it.") sys.exit(1) print("[bayes-data] loading from safetensors ...") raw = load_file(path) self.layers = {} names = {k[: -len("__feat")] for k in raw.keys() if k.endswith("__feat")} for sk in names: param_name = sk.replace("__", ".") self.layers[param_name] = {"feat": raw[f"{sk}__feat"].float().numpy()} print(f"[bayes-data] loaded {len(self.layers)} layers") def get(self, param_name): return self.layers.get(param_name) def num_features_for(self, param_name) -> int: data = self.layers.get(param_name) return 1 if data is None else data["feat"].shape[1] class BayesianFeatureSelector: def __init__(self, n_features: int, device): self.n_features = n_features self.marked_counts = torch.ones(n_features, dtype=torch.float32, device=device) self.unmarked_counts = torch.ones(n_features, dtype=torch.float32, device=device) self.running_mean = torch.zeros(n_features, dtype=torch.float32, device=device) self.running_var = torch.ones(n_features, dtype=torch.float32, device=device) self.n_seen = 0 def select(self, feat_vector: torch.Tensor) -> torch.Tensor: if self.n_seen == 0: return torch.arange(self.n_features, device=feat_vector.device) std = torch.sqrt(self.running_var) + 1e-8 deviation = (feat_vector - self.running_mean).abs() / std marked = torch.where(deviation > 1.0)[0] if marked.numel() == 0: priors = self.marked_counts / (self.marked_counts + self.unmarked_counts) marked = priors.argmax().unsqueeze(0) return marked def update(self, feat_vector: torch.Tensor, marked_idx: torch.Tensor): marked_mask = torch.zeros(self.n_features, dtype=torch.bool, device=feat_vector.device) marked_mask[marked_idx] = True self.marked_counts[marked_mask] += 1 self.unmarked_counts[~marked_mask] += 1 self.n_seen += 1 delta = feat_vector - self.running_mean self.running_mean += delta / self.n_seen delta2 = feat_vector - self.running_mean self.running_var += (delta * delta2 - self.running_var) / self.n_seen self.running_var.clamp_(min=1e-8) class LayerBayesRegistry: def __init__(self, layer_names: list, n_features: int, device): self.selectors = {name: BayesianFeatureSelector(n_features, device) for name in layer_names} self.layer_order = layer_names self.n_features = n_features def select_for(self, layer_name: str, feat_vector: torch.Tensor) -> torch.Tensor: return self.selectors[layer_name].select(feat_vector) def observe(self, layer_name: str, feat_vector: torch.Tensor, marked_idx: torch.Tensor): self.selectors[layer_name].update(feat_vector, marked_idx) def state_summary(self) -> dict: out = {} for name, sel in self.selectors.items(): if sel.n_seen == 0: out[name] = sel.n_features else: priors = sel.marked_counts / (sel.marked_counts + sel.unmarked_counts) out[name] = int((priors > 0.5).sum().item()) return out def build_bayes_registry(model, bayes_data, device) -> LayerBayesRegistry: layer_names = [] n_features = NUM_FEATURES for name, module in model.named_modules(): param_name = f"{name}.weight" if bayes_data.get(param_name) is not None and hasattr(module, "weight"): layer_names.append(param_name) n_features = bayes_data.num_features_for(param_name) print(f"[registry] {len(layer_names)} layers, n_features={n_features}") return LayerBayesRegistry(layer_names, n_features, device) def generate_with_bayes_scalar(model, tokenizer, history, bayes_data, bayes_registry): try: prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True) except Exception: prompt = "\n".join(f"{m['role'].upper()}: {m['content']}" for m in history) + "\nASSISTANT:" inputs = tokenizer(prompt, return_tensors="pt").to(model.device) input_ids = inputs["input_ids"] total_scalars_used = 0 total_elements_all = 0 activation_store = {} hooks = [] layer_meta = {} if BAYES_ENABLED: for name, module in model.named_modules(): param_name = f"{name}.weight" data = bayes_data.get(param_name) if data is None or not hasattr(module, "weight"): continue w = module.weight dim_in = w.shape[1] if w.ndim >= 2 else w.shape[0] n_features = data["feat"].shape[1] layer_meta[param_name] = (dim_in, n_features) def make_hook(pn, di): def hook_fn(module, inp, out): x_in = inp[0] if x_in.ndim == 3: x_t = x_in[0, -1, :] elif x_in.ndim == 2: x_t = x_in[0, :] else: return if x_t.shape[0] == di: activation_store[pn] = x_t.detach() return hook_fn hooks.append(module.register_forward_hook(make_hook(param_name, dim_in))) vocab_size = tokenizer.vocab_size or model.config.vocab_size new_tokens_list = [] step_counter = 0 with torch.no_grad(): past_key_values = None cur_input = input_ids for step in range(MAX_NEW_TOKENS): if step == 0: out = model(input_ids=cur_input, use_cache=True) else: out = model(input_ids=cur_input, past_key_values=past_key_values, use_cache=True) past_key_values = out.past_key_values logits = out.logits[:, -1, :vocab_size].float() torch.nan_to_num_(logits, nan=0.0, posinf=1e4, neginf=-1e4) logits.div_(max(TEMPERATURE, 1e-6)) sorted_logits, sorted_idx = torch.sort(logits, descending=True) probs_sorted = torch.softmax(sorted_logits, dim=-1) cumprobs = torch.cumsum(probs_sorted, dim=-1) mask = (cumprobs - probs_sorted) > 0.9 sorted_logits[mask] = -1e9 probs = torch.softmax(sorted_logits, dim=-1) probs.clamp_(min=0.0) s = probs.sum(dim=-1, keepdim=True) if not (s == 0).any(): probs.div_(s) else: probs.fill_(1.0 / probs.shape[-1]) next_sorted = torch.multinomial(probs, num_samples=1) next_token = sorted_idx.gather(-1, next_sorted) next_id = next_token.item() new_tokens_list.append(next_id) if BAYES_ENABLED and (step_counter % BAYES_EVERY_N == 0) and activation_store: for param_name, x_t in activation_store.items(): dim_in, n_features = layer_meta[param_name] feat_vector = _row_features_torch(x_t, n_features) marked_idx = bayes_registry.select_for(param_name, feat_vector) total_scalars_used += marked_idx.numel() total_elements_all += n_features bayes_registry.observe(param_name, feat_vector, marked_idx) activation_store.clear() step_counter += 1 if next_id == tokenizer.eos_token_id: break cur_input = next_token for h in hooks: h.remove() response_text = tokenizer.decode(new_tokens_list, skip_special_tokens=True) pct = 100.0 * total_scalars_used / total_elements_all if total_elements_all > 0 else 0.0 return response_text, total_scalars_used, total_elements_all, pct BANNER = """ +==================================================================+ | Qwen2-0.5B x Bayesian Minimal Feature Selection | | /stats - stats for the last response | | /bayes - state of the bayesian models (top 10 by k) | | /clear - clear history | | /exit - quit | +==================================================================+ """ def chat(model, tokenizer, bayes_data, bayes_registry): print(BANNER) history = [{"role": "system", "content": SYSTEM_PROMPT}] last_stats = None while True: try: user = input("You: ").strip() except (EOFError, KeyboardInterrupt): print("\nExiting.") break if not user: continue if user == "/exit": break if user == "/clear": history = [{"role": "system", "content": SYSTEM_PROMPT}] print("[history cleared]") continue if user == "/stats": if last_stats: sc, el, pct = last_stats print(f"\n Scalars flagged : {sc:,}") print(f" Total possible : {el:,}") print(f" Fraction flagged : {pct:.4f}%\n") else: print("[no data yet - send a message first]") continue if user == "/bayes": summary = bayes_registry.state_summary() print("\n [bayesian state - top 10 layers by k]") for name, k in sorted(summary.items(), key=lambda x: -x[1])[:10]: print(f" {name:<55} k={k}") print() continue history.append({"role": "user", "content": user}) t0 = time.time() resp, scalars_used, total_elements, pct = generate_with_bayes_scalar( model, tokenizer, history, bayes_data, bayes_registry ) history.append({"role": "assistant", "content": resp}) elapsed = time.time() - t0 last_stats = (scalars_used, total_elements, pct) print(f"\nModel ({elapsed:.1f}s): {resp}") print(f"\n +- Bayesian minimal feature selection -----------------+") print(f" | Flagged : {scalars_used:>15,} |") print(f" | Total : {total_elements:>15,} |") print(f" | Fraction : {pct:>14.4f} % |") print(f" +--------------------------------------------------------+\n") if __name__ == "__main__": if not torch.cuda.is_available(): print("[error] CUDA not available. This script is configured for GPU.") sys.exit(1) device = "cuda" print(f"[start] device: {device}") torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.benchmark = True print(f"\n[1/3] Loading {MODEL_NAME} ...") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, torch_dtype=torch.bfloat16, device_map=device, trust_remote_code=True, ) model.eval() print("\n[2/3] Loading features from analysis file ...") bayes_data = BayesData() print("\n[3/3] Initializing bayesian feature selection registry ...") bayes_registry = build_bayes_registry(model, bayes_data, device) chat(model, tokenizer, bayes_data, bayes_registry) ``` ### `storage_reconstruction_test.py` ```python import torch import numpy as np from transformers import AutoModelForCausalLM, AutoTokenizer import time import json import os from safetensors.torch import load_file # ========================================== # 1. STORAGE / LOADING TEST (NOT A COMPRESSION RESULT) # ========================================== # NOTE: reconstruction below is mean + residual, which is mathematically # exact by construction (mean + (original - mean) = original). # R2 = 1.0 is expected here and does not indicate compression - # it indicates the two files together contain the same information # as the original weight, just split across two files. class FastBayesianStorage: """Weight storage split across two files, for testing a load pipeline""" def __init__(self): self.base_predictions = {} self.layer_residuals = {} self.layer_shapes = {} def decompress_layer(self, name): """Exact reconstruction: mean_val + residual = original (by construction)""" shape = self.layer_shapes[name] mean_val = self.base_predictions[name] residual = self.layer_residuals[name] reconstructed = np.full(residual.shape, mean_val, dtype=np.float32) + residual return torch.from_numpy(reconstructed).view(shape) def load_from_files(self, json_path="/content/bayesian_features.json", safetensors_path="/content/layer_residuals.safetensors"): """Loads scalar features from JSON and residual tensors from Safetensors""" print(f"\n[Import] Loading features and layer residuals from files...") # 1. Load metadata and scalar features with open(json_path, "r", encoding="utf-8") as f: json_data = json.load(f) self.base_predictions = json_data["base_predictions"] self.layer_shapes = json_data["layer_shapes"] print(f" -> Scalar features and shapes loaded from: {json_path}") # 2. Load residual tensors (convert Torch -> NumPy for reconstruction) tensors_dict = load_file(safetensors_path) for name, tensor in tensors_dict.items(): self.layer_residuals[name] = tensor.numpy() print(f" -> Residual tensors loaded from: {safetensors_path}") # ========================================== # 2. BUILD HYBRID MODEL FROM SPLIT FILES # ========================================== def prepare_fast_hybrid_model(model_name="Qwen/Qwen2-0.5B"): start_time = time.time() print(f"Loading base model and tokenizer {model_name}...") tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.float32, device_map="cpu", low_cpu_mem_usage=True ) target_layers = ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'] compressed_layer_names = [] print("\n[Process] Removing original attention/FFN weight tensors (preparing to load from files)...") for name, param in list(model.named_parameters()): if any(target in name for target in target_layers) and "weight" in name: compressed_layer_names.append(name) # remove original weight to simulate a clean storage state delattr(model.get_submodule(name.rsplit('.', 1)[0]), 'weight') print(f"\n[Done] Structure preparation time: {time.time() - start_time:.2f} sec.") return model, tokenizer, compressed_layer_names # ========================================== # 3. TERMINAL CHAT # ========================================== def run_fast_terminal_chat(): # Paths to your prepared files json_path = "/content/bayesian_features.json" safetensors_path = "/content/layer_residuals.safetensors" # Build empty model structure model, tokenizer, compressed_names = prepare_fast_hybrid_model() # Initialize storage and load the already-prepared files (no overwrite) storage = FastBayesianStorage() storage.load_from_files(json_path=json_path, safetensors_path=safetensors_path) # Reconstruct weights from loaded files start_restore = time.time() print("\n[Info] Reconstructing weight tensors from loaded files...") for name in compressed_names: restored_tensor = storage.decompress_layer(name) submodule = model.get_submodule(name.rsplit('.', 1)[0]) submodule.weight = torch.nn.Parameter(restored_tensor) print(f"[Done] All weights reconstructed (R2=1.0 by construction, see note above) in: {time.time() - start_restore:.2f} sec!") print("\n" + "="*50) print(" QWEN-0.5B CHAT - RECONSTRUCTED FROM SPLIT FILES") print(" Type 'exit' to quit.") print("="*50 + "\n") while True: user_input = input("You: ") if user_input.lower() in ['exit', 'quit']: break if not user_input.strip(): continue messages = [{"role": "user", "content": user_input}] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) model_inputs = tokenizer([text], return_tensors="pt") print("Qwen: ", end="", flush=True) generated_ids = model_inputs.input_ids with torch.no_grad(): for _ in range(70): outputs = model(input_ids=generated_ids) next_token_logits = outputs.logits[:, -1, :] next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) token_str = tokenizer.decode(next_token[0], skip_special_tokens=True) print(token_str, end="", flush=True) generated_ids = torch.cat([generated_ids, next_token], dim=-1) if next_token.item() in [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>")]: break print("\n" + "-"*50) if __name__ == "__main__": run_fast_terminal_chat() ``` ### `entropy_gated_compression.py` ```python from __future__ import annotations import hashlib import json import os import struct import zlib from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path from typing import Dict, List, Sequence, Tuple import torch from safetensors import safe_open from safetensors.torch import save_file MAGIC = "__BAYES_PACKET_ZLIB__" VERSION = 4 DEFAULT_PACKET_MB = 8 RAW_ENTROPY_THRESHOLD = 7.90 MIN_COMPRESS_BYTES = 256 * 1024 ZLIB_LEVEL = 1 def _configure_torch() -> None: try: torch.set_num_threads(1) except Exception: pass try: torch.set_num_interop_threads(1) except Exception: pass def dtype_to_name(dtype: torch.dtype) -> str: return str(dtype).replace("torch.", "") def name_to_dtype(name: str) -> torch.dtype: return getattr(torch, name) def tensor_to_raw_bytes(t: torch.Tensor) -> bytes: t = t.detach().contiguous().cpu() return t.view(torch.uint8).numpy().tobytes() def raw_bytes_to_tensor(raw: bytes, dtype: torch.dtype, shape: Sequence[int]) -> torch.Tensor: if not raw: return torch.empty(tuple(shape), dtype=dtype) u8 = torch.frombuffer(memoryview(raw), dtype=torch.uint8).clone() return u8.view(dtype).reshape(tuple(shape)).contiguous() def _sha256(data: bytes) -> str: return hashlib.sha256(data).hexdigest() def _packetize(raw: bytes, packet_size: int) -> List[bytes]: if packet_size <= 0: raise ValueError("packet_size must be positive") if not raw: return [b""] return [raw[i:i + packet_size] for i in range(0, len(raw), packet_size)] def _bayes_features(raw: bytes) -> Dict[str, float]: # NOTE: named for the Laplace-smoothed histogram below, which is a # Bayesian-flavored smoothing technique (equivalent to a uniform # Dirichlet prior). The compression decision itself is a plain # entropy threshold, not Bayesian inference over model parameters. if not raw: return { "n": 0, "mean": 0.0, "std": 0.0, "min": 0, "max": 0, "nonzero": 0, "entropy": 0.0, "top1_mass": 0.0, "hist_sha256": _sha256(b""), } u8 = torch.frombuffer(memoryview(raw), dtype=torch.uint8) n = int(u8.numel()) counts = torch.bincount(u8.to(torch.int64), minlength=256).to(torch.float32) posterior = counts + 1.0 # Laplace smoothing (uniform Dirichlet prior, alpha=1) total = float(posterior.sum().item()) probs = posterior / total entropy = float((-(probs * torch.log2(probs.clamp_min(1e-12)))).sum().item()) top1_mass = float((posterior.max() / total).item()) f = u8.float() mean = float(f.mean().item()) std = float(f.std(unbiased=False).item()) if n > 1 else 0.0 mn = int(u8.min().item()) mx = int(u8.max().item()) nonzero = int((u8 != 0).sum().item()) hist_sha256 = _sha256(counts.to(torch.int32).cpu().numpy().tobytes()) return { "n": n, "mean": mean, "std": std, "min": mn, "max": mx, "nonzero": nonzero, "entropy": entropy, "top1_mass": top1_mass, "hist_sha256": hist_sha256, } def _choose_codec(raw: bytes) -> Tuple[str, bytes, Dict[str, float]]: feats = _bayes_features(raw) if len(raw) < MIN_COMPRESS_BYTES or feats["entropy"] >= RAW_ENTROPY_THRESHOLD: return "raw", raw, feats comp = zlib.compress(raw, level=ZLIB_LEVEL) if len(comp) >= len(raw): return "raw", raw, feats return "zlib", comp, feats def _write_record(out, meta: Dict, payload: bytes) -> None: meta_bytes = json.dumps(meta, ensure_ascii=False, separators=(",", ":")).encode("utf-8") out.write(struct.pack(">I", len(meta_bytes))) out.write(meta_bytes) out.write(struct.pack(">I", len(payload))) out.write(payload) def _read_exact(f, n: int) -> bytes: data = f.read(n) if len(data) != n: raise EOFError("Unexpected end of payload") return data def _read_record(f): head = f.read(4) if not head: return None, None if len(head) != 4: raise EOFError("Corrupted record header") meta_len = struct.unpack(">I", head)[0] meta = json.loads(_read_exact(f, meta_len).decode("utf-8")) payload_len = struct.unpack(">I", _read_exact(f, 4))[0] payload = _read_exact(f, payload_len) return meta, payload def _compress_shard_worker(args): shard_path, tensor_names, packet_size = args shard_name = Path(shard_path).name entries = [] with safe_open(str(shard_path), framework="pt", device="cpu") as f: try: shard_metadata = f.metadata() except Exception: shard_metadata = None for tensor_name in tensor_names: tensor = f.get_tensor(tensor_name) raw = tensor_to_raw_bytes(tensor) packets = _packetize(raw, packet_size) tensor_entry = { "name": tensor_name, "dtype": dtype_to_name(tensor.dtype), "shape": list(tensor.shape), "raw_len": len(raw), "packet_count": len(packets), } for packet_index, packet_raw in enumerate(packets): codec, payload, feats = _choose_codec(packet_raw) packet_meta = { "kind": "packet", "shard_name": shard_name, "tensor_name": tensor_name, "dtype": dtype_to_name(tensor.dtype), "shape": list(tensor.shape), "codec": codec, "packet_index": packet_index, "packet_count": len(packets), "packet_raw_len": len(packet_raw), "sha256": _sha256(packet_raw), "features": feats, "is_last_packet": packet_index == len(packets) - 1, "payload_len": len(payload), } entries.append((packet_meta, payload)) return shard_name, shard_metadata, entries def compress_qwen2_safetensors_fast( model_dir: str, output_bundle_dir: str, packet_mb: int = DEFAULT_PACKET_MB, max_workers: int | None = None, ) -> None: _configure_torch() model_dir = Path(model_dir) out_dir = Path(output_bundle_dir) out_dir.mkdir(parents=True, exist_ok=True) shard_files = sorted(model_dir.glob("*.safetensors")) if not shard_files: raise FileNotFoundError(f"No .safetensors files found in {model_dir}") packet_size = max(64 * 1024, packet_mb * 1024 * 1024) cpu_count = os.cpu_count() or 1 if max_workers is None: max_workers = max(1, min(cpu_count, len(shard_files), 8)) manifest = { "format": MAGIC, "version": VERSION, "source_model_dir": str(model_dir), "packet_size": packet_size, "compression": "zlib", "zlib_level": ZLIB_LEVEL, "files": [], } for aux_name in [ "config.json", "generation_config.json", "tokenizer_config.json", "special_tokens_map.json", "model.safetensors.index.json", ]: aux_path = model_dir / aux_name if aux_path.exists() and aux_path.is_file(): manifest.setdefault("aux_files", []) manifest["aux_files"].append( {"name": aux_name, "text": aux_path.read_text(encoding="utf-8")} ) jobs = [] for shard_path in shard_files: with safe_open(str(shard_path), framework="pt", device="cpu") as f: tensor_names = list(f.keys()) jobs.append((str(shard_path), tensor_names, packet_size)) payload_path = out_dir / "payload.bin" if len(jobs) == 1: results = [_compress_shard_worker(jobs[0])] else: results = [None] * len(jobs) with ProcessPoolExecutor(max_workers=max_workers) as pool: future_map = {pool.submit(_compress_shard_worker, job): i for i, job in enumerate(jobs)} for fut in as_completed(future_map): idx = future_map[fut] results[idx] = fut.result() with open(payload_path, "wb") as payload_out: for shard_name, shard_meta, entries in results: shard_entry = { "name": shard_name, "metadata": shard_meta, "records": len(entries), "tensors": [], } tensor_map = {} for packet_meta, payload in entries: _write_record(payload_out, packet_meta, payload) tname = packet_meta["tensor_name"] if tname not in tensor_map: tensor_map[tname] = { "name": tname, "dtype": packet_meta["dtype"], "shape": packet_meta["shape"], "raw_len": 0, "packet_count": packet_meta["packet_count"], } tensor_map[tname]["raw_len"] = packet_meta["packet_raw_len"] shard_entry["tensors"] = list(tensor_map.values()) manifest["files"].append(shard_entry) with open(out_dir / "manifest.json", "w", encoding="utf-8") as f: json.dump(manifest, f, ensure_ascii=False, separators=(",", ":")) print("Compression finished.") print(f"Payload: {payload_path}") print(f"Manifest: {out_dir / 'manifest.json'}") print(f"Packet: {packet_size / (1024 * 1024):.1f} MB") def _verify_packet(meta: Dict, raw: bytes) -> None: if len(raw) != int(meta["packet_raw_len"]): raise ValueError( f"Length mismatch for {meta.get('tensor_name')} packet {meta.get('packet_index')}" ) if _sha256(raw) != meta["sha256"]: raise ValueError( f"SHA256 mismatch for {meta.get('tensor_name')} packet {meta.get('packet_index')}" ) feats = _bayes_features(raw) exp = meta["features"] if feats["hist_sha256"] != exp["hist_sha256"]: raise ValueError( f"Histogram signature mismatch for {meta.get('tensor_name')} packet {meta.get('packet_index')}" ) if feats["n"] != exp["n"]: raise ValueError( f"Feature length mismatch for {meta.get('tensor_name')} packet {meta.get('packet_index')}" ) if int(feats["min"]) != int(exp["min"]) or int(feats["max"]) != int(exp["max"]): raise ValueError( f"Range feature mismatch for {meta.get('tensor_name')} packet {meta.get('packet_index')}" ) def decompress_qwen2_safetensors_fast( bundle_dir: str, restored_model_dir: str, ) -> None: _configure_torch() bundle_dir = Path(bundle_dir) restored_model_dir = Path(restored_model_dir) restored_model_dir.mkdir(parents=True, exist_ok=True) manifest_path = bundle_dir / "manifest.json" payload_path = bundle_dir / "payload.bin" if not manifest_path.exists(): raise FileNotFoundError(f"Missing manifest.json: {manifest_path}") if not payload_path.exists(): raise FileNotFoundError(f"Missing payload.bin: {payload_path}") with open(manifest_path, "r", encoding="utf-8") as f: manifest = json.load(f) for aux in manifest.get("aux_files", []): (restored_model_dir / aux["name"]).write_text(aux["text"], encoding="utf-8") shard_meta_map = {entry["name"]: entry.get("metadata") for entry in manifest["files"]} current_shard_name = None current_state_dict = {} current_tensor_parts: Dict[str, List[bytes]] = {} current_tensor_meta: Dict[str, Dict] = {} def flush_current_shard(): nonlocal current_state_dict, current_tensor_parts, current_tensor_meta, current_shard_name if current_shard_name is None: return for tensor_name, parts in current_tensor_parts.items(): meta = current_tensor_meta[tensor_name] raw = b"".join(parts) dtype = name_to_dtype(meta["dtype"]) shape = tuple(meta["shape"]) current_state_dict[tensor_name] = raw_bytes_to_tensor(raw, dtype=dtype, shape=shape) out_shard = restored_model_dir / current_shard_name save_file(current_state_dict, str(out_shard), metadata=shard_meta_map.get(current_shard_name)) current_state_dict = {} current_tensor_parts = {} current_tensor_meta = {} current_shard_name = None with open(payload_path, "rb") as payload_in: while True: meta, payload = _read_record(payload_in) if meta is None: break codec = meta["codec"] if codec == "zlib": raw = zlib.decompress(payload) elif codec == "raw": raw = payload else: raise ValueError(f"Unknown codec: {codec}") _verify_packet(meta, raw) shard_name = meta["shard_name"] tensor_name = meta["tensor_name"] if current_shard_name is None: current_shard_name = shard_name elif current_shard_name != shard_name: flush_current_shard() current_shard_name = shard_name if tensor_name not in current_tensor_parts: current_tensor_parts[tensor_name] = [] current_tensor_meta[tensor_name] = meta current_tensor_parts[tensor_name].append(raw) flush_current_shard() print(f"Restored to: {restored_model_dir}") if __name__ == "__main__": source_model_dir = "/content/Qwen2-0.5B" bundle_dir = "qwen2_0_5b_bayes_zlib_bundle" restored_dir = "qwen2_0_5b_restored" compress_qwen2_safetensors_fast( source_model_dir, bundle_dir, packet_mb=8, ) decompress_qwen2_safetensors_fast(bundle_dir, restored_dir) ``` ## Results **Entropy-gated compression — measured, single run:** 19.7% size reduction (943 MB → 757 MB) on Qwen2-0.5B, with 100% bit-exact decompression verified via SHA256. See "A third script" section above for the full breakdown. This has been run once — repeating across other model sizes and dtypes (fp16 vs bf16 vs fp32) would show whether 19.7% generalizes or is specific to this checkpoint. **Bayesian feature selection during inference — still placeholder.** The activation-tracking experiment (main script) has not yet been benchmarked systematically: - [ ] Fraction of features flagged per layer, averaged across a test set of prompts - [ ] How the flagged fraction changes over the course of a generation (early tokens vs. late tokens) - [ ] Per-layer comparison: which layers have consistently high vs. low flagged fractions - [ ] Any correlation (or lack of one) between flagged fraction and output quality — this is the test that would actually justify calling the flagged subset "informative" ## Open questions / next steps - Does the flagged feature subset stay stable across different prompts, or does it change drastically session to session? - Is there a relationship between which features get flagged and attention patterns in the same layer? - Right now `BAYES_EVERY_N = 8` and the deviation threshold (`> 1.0` std) are picked without tuning — sweeping these would show whether the flagged fraction is a real signal or just a threshold artifact. ## License Apache 2.0, matching the base model license. This repo builds on [Qwen/Qwen2-0.5B](https://huggingface.co/Qwen/Qwen2-0.5B).