amachine.am_transformers.am_basis_eval

  1import contextlib
  2import math
  3from dataclasses import dataclass
  4from typing import Any, cast
  5
  6import torch
  7
  8# --- Locating Decoder Blocks -----------------------------------------------
  9
 10from .am_module_resolver import (
 11    find_decoder_layers
 12)
 13
 14# --- Statistics & Math -----------------------------------------------------
 15
 16def _moments_inplace(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
 17    """Per-row mean, variance, Pearson skewness, and Pearson kurtosis, in one pass.
 18
 19    Destructive to save VRAM: `x` is overwritten with centered values and then
 20    with their square, so only one extra same-sized buffer (`sq`) is ever
 21    allocated instead of one per moment.
 22    """
 23    mean = x.mean(dim=-1)
 24    x.sub_(mean.unsqueeze(-1))           # x now holds centered values
 25    sq = x.square()                      # sq = centered^2
 26    var = sq.mean(dim=-1).clamp_min(1e-12)
 27    skew = (sq * x).mean(dim=-1) / var.pow(1.5)   # sq * x = centered^3
 28    sq.square_()                         # reuse buffer: sq now = centered^4
 29    kurt = sq.mean(dim=-1) / var.square()
 30    return mean, var, skew, kurt
 31
 32def _hoyer_sparsity(x: torch.Tensor) -> torch.Tensor:
 33    """Per-row Hoyer sparseness on the raw (uncentered) values: 0 = dense/uniform
 34    magnitude across dims, 1 = maximally sparse (a single dominant dim).
 35
 36    Non-destructive and cheap (two reductions) since it must be measured before
 37    `_moments_inplace` overwrites the row with centered values.
 38    """
 39    d = x.shape[-1]
 40    l1 = x.abs().sum(dim=-1)
 41    l2 = x.square().sum(dim=-1).sqrt().clamp_min(1e-12)
 42    return (math.sqrt(d) - l1 / l2) / (math.sqrt(d) - 1)
 43
 44def _rank(x: torch.Tensor) -> torch.Tensor:
 45    """Assigns 0..n-1 ranks to the values in x. Ties are broken by index,
 46    which is fine here: this is only ever used on continuous-valued
 47    per-channel statistics, where exact ties are effectively never seen.
 48    """
 49    order = torch.argsort(x)
 50    ranks = torch.empty_like(order, dtype=torch.float64)
 51    ranks[order] = torch.arange(x.shape[0], dtype=torch.float64)
 52    return ranks
 53
 54def _spearman(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
 55    """Spearman rank correlation between two equal-length vectors."""
 56    ra_raw, rb_raw = _rank(a), _rank(b)
 57    ra, rb = ra_raw - ra_raw.mean(), rb_raw - rb_raw.mean()
 58    denom = (ra.square().sum().sqrt() * rb.square().sum().sqrt()).clamp_min(1e-12)
 59    return (ra * rb).sum() / denom
 60
 61def _haar_orthogonal(d: int, seed: int) -> torch.Tensor:
 62    """Generate a CPU float32 Haar-distributed orthogonal matrix."""
 63    generator = torch.Generator(device="cpu").manual_seed(seed % (2**63))
 64    q, r = torch.linalg.qr(torch.randn(d, d, generator=generator, dtype=torch.float32))
 65    signs = torch.where(torch.diagonal(r) >= 0, 1.0, -1.0)
 66    return q * signs
 67
 68# --- State Management ------------------------------------------------------
 69
 70@dataclass
 71class DeviceState:
 72    """Consolidates running metrics for a specific device."""
 73    native: torch.Tensor
 74    rotated: torch.Tensor
 75    count: torch.Tensor
 76    mean_native: torch.Tensor
 77    mean_rotated: torch.Tensor
 78    var_native: torch.Tensor
 79    var_rotated: torch.Tensor
 80    skew_native: torch.Tensor
 81    skew_rotated: torch.Tensor
 82    sparsity_native: torch.Tensor
 83    sparsity_rotated: torch.Tensor
 84    mask: torch.Tensor | None = None
 85    blocks: tuple[torch.Tensor, ...] | None = None
 86    sum_x: torch.Tensor | None = None
 87    sum_xx: torch.Tensor | None = None
 88    running_max: torch.Tensor | None = None
 89    running_min: torch.Tensor | None = None
 90    # Per-channel power sums (native basis only -- see PrivilegedBasisEvaluator
 91    # docstring on compute_channel_stats for why rotated channels aren't tracked)
 92    chan_sum_x: torch.Tensor | None = None
 93    chan_sum_x2: torch.Tensor | None = None
 94    chan_sum_x3: torch.Tensor | None = None
 95    chan_sum_x4: torch.Tensor | None = None
 96    # Same, split into two disjoint halves of the token stream (alternated per
 97    # observe() call) for the channel-persistence check
 98    chan_sum_x_a: torch.Tensor | None = None
 99    chan_sum_x2_a: torch.Tensor | None = None
100    chan_count_a: torch.Tensor | None = None
101    chan_sum_x_b: torch.Tensor | None = None
102    chan_sum_x2_b: torch.Tensor | None = None
103    chan_count_b: torch.Tensor | None = None
104
105# --- Evaluator -------------------------------------------------------------
106
107class PrivilegedBasisEvaluator(contextlib.AbstractContextManager):
108    _AUTO_ROTATION_BUDGET = 256
109    _MAX_CHUNK_ELEMENTS = 8_000_000
110    _MAX_ROTATION_BLOCK_ELEMENTS = 8_000_000
111
112    def __init__(
113        self,
114        model,
115        pad_id: int | None = None,
116        seed: int = 0,
117        num_rotations: int | None = None,
118        compute_effective_rank: bool = False,
119        compute_channel_stats: bool = False
120    ):
121        self.pad_id = pad_id
122        self.seed = int(seed)
123        self.num_rotations = num_rotations
124        self.compute_effective_rank = compute_effective_rank
125        # Optional and off by default: per-channel (across-token) statistics
126        # are a handful of extra reductions per chunk plus 10 small
127        # (n_layers x hidden_size or n_layers) accumulators -- cheap relative
128        # to the effective-rank covariance matrix, but not free, so this is
129        # opt-in rather than bundled into compute_effective_rank.
130        self.compute_channel_stats = compute_channel_stats
131
132        self.layers = find_decoder_layers(model)
133        self.n_layers = len(self.layers)
134
135        # Public CPU snapshots updated by finalize()
136        self.sum_kurt = torch.zeros(self.n_layers, dtype=torch.float64)
137        self.sum_kurt_rot = torch.zeros(self.n_layers, dtype=torch.float64)
138        self.sum_mean = torch.zeros(self.n_layers, dtype=torch.float64)
139        self.sum_mean_rot = torch.zeros(self.n_layers, dtype=torch.float64)
140        self.sum_var = torch.zeros(self.n_layers, dtype=torch.float64)
141        self.sum_var_rot = torch.zeros(self.n_layers, dtype=torch.float64)
142        self.sum_skew = torch.zeros(self.n_layers, dtype=torch.float64)
143        self.sum_skew_rot = torch.zeros(self.n_layers, dtype=torch.float64)
144        self.sum_sparsity = torch.zeros(self.n_layers, dtype=torch.float64)
145        self.sum_sparsity_rot = torch.zeros(self.n_layers, dtype=torch.float64)
146        self.token_count = torch.zeros(self.n_layers, dtype=torch.float64)
147        # Raw per-channel moment vectors (n_layers x hidden_size), populated by
148        # finalize() only if compute_channel_stats=True. Intended for on-demand
149        # visual inspection (see plot_channel_concentration) rather than as
150        # flat per-layer scalars like the rest of the results dict.
151        self.channel_moments: dict[str, torch.Tensor] | None = None
152        # self.outlier_dims: torch.Tensor | None = None
153
154        self._handles: list = []
155        self._token_mask_cpu: torch.Tensor | None = None
156        self._states: dict[torch.device, DeviceState] = {}
157        self._effective_rotations: int | None = None
158        self._mask_all: bool = False
159        # Alternates every observe() call, splitting the token stream into two
160        # disjoint halves for the channel-persistence check.
161        self._split_toggle: bool = False
162        self._current_split: bool = False
163
164    def __enter__(self):
165        if self._handles:
166            raise RuntimeError("Context already entered.")
167        self._handles = [
168            layer.register_forward_hook(self._make_hook(i))
169            for i, layer in enumerate(self.layers)
170        ]
171        return self
172
173    def __exit__(self, *exc):
174        for handle in self._handles:
175            handle.remove()
176        self._handles.clear()
177        self._token_mask_cpu = None
178        self._mask_all = False
179        # Only release the transient, memory-heavy per-pass caches (rotation
180        # blocks, device mask). The running accumulators (native, rotated,
181        # count, outliers) must survive __exit__ since finalize() is called
182        # by the caller *after* the `with` block ends.
183        for state in self._states.values():
184            state.mask = None
185            state.blocks = None
186        return False
187
188    def observe(self, input_ids=None, attention_mask=None) -> None:
189        """Set the valid-token mask for the next forward pass."""
190        if attention_mask is not None:
191            self._token_mask_cpu = attention_mask.detach().reshape(-1).bool().cpu()
192        elif input_ids is not None and self.pad_id is None:
193            raise ValueError("pad_id required if attention_mask is omitted.")
194        elif input_ids is not None:
195            self._token_mask_cpu = (input_ids.detach() != self.pad_id).reshape(-1).cpu()
196        else:
197            raise ValueError("Provide attention_mask or input_ids.")
198        
199        # Clear lazily loaded per-device masks
200        for state in self._states.values():
201            state.mask = None
202
203        # Alternate which half of the token stream this forward pass counts
204        # toward, for the channel-persistence check (see compute_channel_stats)
205        self._current_split = self._split_toggle
206        self._split_toggle = not self._split_toggle
207
208        # Check if we can run the unmasked fast-path
209
210        assert self._token_mask_cpu is not None
211        self._mask_all = bool(self._token_mask_cpu.all())
212
213    def _get_device_state(self, device: torch.device, hidden_size: int) -> DeviceState:
214        """Retrieve or initialize consolidated state for a device."""
215        if device not in self._states:
216            
217            sum_x = torch.zeros ( 
218                (self.n_layers, hidden_size), 
219                device=device, dtype=torch.float64
220            ) if self.compute_effective_rank else None
221
222            sum_xx = torch.zeros(
223                ( self.n_layers, hidden_size, hidden_size), 
224                device=device, dtype=torch.float64
225            ) if self.compute_effective_rank else None
226
227            def _chan_vec():
228                return torch.zeros((self.n_layers, hidden_size), device=device, dtype=torch.float64) \
229                    if self.compute_channel_stats else None
230
231            def _chan_count():
232                return torch.zeros(self.n_layers, device=device, dtype=torch.float64) \
233                    if self.compute_channel_stats else None
234
235            self._states[device] = DeviceState(
236                native=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
237                rotated=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
238                count=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
239                mean_native=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
240                mean_rotated=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
241                var_native=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
242                var_rotated=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
243                skew_native=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
244                skew_rotated=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
245                sparsity_native=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
246                sparsity_rotated=torch.zeros(self.n_layers, device=device, dtype=torch.float64),
247                # outliers=torch.zeros(self.n_layers, hidden_size, device=device, dtype=torch.bool),
248                sum_x=sum_x,
249                sum_xx=sum_xx,
250                chan_sum_x=_chan_vec(),
251                chan_sum_x2=_chan_vec(),
252                chan_sum_x3=_chan_vec(),
253                chan_sum_x4=_chan_vec(),
254                chan_sum_x_a=_chan_vec(),
255                chan_sum_x2_a=_chan_vec(),
256                chan_count_a=_chan_count(),
257                chan_sum_x_b=_chan_vec(),
258                chan_sum_x2_b=_chan_vec(),
259                chan_count_b=_chan_count(),
260                running_max=torch.full(
261                    (self.n_layers, hidden_size), 
262                    float('-inf'), 
263                    device=device, 
264                    dtype=torch.float32
265                ),
266                running_min=torch.full(
267                    (self.n_layers, hidden_size), 
268                    float('inf'), 
269                    device=device, 
270                    dtype=torch.float32
271                )
272            )
273            
274            # Initialize batched rotations once per device
275            count = self.num_rotations or max(1, math.ceil(self._AUTO_ROTATION_BUDGET / hidden_size))
276            self._effective_rotations = count
277            
278            # Pack rotations into blocks to optimize GEMMs
279            by_matrix_budget = max(1, self._MAX_ROTATION_BLOCK_ELEMENTS // (hidden_size * hidden_size))
280            block_size = max(1, min(count, by_matrix_budget))
281            
282            blocks = []
283            for start in range(0, count, block_size):
284                block_rots = min(block_size, count - start)
285                mats = []
286                for idx in range(start, start + block_rots):
287                    seed = (self.seed + idx * 0x9E3779B97F4A7C15) % (2**63)
288                    mats.append(_haar_orthogonal(hidden_size, seed))
289                
290                packed_cpu = torch.cat(mats, dim=1).contiguous()
291                blocks.append(packed_cpu.to(device=device, dtype=torch.float32, non_blocking=True))
292            
293            self._states[device].blocks = tuple(blocks)
294
295        return self._states[device]
296
297    def _make_hook(self, layer_idx: int):
298
299        def hook(_module, _inputs, output):
300            
301            if self._token_mask_cpu is None:
302                return
303
304            hidden = output[0] if isinstance(output, (tuple, list)) else output
305            
306            # Safety check against unexpected output structures
307            if not isinstance(hidden, torch.Tensor) or getattr(hidden, "ndim", 0) < 1:
308                return
309
310            device = hidden.device
311            hidden_size = hidden.shape[-1]
312
313            amp_ctx = (
314                torch.autocast(device_type=device.type, enabled=False)
315                if device.type in ("cuda", "cpu")
316                else contextlib.nullcontext()
317            )
318
319            with torch.inference_mode(), amp_ctx:
320    
321                flat  = hidden.detach().reshape(-1, hidden_size)
322                state = self._get_device_state(device, hidden_size)
323
324                total_rows = flat.shape[0]
325
326                if self._token_mask_cpu.numel() != total_rows:
327                    raise ValueError(f"Mask mismatch: {self._token_mask_cpu.numel()} vs {total_rows}")
328
329                # Prove to static analyzer that variables initialized in _get_device_state are not None
330                assert state.blocks is not None
331                assert self._effective_rotations is not None
332
333                # Calculate memory-safe chunk sizes considering the batched rotations
334                max_block_rots = max(b.shape[1] // hidden_size for b in state.blocks)
335                denominator = hidden_size * (1.0 + max_block_rots) + 3.0 * max_block_rots
336                chunk_tokens = min(total_rows, max(1, int(self._MAX_CHUNK_ELEMENTS // max(1.0, denominator))))
337                
338                # Copy the mask lazily only if there are tokens to mask
339                mask = None
340                
341                if not self._mask_all:
342                    if state.mask is None:
343                        state.mask = self._token_mask_cpu.to(device, non_blocking=True)
344                    
345                    mask = state.mask
346
347                for start in range(0, total_rows, chunk_tokens):
348
349                    end   = min(total_rows, start + chunk_tokens)
350                    chunk = flat[start:end].to(dtype=torch.float32, copy=True)
351                    
352                    m_chunk: torch.Tensor | None = mask[start:end] if mask is not None else None
353                    
354                    # Rotated Moments (Batched): mean, variance, skewness, kurtosis, sparsity
355                    rotated_mean_total = torch.zeros((), device=device, dtype=torch.float64)
356                    rotated_var_total = torch.zeros((), device=device, dtype=torch.float64)
357                    rotated_skew_total = torch.zeros((), device=device, dtype=torch.float64)
358                    rotated_kurt_total = torch.zeros((), device=device, dtype=torch.float64)
359                    rotated_sparsity_total = torch.zeros((), device=device, dtype=torch.float64)
360
361                    for block in state.blocks:
362                        
363                        block_rots = block.shape[1] // hidden_size
364                        rotated = (chunk @ block).view(chunk.shape[0], block_rots, hidden_size)
365                        
366                        # Sparsity reads raw values, so it must run before the
367                        # destructive moments call overwrites `rotated` below.
368                        sparsity_rot = _hoyer_sparsity(rotated)
369                        mean_rot, var_rot, skew_rot, kurt_rot = _moments_inplace(rotated)
370                        if m_chunk is not None:
371
372                            invalid = (~m_chunk).view(-1, 1).expand_as(kurt_rot)
373                            mean_rot.masked_fill_(invalid, 0)
374                            var_rot.masked_fill_(invalid, 0)
375                            skew_rot.masked_fill_(invalid, 0)
376                            kurt_rot.masked_fill_(invalid, 0)
377                            sparsity_rot.masked_fill_(invalid, 0)
378                        
379                        rotated_mean_total.add_(mean_rot.sum(dtype=torch.float64))
380                        rotated_var_total.add_(var_rot.sum(dtype=torch.float64))
381                        rotated_skew_total.add_(skew_rot.sum(dtype=torch.float64))
382                        rotated_kurt_total.add_(kurt_rot.sum(dtype=torch.float64))
383                        rotated_sparsity_total.add_(sparsity_rot.sum(dtype=torch.float64))
384
385                    # Row bookkeeping and valid-token extraction, done on the
386                    # pristine chunk *before* native moments are computed
387                    # in-place below (which overwrites chunk's contents).
388                    if m_chunk is None:
389                        state.count[layer_idx].add_(chunk.shape[0])
390                        valid_chunk = chunk
391                    else:
392                        state.count[layer_idx].add_(m_chunk.sum(dtype=torch.float64))
393                        valid_chunk = chunk[m_chunk] # Extract only valid tokens (copy)
394
395                    valid_chunk_64 = None
396
397                    # Accumulate for Covariance
398                    if self.compute_effective_rank and valid_chunk.shape[0] > 0:
399
400                        assert state.sum_x is not None
401                        assert state.sum_xx is not None
402
403                        # Cast to float64 to prevent numerical overflow across millions of tokens
404                        # (also copies, so this is safe from the in-place native moments below
405                        # even when valid_chunk aliases chunk directly, i.e. the unmasked case)
406                        valid_chunk_64 = valid_chunk.to(torch.float64)
407                        state.sum_x[layer_idx].add_(valid_chunk_64.sum(dim=0))
408
409                        # Efficiently compute sum of outer products: X^T @ X
410                        state.sum_xx[layer_idx].add_(valid_chunk_64.t() @ valid_chunk_64)
411
412                        assert state.running_min is not None
413                        assert state.running_max is not None
414
415                        state.running_max[layer_idx] = torch.maximum(state.running_max[layer_idx], valid_chunk.max(dim=0)[0])
416
417                        state.running_min[layer_idx] = torch.minimum(state.running_min[layer_idx], valid_chunk.min(dim=0)[0])
418
419                    # Per-channel power sums (native basis only: channel identity
420                    # isn't preserved under a random rotation, so there's nothing
421                    # meaningful to accumulate for a "rotated channel").
422                    if self.compute_channel_stats and valid_chunk.shape[0] > 0:
423
424                        assert state.chan_sum_x is not None and state.chan_sum_x2 is not None
425                        assert state.chan_sum_x3 is not None and state.chan_sum_x4 is not None
426                        assert state.chan_sum_x_a is not None and state.chan_sum_x2_a is not None
427                        assert state.chan_sum_x_b is not None and state.chan_sum_x2_b is not None
428                        assert state.chan_count_a is not None and state.chan_count_b is not None
429
430                        # Reuse the float64 cast from the covariance block above
431                        # when available; otherwise compute it fresh here.
432
433                        if valid_chunk_64 is None :
434                            valid_chunk_64 = valid_chunk.to(torch.float64)
435
436                        valid_chunk_64_ch = valid_chunk_64 if self.compute_effective_rank else valid_chunk.to(torch.float64)
437                        sq64 = valid_chunk_64_ch.square()
438
439                        state.chan_sum_x[layer_idx].add_(valid_chunk_64_ch.sum(dim=0))
440                        state.chan_sum_x2[layer_idx].add_(sq64.sum(dim=0))
441                        state.chan_sum_x3[layer_idx].add_((sq64 * valid_chunk_64_ch).sum(dim=0))
442                        state.chan_sum_x4[layer_idx].add_(sq64.square().sum(dim=0))
443
444                        if self._current_split:
445                            state.chan_sum_x_a[layer_idx].add_(valid_chunk_64_ch.sum(dim=0))
446                            state.chan_sum_x2_a[layer_idx].add_(sq64.sum(dim=0))
447                            state.chan_count_a[layer_idx].add_(valid_chunk.shape[0])
448                        else:
449                            state.chan_sum_x_b[layer_idx].add_(valid_chunk_64_ch.sum(dim=0))
450                            state.chan_sum_x2_b[layer_idx].add_(sq64.sum(dim=0))
451                            state.chan_count_b[layer_idx].add_(valid_chunk.shape[0])
452
453                    # Native Moments (sparsity + moments both need the pristine
454                    # chunk; sparsity is non-destructive but must still be read
455                    # before the in-place moments call below overwrites it)
456                    sparsity_native = _hoyer_sparsity(chunk)
457                    mean_native, var_native, skew_native, kurt_native = _moments_inplace(chunk)
458                    if m_chunk is not None:
459                        mean_native.masked_fill_(~m_chunk, 0)
460                        var_native.masked_fill_(~m_chunk, 0)
461                        skew_native.masked_fill_(~m_chunk, 0)
462                        kurt_native.masked_fill_(~m_chunk, 0)
463                        sparsity_native.masked_fill_(~m_chunk, 0)
464
465                    # Update States
466                    state.native[layer_idx].add_(kurt_native.sum(dtype=torch.float64))
467                    state.rotated[layer_idx].add_(rotated_kurt_total / self._effective_rotations)
468
469                    state.mean_native[layer_idx].add_(mean_native.sum(dtype=torch.float64))
470                    state.mean_rotated[layer_idx].add_(rotated_mean_total / self._effective_rotations)
471
472                    state.var_native[layer_idx].add_(var_native.sum(dtype=torch.float64))
473                    state.var_rotated[layer_idx].add_(rotated_var_total / self._effective_rotations)
474
475                    state.skew_native[layer_idx].add_(skew_native.sum(dtype=torch.float64))
476                    state.skew_rotated[layer_idx].add_(rotated_skew_total / self._effective_rotations)
477
478                    state.sparsity_native[layer_idx].add_(sparsity_native.sum(dtype=torch.float64))
479                    state.sparsity_rotated[layer_idx].add_(rotated_sparsity_total / self._effective_rotations)
480
481        return hook
482
483    def finalize(self) -> dict[str, Any]:
484        """Aggregate metrics across all devices to the CPU and compute statistics."""
485        
486        if not self._states:
487            return {}
488
489        sum_kurt = torch.zeros(self.n_layers, dtype=torch.float64)
490        sum_kurt_rot = torch.zeros(self.n_layers, dtype=torch.float64)
491        sum_mean = torch.zeros(self.n_layers, dtype=torch.float64)
492        sum_mean_rot = torch.zeros(self.n_layers, dtype=torch.float64)
493        sum_var = torch.zeros(self.n_layers, dtype=torch.float64)
494        sum_var_rot = torch.zeros(self.n_layers, dtype=torch.float64)
495        sum_skew = torch.zeros(self.n_layers, dtype=torch.float64)
496        sum_skew_rot = torch.zeros(self.n_layers, dtype=torch.float64)
497        sum_sparsity = torch.zeros(self.n_layers, dtype=torch.float64)
498        sum_sparsity_rot = torch.zeros(self.n_layers, dtype=torch.float64)
499        token_count = torch.zeros(self.n_layers, dtype=torch.float64)
500
501        for state in self._states.values():
502            sum_kurt += state.native.cpu().to(dtype=torch.float64)
503            sum_kurt_rot += state.rotated.cpu().to(dtype=torch.float64)
504            sum_mean += state.mean_native.cpu().to(dtype=torch.float64)
505            sum_mean_rot += state.mean_rotated.cpu().to(dtype=torch.float64)
506            sum_var += state.var_native.cpu().to(dtype=torch.float64)
507            sum_var_rot += state.var_rotated.cpu().to(dtype=torch.float64)
508            sum_skew += state.skew_native.cpu().to(dtype=torch.float64)
509            sum_skew_rot += state.skew_rotated.cpu().to(dtype=torch.float64)
510            sum_sparsity += state.sparsity_native.cpu().to(dtype=torch.float64)
511            sum_sparsity_rot += state.sparsity_rotated.cpu().to(dtype=torch.float64)
512            token_count += state.count.cpu().to(dtype=torch.float64)
513
514        self.sum_kurt = sum_kurt
515        self.sum_kurt_rot = sum_kurt_rot
516        self.sum_mean = sum_mean
517        self.sum_mean_rot = sum_mean_rot
518        self.sum_var = sum_var
519        self.sum_var_rot = sum_var_rot
520        self.sum_skew = sum_skew
521        self.sum_skew_rot = sum_skew_rot
522        self.sum_sparsity = sum_sparsity
523        self.sum_sparsity_rot = sum_sparsity_rot
524        self.token_count = token_count
525
526        seen = token_count > 0
527        if not seen.any().item():
528            return {}
529
530        mean_kurt = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
531        mean_kurt_rot = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
532        mean_mean = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
533        mean_mean_rot = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
534        mean_var = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
535        mean_var_rot = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
536        mean_skew = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
537        mean_skew_rot = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
538        mean_sparsity = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
539        mean_sparsity_rot = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
540
541        mean_kurt[seen] = sum_kurt[seen] / token_count[seen]
542        mean_kurt_rot[seen] = sum_kurt_rot[seen] / token_count[seen]
543        mean_mean[seen] = sum_mean[seen] / token_count[seen]
544        mean_mean_rot[seen] = sum_mean_rot[seen] / token_count[seen]
545        mean_var[seen] = sum_var[seen] / token_count[seen]
546        mean_var_rot[seen] = sum_var_rot[seen] / token_count[seen]
547        mean_skew[seen] = sum_skew[seen] / token_count[seen]
548        mean_skew_rot[seen] = sum_skew_rot[seen] / token_count[seen]
549        mean_sparsity[seen] = sum_sparsity[seen] / token_count[seen]
550        mean_sparsity_rot[seen] = sum_sparsity_rot[seen] / token_count[seen]
551
552        excess_per_layer = mean_kurt - mean_kurt_rot
553
554        ratio_per_layer = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
555        ratio_per_layer[seen] = mean_kurt[seen] / mean_kurt_rot[seen].clamp_min(1e-12)
556
557        seen_indices = torch.where(seen)[0]
558        local_max_index = excess_per_layer[seen].argmax()
559        
560        # Explicit int cast for tensor indices to satisfy static analyzer
561        max_excess_layer = int(seen_indices[local_max_index].item())
562
563        max_kurt = mean_kurt[seen].max().item()
564        rotated_baseline = mean_kurt_rot[seen].mean().item()
565        max_excess = excess_per_layer[max_excess_layer].item()
566        max_excess_ratio = ratio_per_layer[max_excess_layer].item()
567
568        if max_excess >= 1.0 and max_excess_ratio >= 2.0:
569            verdict = "strong privileged-basis signal"
570        elif max_excess >= 0.25 and max_excess_ratio >= 1.2:
571            verdict = "moderate privileged-basis signal"
572        else:
573            verdict = "no clear privileged-basis signal"
574
575        results = {
576            "pb_mean_per_layer": mean_mean.tolist(),
577            "pb_mean_per_layer_rotated": mean_mean_rot.tolist(),
578            "pb_variance_per_layer": mean_var.tolist(),
579            "pb_variance_per_layer_rotated": mean_var_rot.tolist(),
580            "pb_std_per_layer": mean_var.sqrt().tolist(),
581            "pb_std_per_layer_rotated": mean_var_rot.sqrt().tolist(),
582            "pb_skewness_per_layer": mean_skew.tolist(),
583            "pb_skewness_per_layer_rotated": mean_skew_rot.tolist(),
584            "pb_kurtosis_per_layer": mean_kurt.tolist(),
585            "pb_kurtosis_per_layer_rotated": mean_kurt_rot.tolist(),
586            "pb_sparsity_per_layer": mean_sparsity.tolist(),
587            "pb_sparsity_per_layer_rotated": mean_sparsity_rot.tolist(),
588            "pb_excess_kurtosis_per_layer": excess_per_layer.tolist(),
589            "pb_kurtosis_ratio_per_layer": ratio_per_layer.tolist(),
590            "pb_num_rotations": self._effective_rotations,
591            "pb_max_kurtosis": max_kurt,
592            "pb_rotated_kurtosis_baseline": rotated_baseline,
593            "pb_max_excess_kurtosis": max_excess,
594            "pb_max_excess_ratio": max_excess_ratio,
595            "pb_max_excess_layer": max_excess_layer,
596            "pb_verdict": verdict
597        }
598
599        # Effective Rank
600        if self.compute_effective_rank:
601
602            first_state = next(iter(self._states.values()))
603
604            assert first_state.sum_x is not None
605            assert first_state.sum_xx is not None
606
607            hidden_size = first_state.sum_x.shape[-1]
608            
609            sum_x_cpu  : torch.Tensor = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
610            sum_xx_cpu : torch.Tensor = torch.zeros((self.n_layers, hidden_size, hidden_size), dtype=torch.float64)
611
612            for state in self._states.values() :
613
614                assert state.sum_x is not None
615                assert state.sum_xx is not None
616
617                sum_x_cpu  += state.sum_x.cpu()
618                sum_xx_cpu += state.sum_xx.cpu()
619
620            effective_ranks = []
621            axis_alignment = []
622            for l in range(self.n_layers):
623                
624                if token_count[l] == 0:
625                    effective_ranks.append(float("nan"))
626                    axis_alignment.append(float("nan"))
627                    continue
628                    
629                N = token_count[l].item()
630                mean_x = sum_x_cpu[l] / N
631                
632                # Covariance matrix
633                cov = (sum_xx_cpu[l] / N) - torch.outer(mean_x, mean_x)
634                
635                # Symmetric matrix: eigh gives eigenvalues *and* eigenvectors,
636                # ascending order. The eigenvectors are what let us also check
637                # axis-alignment below, at no extra decomposition cost.
638                eigvals, eigvecs = torch.linalg.eigh(cov)
639                
640                # Relative floor: absolute clamps (e.g. a flat 1e-12) are wrong
641                # once you compare layers whose activation scale differs a lot
642                # (residual-stream norm grows with depth) — a fixed floor is
643                # either a no-op on large-scale layers or an inflating one on
644                # small-scale layers. Floor relative to this layer's own top
645                # eigenvalue instead, with a tiny absolute fallback only to
646                # guard against an all-zero covariance.
647                floor = torch.clamp_min(eigvals.max().clamp_min(0) * 1e-8, 1e-12)
648                eigvals = torch.clamp_min(eigvals, floor)
649                
650                # Normalize to probability distribution
651                p = eigvals / eigvals.sum()
652                
653                # Compute Shannon entropy
654                entropy = -(p * torch.log(p)).sum()
655                eff_rank = torch.exp(entropy).item()
656                
657                effective_ranks.append(eff_rank)
658
659                # Axis alignment: unlike the eigenvalues themselves, the
660                # eigenvectors *do* change under rotation of the underlying
661                # basis, so this is a direct (and much cheaper, since the
662                # decomposition is already in hand) complement to the
663                # kurtosis-based privileged-basis test above. For each
664                # eigenvector, the participation ratio of its own components
665                # (1 / sum(v_i^4), since ||v||=1) is 1 when it points along a
666                # single standard axis and hidden_size when it's spread
667                # evenly across all of them. We average this over eigenvectors,
668                # weighted by their eigenvalue share `p`, so directions that
669                # carry more of the variance dominate the score.
670                eigvec_ipr = 1.0 / eigvecs.pow(4).sum(dim=0).clamp_min(1e-12)
671                axis_alignment.append((p * eigvec_ipr).sum().item())
672
673            results["pb_effective_rank_per_layer"] = effective_ranks
674            results["pb_axis_alignment_per_layer"] = axis_alignment
675
676        # Per-channel statistics (native basis only -- see compute_channel_stats)
677        if self.compute_channel_stats:
678
679            first_state = next(iter(self._states.values()))
680            assert first_state.chan_sum_x is not None
681            hidden_size = first_state.chan_sum_x.shape[-1]
682
683            chan_sum_x  = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
684            chan_sum_x2 = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
685            chan_sum_x3 = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
686            chan_sum_x4 = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
687            chan_sum_x_a  = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
688            chan_sum_x2_a = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
689            chan_sum_x_b  = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
690            chan_sum_x2_b = torch.zeros((self.n_layers, hidden_size), dtype=torch.float64)
691            chan_count_a = torch.zeros(self.n_layers, dtype=torch.float64)
692            chan_count_b = torch.zeros(self.n_layers, dtype=torch.float64)
693
694            for state in self._states.values():
695
696                assert state.chan_sum_x is not None and state.chan_sum_x2 is not None
697                assert state.chan_sum_x3 is not None and state.chan_sum_x4 is not None
698                assert state.chan_sum_x_a is not None and state.chan_sum_x2_a is not None
699                assert state.chan_sum_x_b is not None and state.chan_sum_x2_b is not None
700                assert state.chan_count_a is not None and state.chan_count_b is not None
701
702                chan_sum_x    += state.chan_sum_x.cpu()
703                chan_sum_x2   += state.chan_sum_x2.cpu()
704                chan_sum_x3   += state.chan_sum_x3.cpu()
705                chan_sum_x4   += state.chan_sum_x4.cpu()
706                chan_sum_x_a  += state.chan_sum_x_a.cpu()
707                chan_sum_x2_a += state.chan_sum_x2_a.cpu()
708                chan_sum_x_b  += state.chan_sum_x_b.cpu()
709                chan_sum_x2_b += state.chan_sum_x2_b.cpu()
710                chan_count_a  += state.chan_count_a.cpu()
711                chan_count_b  += state.chan_count_b.cpu()
712
713            channel_mean = torch.full((self.n_layers, hidden_size), float("nan"), dtype=torch.float64)
714            channel_var  = torch.full((self.n_layers, hidden_size), float("nan"), dtype=torch.float64)
715            channel_skew = torch.full((self.n_layers, hidden_size), float("nan"), dtype=torch.float64)
716            channel_kurt = torch.full((self.n_layers, hidden_size), float("nan"), dtype=torch.float64)
717
718            eff_channel_count = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
719            channel_gini = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
720            channel_powerlaw_exponent = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
721            channel_persistence = torch.full((self.n_layers,), float("nan"), dtype=torch.float64)
722
723            n = hidden_size
724            log_rank_desc = torch.log(torch.arange(1, n + 1, dtype=torch.float64))
725            log_rank_desc_centered = log_rank_desc - log_rank_desc.mean()
726            rank_1_to_n = torch.arange(1, n + 1, dtype=torch.float64)
727
728            for l in range(self.n_layers):
729
730                N = token_count[l].item()
731                if N == 0:
732                    continue
733
734                mean_l = chan_sum_x[l] / N
735                e_x2 = chan_sum_x2[l] / N
736                e_x3 = chan_sum_x3[l] / N
737                e_x4 = chan_sum_x4[l] / N
738
739                # Raw -> central moment conversion (single pass at finalize
740                # time, so the usual cancellation risk of this formula is a
741                # non-issue here -- it's not being iterated per chunk).
742                m2 = (e_x2 - mean_l.square()).clamp_min(1e-12)
743                m3 = e_x3 - 3 * mean_l * e_x2 + 2 * mean_l.pow(3)
744                m4 = e_x4 - 4 * mean_l * e_x3 + 6 * mean_l.square() * e_x2 - 3 * mean_l.pow(4)
745
746                channel_mean[l] = mean_l
747                channel_var[l] = m2
748                channel_skew[l] = m3 / m2.pow(1.5)
749                channel_kurt[l] = m4 / m2.square()
750
751                var_l = m2.clamp_min(0)
752                s1, s2 = var_l.sum(), var_l.square().sum()
753
754                # Effective channel count: participation ratio of per-channel
755                # variance. No privileged basis -> variance spread ~uniformly
756                # across channels -> this saturates near hidden_size;
757                # concentration in a few channels pulls it down. Threshold-free
758                # counterpart to "how many outlier channels are there".
759                if s2 > 0:
760                    eff_channel_count[l] = (s1.square() / s2).item()
761
762                # Gini coefficient of per-channel variance (0 = perfectly even
763                # across channels, ~1 = all variance in one channel).
764                sorted_var, _ = torch.sort(var_l)
765                total = sorted_var.sum().clamp_min(1e-24)
766                channel_gini[l] = ((2 * (rank_1_to_n * sorted_var).sum()) / (n * total) - (n + 1) / n).item()
767
768                # Power-law decay exponent of the sorted (descending) variance
769                # curve: fit slope of log(variance) vs log(rank). A shape
770                # descriptor, not a cutoff -- two layers can share a Gini value
771                # while one is a clean power law and the other a step function.
772                desc_var, _ = torch.sort(var_l, descending=True)
773                floor = torch.clamp_min(desc_var.max() * 1e-8, 1e-24)
774                log_v = torch.log(desc_var.clamp_min(floor))
775                log_v_centered = log_v - log_v.mean()
776                slope = (log_rank_desc_centered * log_v_centered).sum() / log_rank_desc_centered.square().sum()
777                channel_powerlaw_exponent[l] = (-slope).item()
778
779                # Persistence: Spearman rank correlation of per-channel variance
780                # between two disjoint halves of the token stream. High ->
781                # same channels dominate dataset-wide (a genuinely fixed
782                # privileged set); low -> the "privileged" channels drift
783                # between batches and the concentration seen above may just be
784                # sampling noise rather than a stable structural feature.
785                na, nb = chan_count_a[l].item(), chan_count_b[l].item()
786                if na > 0 and nb > 0:
787                    mean_a = chan_sum_x_a[l] / na
788                    var_a = (chan_sum_x2_a[l] / na - mean_a.square()).clamp_min(0)
789                    mean_b = chan_sum_x_b[l] / nb
790                    var_b = (chan_sum_x2_b[l] / nb - mean_b.square()).clamp_min(0)
791                    channel_persistence[l] = _spearman(var_a, var_b).item()
792
793            results["pb_effective_channel_count_per_layer"] = eff_channel_count.tolist()
794            results["pb_channel_gini_per_layer"] = channel_gini.tolist()
795            results["pb_channel_powerlaw_exponent_per_layer"] = channel_powerlaw_exponent.tolist()
796            results["pb_channel_persistence_per_layer"] = channel_persistence.tolist()
797
798            # Raw per-channel moment vectors, for on-demand visual inspection
799            # (see plot_channel_concentration) -- not meant to be read as flat
800            # per-layer scalars like the rest of `results`.
801            self.channel_moments = {
802                "mean": channel_mean,
803                "variance": channel_var,
804                "skewness": channel_skew,
805                "kurtosis": channel_kurt,
806            }
807
808        return results
809
810# --- On-demand visual inspection --------------------------------------------
811
812def plot_channel_concentration(channel_variance: torch.Tensor, label: str = "layer", ax=None):
813    """Sorted-magnitude curve of per-channel variance for a single layer
814    (e.g. `evaluator.channel_moments["variance"][layer_idx]`), on log-log axes.
815
816    This is the reality check for the channel-stats summary scalars above: a
817    clean straight line means the power-law exponent is actually describing
818    the shape (a genuine power law); a flat-then-cliff curve means a small,
819    sharply-separated outlier set (matches the effective-channel-count/Gini
820    story cleanly); anything else -- e.g. a single wildly dominant point, or a
821    long plateau with no clear knee -- is a sign the summary scalars for that
822    layer might be misleading (or at least need a specific story), and is
823    worth a look before trusting them. Cheap enough to call routinely on
824    layers that look unremarkable, not just ones that look suspicious.
825
826    Not wired into the evaluator: call it yourself, on whichever layer(s) or
827    checkpoints you want to inspect. Pass the same `ax` across calls to
828    overlay multiple curves (e.g. across layers, or the same layer across
829    training checkpoints) for comparison.
830    """
831    import matplotlib.pyplot as plt
832
833    if ax is None:
834        _, ax = plt.subplots(figsize=(6, 4))
835
836    sorted_desc, _ = torch.sort(channel_variance.detach().float().clamp_min(1e-12), descending=True)
837    rank = torch.arange(1, sorted_desc.shape[0] + 1)
838
839    ax.plot(rank.numpy(), sorted_desc.cpu().numpy(), marker="o", markersize=2, linewidth=1, label=label)
840    ax.set_xscale("log")
841    ax.set_yscale("log")
842    ax.set_xlabel("channel rank (descending variance)")
843    ax.set_ylabel("variance")
844    ax.legend()
845    return ax