Update README.md
Browse files---
license: apache-2.0
base_model: Qwen/Qwen2-0.5B
tags:
- experimental
- inference
- feature-selection
- bayesian
---
# 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.
## 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. |
## 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.
## 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[la