arcticoneai commited on
Commit
4ecde8e
·
verified ·
1 Parent(s): f93e0b4

Upload README.md

Browse files
Files changed (1) hide show
  1. README.md +601 -4
README.md CHANGED
@@ -1,5 +1,602 @@
1
  ---
2
- license: other
3
- license_name: custom-source-available
4
- license_link: LICENSE
5
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ license: apache-2.0
3
+ base_model: Qwen/Qwen2-0.5B
4
+ tags:
5
+ - experimental
6
+ - inference
7
+ - feature-selection
8
+ - bayesian
9
+ ---
10
+
11
+ # Adaptive Sparse Feature Selection at Inference Time (Qwen2-0.5B)
12
+
13
+ **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.
14
+
15
+ ## What this actually is
16
+
17
+ 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.
18
+
19
+ 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.
20
+
21
+ ## What this is *not*
22
+
23
+ To be upfront about scope, since it's easy to over-read either script:
24
+
25
+ - **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.
26
+ - **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.
27
+ - **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.
28
+
29
+ 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.
30
+
31
+ ## Why this might be interesting anyway
32
+
33
+ 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.
34
+
35
+ That's the honest pitch. No claims beyond it yet.
36
+
37
+ ## Files
38
+
39
+ | File | What it does |
40
+ |---|---|
41
+ | `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). |
42
+ | `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. |
43
+
44
+ ## Requirements
45
+
46
+ ```bash
47
+ pip install torch transformers safetensors numpy
48
+ ```
49
+
50
+ 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.
51
+
52
+ ## How to run
53
+
54
+ ### 1. Bayesian feature selector chat (main experiment)
55
+
56
+ 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.
57
+
58
+ ```bash
59
+ python terminal_chat_bayesian.py
60
+ ```
61
+
62
+ In the chat session:
63
+ - Type normally to talk to the model
64
+ - `/stats` — shows how many features were flagged vs. total possible in the last response
65
+ - `/bayes` — shows the top 10 layers by number of currently-flagged features
66
+ - `/clear` — resets conversation history
67
+ - `/exit` — quit
68
+
69
+ ### 2. Storage/reconstruction test (secondary, not a compression result)
70
+
71
+ 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).
72
+
73
+ ```bash
74
+ python storage_reconstruction_test.py
75
+ ```
76
+
77
+ 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.
78
+
79
+ ## Code
80
+
81
+ ### `terminal_chat_bayesian.py`
82
+
83
+ ```python
84
+ import torch
85
+ import numpy as np
86
+ from safetensors.torch import load_file
87
+ from transformers import AutoTokenizer, AutoModelForCausalLM
88
+ import time
89
+ import os
90
+ import sys
91
+
92
+ MODEL_NAME = "Qwen/Qwen2-0.5B"
93
+ MAX_NEW_TOKENS = 200
94
+ TEMPERATURE = 0.7
95
+ ANALYSIS_FILE = "bayes_analysis.safetensors"
96
+ SYSTEM_PROMPT = "You are a helpful assistant."
97
+ NUM_FEATURES = 64
98
+ BAYES_EVERY_N = 8 # compute bayes stats every N tokens instead of every token
99
+ BAYES_ENABLED = True # can be fully disabled with this flag
100
+
101
+
102
+ def _row_features_torch(x: torch.Tensor, n_features: int = NUM_FEATURES) -> torch.Tensor:
103
+ x = x.float()
104
+ L = x.shape[0]
105
+ mean = x.mean()
106
+ std = x.std(unbiased=False)
107
+ abs_x = x.abs()
108
+
109
+ feats = torch.zeros(n_features, dtype=torch.float32, device=x.device)
110
+ feats[0] = mean
111
+ feats[1] = std
112
+ feats[2] = x.max()
113
+ feats[3] = x.min()
114
+
115
+ q = torch.quantile(x, torch.tensor([0.25, 0.5, 0.75, 0.05, 0.10, 0.90, 0.95], device=x.device))
116
+ feats[4], feats[5], feats[6] = q[0], q[1], q[2]
117
+ feats[16], feats[17], feats[18], feats[19] = q[3], q[4], q[5], q[6]
118
+
119
+ feats[7] = (x > mean + std).sum()
120
+ feats[8] = (x < mean - std).sum()
121
+ feats[9] = abs_x.mean()
122
+ feats[10] = abs_x.median()
123
+
124
+ w = 8
125
+ if L >= w:
126
+ wins = x.unfold(0, w, 1)
127
+ feats[11] = wins.mean(dim=1).mean()
128
+ feats[12] = wins.std(dim=1, unbiased=False).mean()
129
+ feats[13] = wins.max(dim=1).values.mean()
130
+ feats[14] = wins.min(dim=1).values.mean()
131
+ feats[15] = x.diff().abs().mean()
132
+ else:
133
+ feats[11], feats[12], feats[13], feats[14], feats[15] = mean, std, x.max(), x.min(), 0.0
134
+
135
+ if L > 1 and std > 1e-12:
136
+ a, b = x[:-1], x[1:]
137
+ a_c, b_c = a - a.mean(), b - b.mean()
138
+ denom = torch.sqrt((a_c * a_c).sum() * (b_c * b_c).sum())
139
+ feats[20] = (a_c * b_c).sum() / denom if denom > 1e-12 else 0.0
140
+ else:
141
+ feats[20] = 0.0
142
+
143
+ return feats[:n_features]
144
+
145
+
146
+ class BayesData:
147
+ def __init__(self, path: str = ANALYSIS_FILE):
148
+ if not os.path.exists(path):
149
+ print(f"[error] {path} not found. Run the analysis pass first to generate it.")
150
+ sys.exit(1)
151
+ print("[bayes-data] loading from safetensors ...")
152
+ raw = load_file(path)
153
+ self.layers = {}
154
+ names = {k[: -len("__feat")] for k in raw.keys() if k.endswith("__feat")}
155
+ for sk in names:
156
+ param_name = sk.replace("__", ".")
157
+ self.layers[param_name] = {"feat": raw[f"{sk}__feat"].float().numpy()}
158
+ print(f"[bayes-data] loaded {len(self.layers)} layers")
159
+
160
+ def get(self, param_name):
161
+ return self.layers.get(param_name)
162
+
163
+ def num_features_for(self, param_name) -> int:
164
+ data = self.layers.get(param_name)
165
+ return 1 if data is None else data["feat"].shape[1]
166
+
167
+
168
+ class BayesianFeatureSelector:
169
+ def __init__(self, n_features: int, device):
170
+ self.n_features = n_features
171
+ self.marked_counts = torch.ones(n_features, dtype=torch.float32, device=device)
172
+ self.unmarked_counts = torch.ones(n_features, dtype=torch.float32, device=device)
173
+ self.running_mean = torch.zeros(n_features, dtype=torch.float32, device=device)
174
+ self.running_var = torch.ones(n_features, dtype=torch.float32, device=device)
175
+ self.n_seen = 0
176
+
177
+ def select(self, feat_vector: torch.Tensor) -> torch.Tensor:
178
+ if self.n_seen == 0:
179
+ return torch.arange(self.n_features, device=feat_vector.device)
180
+ std = torch.sqrt(self.running_var) + 1e-8
181
+ deviation = (feat_vector - self.running_mean).abs() / std
182
+ marked = torch.where(deviation > 1.0)[0]
183
+ if marked.numel() == 0:
184
+ priors = self.marked_counts / (self.marked_counts + self.unmarked_counts)
185
+ marked = priors.argmax().unsqueeze(0)
186
+ return marked
187
+
188
+ def update(self, feat_vector: torch.Tensor, marked_idx: torch.Tensor):
189
+ marked_mask = torch.zeros(self.n_features, dtype=torch.bool, device=feat_vector.device)
190
+ marked_mask[marked_idx] = True
191
+ self.marked_counts[marked_mask] += 1
192
+ self.unmarked_counts[~marked_mask] += 1
193
+
194
+ self.n_seen += 1
195
+ delta = feat_vector - self.running_mean
196
+ self.running_mean += delta / self.n_seen
197
+ delta2 = feat_vector - self.running_mean
198
+ self.running_var += (delta * delta2 - self.running_var) / self.n_seen
199
+ self.running_var.clamp_(min=1e-8)
200
+
201
+
202
+ class LayerBayesRegistry:
203
+ def __init__(self, layer_names: list, n_features: int, device):
204
+ self.selectors = {name: BayesianFeatureSelector(n_features, device) for name in layer_names}
205
+ self.layer_order = layer_names
206
+ self.n_features = n_features
207
+
208
+ def select_for(self, layer_name: str, feat_vector: torch.Tensor) -> torch.Tensor:
209
+ return self.selectors[layer_name].select(feat_vector)
210
+
211
+ def observe(self, layer_name: str, feat_vector: torch.Tensor, marked_idx: torch.Tensor):
212
+ self.selectors[layer_name].update(feat_vector, marked_idx)
213
+
214
+ def state_summary(self) -> dict:
215
+ out = {}
216
+ for name, sel in self.selectors.items():
217
+ if sel.n_seen == 0:
218
+ out[name] = sel.n_features
219
+ else:
220
+ priors = sel.marked_counts / (sel.marked_counts + sel.unmarked_counts)
221
+ out[name] = int((priors > 0.5).sum().item())
222
+ return out
223
+
224
+
225
+ def build_bayes_registry(model, bayes_data, device) -> LayerBayesRegistry:
226
+ layer_names = []
227
+ n_features = NUM_FEATURES
228
+ for name, module in model.named_modules():
229
+ param_name = f"{name}.weight"
230
+ if bayes_data.get(param_name) is not None and hasattr(module, "weight"):
231
+ layer_names.append(param_name)
232
+ n_features = bayes_data.num_features_for(param_name)
233
+ print(f"[registry] {len(layer_names)} layers, n_features={n_features}")
234
+ return LayerBayesRegistry(layer_names, n_features, device)
235
+
236
+
237
+ def generate_with_bayes_scalar(model, tokenizer, history, bayes_data, bayes_registry):
238
+ try:
239
+ prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
240
+ except Exception:
241
+ prompt = "\n".join(f"{m['role'].upper()}: {m['content']}" for m in history) + "\nASSISTANT:"
242
+
243
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
244
+ input_ids = inputs["input_ids"]
245
+
246
+ total_scalars_used = 0
247
+ total_elements_all = 0
248
+ activation_store = {}
249
+ hooks = []
250
+ layer_meta = {}
251
+
252
+ if BAYES_ENABLED:
253
+ for name, module in model.named_modules():
254
+ param_name = f"{name}.weight"
255
+ data = bayes_data.get(param_name)
256
+ if data is None or not hasattr(module, "weight"):
257
+ continue
258
+ w = module.weight
259
+ dim_in = w.shape[1] if w.ndim >= 2 else w.shape[0]
260
+ n_features = data["feat"].shape[1]
261
+ layer_meta[param_name] = (dim_in, n_features)
262
+
263
+ def make_hook(pn, di):
264
+ def hook_fn(module, inp, out):
265
+ x_in = inp[0]
266
+ if x_in.ndim == 3:
267
+ x_t = x_in[0, -1, :]
268
+ elif x_in.ndim == 2:
269
+ x_t = x_in[0, :]
270
+ else:
271
+ return
272
+ if x_t.shape[0] == di:
273
+ activation_store[pn] = x_t.detach()
274
+ return hook_fn
275
+
276
+ hooks.append(module.register_forward_hook(make_hook(param_name, dim_in)))
277
+
278
+ vocab_size = tokenizer.vocab_size or model.config.vocab_size
279
+ new_tokens_list = []
280
+ step_counter = 0
281
+
282
+ with torch.no_grad():
283
+ past_key_values = None
284
+ cur_input = input_ids
285
+
286
+ for step in range(MAX_NEW_TOKENS):
287
+ if step == 0:
288
+ out = model(input_ids=cur_input, use_cache=True)
289
+ else:
290
+ out = model(input_ids=cur_input, past_key_values=past_key_values, use_cache=True)
291
+
292
+ past_key_values = out.past_key_values
293
+ logits = out.logits[:, -1, :vocab_size].float()
294
+
295
+ torch.nan_to_num_(logits, nan=0.0, posinf=1e4, neginf=-1e4)
296
+ logits.div_(max(TEMPERATURE, 1e-6))
297
+
298
+ sorted_logits, sorted_idx = torch.sort(logits, descending=True)
299
+ probs_sorted = torch.softmax(sorted_logits, dim=-1)
300
+ cumprobs = torch.cumsum(probs_sorted, dim=-1)
301
+ mask = (cumprobs - probs_sorted) > 0.9
302
+ sorted_logits[mask] = -1e9
303
+
304
+ probs = torch.softmax(sorted_logits, dim=-1)
305
+ probs.clamp_(min=0.0)
306
+ s = probs.sum(dim=-1, keepdim=True)
307
+ if not (s == 0).any():
308
+ probs.div_(s)
309
+ else:
310
+ probs.fill_(1.0 / probs.shape[-1])
311
+
312
+ next_sorted = torch.multinomial(probs, num_samples=1)
313
+ next_token = sorted_idx.gather(-1, next_sorted)
314
+ next_id = next_token.item()
315
+ new_tokens_list.append(next_id)
316
+
317
+ if BAYES_ENABLED and (step_counter % BAYES_EVERY_N == 0) and activation_store:
318
+ for param_name, x_t in activation_store.items():
319
+ dim_in, n_features = layer_meta[param_name]
320
+ feat_vector = _row_features_torch(x_t, n_features)
321
+ marked_idx = bayes_registry.select_for(param_name, feat_vector)
322
+ total_scalars_used += marked_idx.numel()
323
+ total_elements_all += n_features
324
+ bayes_registry.observe(param_name, feat_vector, marked_idx)
325
+ activation_store.clear()
326
+ step_counter += 1
327
+
328
+ if next_id == tokenizer.eos_token_id:
329
+ break
330
+
331
+ cur_input = next_token
332
+
333
+ for h in hooks:
334
+ h.remove()
335
+
336
+ response_text = tokenizer.decode(new_tokens_list, skip_special_tokens=True)
337
+ pct = 100.0 * total_scalars_used / total_elements_all if total_elements_all > 0 else 0.0
338
+ return response_text, total_scalars_used, total_elements_all, pct
339
+
340
+
341
+ BANNER = """
342
+ +==================================================================+
343
+ | Qwen2-0.5B x Bayesian Minimal Feature Selection |
344
+ | /stats - stats for the last response |
345
+ | /bayes - state of the bayesian models (top 10 by k) |
346
+ | /clear - clear history |
347
+ | /exit - quit |
348
+ +==================================================================+
349
+ """
350
+
351
+
352
+ def chat(model, tokenizer, bayes_data, bayes_registry):
353
+ print(BANNER)
354
+ history = [{"role": "system", "content": SYSTEM_PROMPT}]
355
+ last_stats = None
356
+
357
+ while True:
358
+ try:
359
+ user = input("You: ").strip()
360
+ except (EOFError, KeyboardInterrupt):
361
+ print("\nExiting.")
362
+ break
363
+
364
+ if not user:
365
+ continue
366
+ if user == "/exit":
367
+ break
368
+ if user == "/clear":
369
+ history = [{"role": "system", "content": SYSTEM_PROMPT}]
370
+ print("[history cleared]")
371
+ continue
372
+ if user == "/stats":
373
+ if last_stats:
374
+ sc, el, pct = last_stats
375
+ print(f"\n Scalars flagged : {sc:,}")
376
+ print(f" Total possible : {el:,}")
377
+ print(f" Fraction flagged : {pct:.4f}%\n")
378
+ else:
379
+ print("[no data yet - send a message first]")
380
+ continue
381
+ if user == "/bayes":
382
+ summary = bayes_registry.state_summary()
383
+ print("\n [bayesian state - top 10 layers by k]")
384
+ for name, k in sorted(summary.items(), key=lambda x: -x[1])[:10]:
385
+ print(f" {name:<55} k={k}")
386
+ print()
387
+ continue
388
+
389
+ history.append({"role": "user", "content": user})
390
+ t0 = time.time()
391
+
392
+ resp, scalars_used, total_elements, pct = generate_with_bayes_scalar(
393
+ model, tokenizer, history, bayes_data, bayes_registry
394
+ )
395
+
396
+ history.append({"role": "assistant", "content": resp})
397
+ elapsed = time.time() - t0
398
+ last_stats = (scalars_used, total_elements, pct)
399
+
400
+ print(f"\nModel ({elapsed:.1f}s): {resp}")
401
+ print(f"\n +- Bayesian minimal feature selection -----------------+")
402
+ print(f" | Flagged : {scalars_used:>15,} |")
403
+ print(f" | Total : {total_elements:>15,} |")
404
+ print(f" | Fraction : {pct:>14.4f} % |")
405
+ print(f" +--------------------------------------------------------+\n")
406
+
407
+
408
+ if __name__ == "__main__":
409
+ if not torch.cuda.is_available():
410
+ print("[error] CUDA not available. This script is configured for GPU.")
411
+ sys.exit(1)
412
+
413
+ device = "cuda"
414
+ print(f"[start] device: {device}")
415
+
416
+ torch.backends.cuda.matmul.allow_tf32 = True
417
+ torch.backends.cudnn.allow_tf32 = True
418
+ torch.backends.cudnn.benchmark = True
419
+
420
+ print(f"\n[1/3] Loading {MODEL_NAME} ...")
421
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
422
+ model = AutoModelForCausalLM.from_pretrained(
423
+ MODEL_NAME,
424
+ torch_dtype=torch.bfloat16,
425
+ device_map=device,
426
+ trust_remote_code=True,
427
+ )
428
+ model.eval()
429
+
430
+ print("\n[2/3] Loading features from analysis file ...")
431
+ bayes_data = BayesData()
432
+
433
+ print("\n[3/3] Initializing bayesian feature selection registry ...")
434
+ bayes_registry = build_bayes_registry(model, bayes_data, device)
435
+
436
+ chat(model, tokenizer, bayes_data, bayes_registry)
437
+ ```
438
+
439
+ ### `storage_reconstruction_test.py`
440
+
441
+ ```python
442
+ import torch
443
+ import numpy as np
444
+ from transformers import AutoModelForCausalLM, AutoTokenizer
445
+ import time
446
+ import json
447
+ import os
448
+ from safetensors.torch import load_file
449
+
450
+ # ==========================================
451
+ # 1. STORAGE / LOADING TEST (NOT A COMPRESSION RESULT)
452
+ # ==========================================
453
+ # NOTE: reconstruction below is mean + residual, which is mathematically
454
+ # exact by construction (mean + (original - mean) = original).
455
+ # R2 = 1.0 is expected here and does not indicate compression -
456
+ # it indicates the two files together contain the same information
457
+ # as the original weight, just split across two files.
458
+
459
+ class FastBayesianStorage:
460
+ """Weight storage split across two files, for testing a load pipeline"""
461
+ def __init__(self):
462
+ self.base_predictions = {}
463
+ self.layer_residuals = {}
464
+ self.layer_shapes = {}
465
+
466
+ def decompress_layer(self, name):
467
+ """Exact reconstruction: mean_val + residual = original (by construction)"""
468
+ shape = self.layer_shapes[name]
469
+ mean_val = self.base_predictions[name]
470
+ residual = self.layer_residuals[name]
471
+
472
+ reconstructed = np.full(residual.shape, mean_val, dtype=np.float32) + residual
473
+
474
+ return torch.from_numpy(reconstructed).view(shape)
475
+
476
+ def load_from_files(self, json_path="/content/bayesian_features.json", safetensors_path="/content/layer_residuals.safetensors"):
477
+ """Loads scalar features from JSON and residual tensors from Safetensors"""
478
+ print(f"\n[Import] Loading features and layer residuals from files...")
479
+
480
+ # 1. Load metadata and scalar features
481
+ with open(json_path, "r", encoding="utf-8") as f:
482
+ json_data = json.load(f)
483
+
484
+ self.base_predictions = json_data["base_predictions"]
485
+ self.layer_shapes = json_data["layer_shapes"]
486
+ print(f" -> Scalar features and shapes loaded from: {json_path}")
487
+
488
+ # 2. Load residual tensors (convert Torch -> NumPy for reconstruction)
489
+ tensors_dict = load_file(safetensors_path)
490
+ for name, tensor in tensors_dict.items():
491
+ self.layer_residuals[name] = tensor.numpy()
492
+ print(f" -> Residual tensors loaded from: {safetensors_path}")
493
+
494
+
495
+ # ==========================================
496
+ # 2. BUILD HYBRID MODEL FROM SPLIT FILES
497
+ # ==========================================
498
+
499
+ def prepare_fast_hybrid_model(model_name="Qwen/Qwen2-0.5B"):
500
+ start_time = time.time()
501
+ print(f"Loading base model and tokenizer {model_name}...")
502
+
503
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
504
+ model = AutoModelForCausalLM.from_pretrained(
505
+ model_name, torch_dtype=torch.float32, device_map="cpu", low_cpu_mem_usage=True
506
+ )
507
+
508
+ target_layers = ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj']
509
+ compressed_layer_names = []
510
+
511
+ print("\n[Process] Removing original attention/FFN weight tensors (preparing to load from files)...")
512
+ for name, param in list(model.named_parameters()):
513
+ if any(target in name for target in target_layers) and "weight" in name:
514
+ compressed_layer_names.append(name)
515
+
516
+ # remove original weight to simulate a clean storage state
517
+ delattr(model.get_submodule(name.rsplit('.', 1)[0]), 'weight')
518
+
519
+ print(f"\n[Done] Structure preparation time: {time.time() - start_time:.2f} sec.")
520
+ return model, tokenizer, compressed_layer_names
521
+
522
+ # ==========================================
523
+ # 3. TERMINAL CHAT
524
+ # ==========================================
525
+
526
+ def run_fast_terminal_chat():
527
+ # Paths to your prepared files
528
+ json_path = "/content/bayesian_features.json"
529
+ safetensors_path = "/content/layer_residuals.safetensors"
530
+
531
+ # Build empty model structure
532
+ model, tokenizer, compressed_names = prepare_fast_hybrid_model()
533
+
534
+ # Initialize storage and load the already-prepared files (no overwrite)
535
+ storage = FastBayesianStorage()
536
+ storage.load_from_files(json_path=json_path, safetensors_path=safetensors_path)
537
+
538
+ # Reconstruct weights from loaded files
539
+ start_restore = time.time()
540
+ print("\n[Info] Reconstructing weight tensors from loaded files...")
541
+ for name in compressed_names:
542
+ restored_tensor = storage.decompress_layer(name)
543
+ submodule = model.get_submodule(name.rsplit('.', 1)[0])
544
+ submodule.weight = torch.nn.Parameter(restored_tensor)
545
+ print(f"[Done] All weights reconstructed (R2=1.0 by construction, see note above) in: {time.time() - start_restore:.2f} sec!")
546
+
547
+ print("\n" + "="*50)
548
+ print(" QWEN-0.5B CHAT - RECONSTRUCTED FROM SPLIT FILES")
549
+ print(" Type 'exit' to quit.")
550
+ print("="*50 + "\n")
551
+
552
+ while True:
553
+ user_input = input("You: ")
554
+ if user_input.lower() in ['exit', 'quit']:
555
+ break
556
+
557
+ if not user_input.strip():
558
+ continue
559
+
560
+ messages = [{"role": "user", "content": user_input}]
561
+ text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
562
+ model_inputs = tokenizer([text], return_tensors="pt")
563
+
564
+ print("Qwen: ", end="", flush=True)
565
+ generated_ids = model_inputs.input_ids
566
+
567
+ with torch.no_grad():
568
+ for _ in range(70):
569
+ outputs = model(input_ids=generated_ids)
570
+ next_token_logits = outputs.logits[:, -1, :]
571
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
572
+
573
+ token_str = tokenizer.decode(next_token[0], skip_special_tokens=True)
574
+ print(token_str, end="", flush=True)
575
+
576
+ generated_ids = torch.cat([generated_ids, next_token], dim=-1)
577
+ if next_token.item() in [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>")]:
578
+ break
579
+ print("\n" + "-"*50)
580
+
581
+ if __name__ == "__main__":
582
+ run_fast_terminal_chat()
583
+ ```
584
+
585
+ ## Results
586
+
587
+ *Placeholder — to be filled in with real numbers from benchmark runs.*
588
+
589
+ - [ ] Fraction of features flagged per layer, averaged across a test set of prompts
590
+ - [ ] How the flagged fraction changes over the course of a generation (early tokens vs. late tokens)
591
+ - [ ] Per-layer comparison: which layers have consistently high vs. low flagged fractions
592
+ - [ ] 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"
593
+
594
+ ## Open questions / next steps
595
+
596
+ - Does the flagged feature subset stay stable across different prompts, or does it change drastically session to session?
597
+ - Is there a relationship between which features get flagged and attention patterns in the same layer?
598
+ - 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.
599
+
600
+ ## License
601
+
602
+ Apache 2.0, matching the base model license. This repo builds on [Qwen/Qwen2-0.5B](https://huggingface.co/Qwen/Qwen2-0.5B).