LEM/pkg/lem/attention.go

// Q/K Bone Orientation analysis engine.
//
// Computes attention coherence metrics from KV cache snapshots.
// Pure Go CPU math — no GPU, no CGO dependencies.
package lem

import (
	"math"

	"forge.lthn.ai/core/go-inference"
)

// BOResult holds Q/K Bone Orientation metrics for a single inference.
type BOResult struct {
	MeanCoherence       float64   `json:"mean_coherence"`        // Mean pairwise head coherence (0-1)
	MeanCrossAlignment  float64   `json:"mean_cross_alignment"`  // Mean adjacent-layer alignment (0-1)
	MeanHeadEntropy     float64   `json:"mean_head_entropy"`     // Mean attention entropy per head (0-1)
	PhaseLockScore      float64   `json:"phase_lock_score"`      // Fraction of head pairs above coherence threshold
	JointCollapseCount  int       `json:"joint_collapse_count"`  // Layers where cross-alignment drops below threshold
	LayerCoherence      []float64 `json:"layer_coherence"`       // Per-layer mean head coherence
	LayerCrossAlignment []float64 `json:"layer_cross_alignment"` // Per-layer cross-alignment (len = layers-1)
}

// Composite returns a 0-100 score from BO metrics.
func (r *BOResult) Composite() float64 {
	score := (0.30*r.MeanCoherence +
		0.25*r.MeanCrossAlignment +
		0.20*r.PhaseLockScore +
		0.15*r.MeanHeadEntropy +
		0.10*math.Max(0, 1.0-float64(r.JointCollapseCount)*0.2)) * 100.0
	return min(100, max(0, score))
}

const (
	coherenceThreshold = 0.7 // Minimum cosine sim for "phase-locked" head pair
	collapseThreshold  = 0.5 // Below this cross-alignment = joint collapse
)

// AnalyseAttention computes Q/K Bone Orientation metrics from a KV cache snapshot.
func AnalyseAttention(snap *inference.AttentionSnapshot) *BOResult {
	if snap == nil || len(snap.Keys) == 0 {
		return &BOResult{}
	}

	result := &BOResult{
		LayerCoherence:      make([]float64, snap.NumLayers),
		LayerCrossAlignment: make([]float64, max(0, snap.NumLayers-1)),
	}

	var totalCoherence, totalEntropy float64
	var totalPairsLocked, totalPairs int
	layerMeans := make([][]float32, snap.NumLayers) // mean K vector per layer

	for layer := 0; layer < snap.NumLayers; layer++ {
		if layer >= len(snap.Keys) || snap.Keys[layer] == nil {
			continue
		}
		heads := snap.Keys[layer]
		nHeads := len(heads)

		// Compute mean K vector for this layer (average over heads).
		layerMeans[layer] = meanVector(heads)

		// Pairwise head coherence within layer.
		var layerCoh float64
		var pairs int
		for i := 0; i < nHeads; i++ {
			for j := i + 1; j < nHeads; j++ {
				sim := cosineSim32(heads[i], heads[j])
				layerCoh += sim
				pairs++
				if sim >= coherenceThreshold {
					totalPairsLocked++
				}
				totalPairs++
			}
		}
		if pairs > 0 {
			layerCoh /= float64(pairs)
		}
		result.LayerCoherence[layer] = layerCoh
		totalCoherence += layerCoh

		// Per-head entropy (magnitude distribution across positions).
		for _, head := range heads {
			totalEntropy += headEntropy(head, snap.SeqLen, snap.HeadDim)
		}
	}

	// Cross-layer alignment.
	var totalCross float64
	for i := 0; i < snap.NumLayers-1; i++ {
		if layerMeans[i] == nil || layerMeans[i+1] == nil {
			continue
		}
		alignment := cosineSim32(layerMeans[i], layerMeans[i+1])
		result.LayerCrossAlignment[i] = alignment
		totalCross += alignment
		if alignment < collapseThreshold {
			result.JointCollapseCount++
		}
	}

	if snap.NumLayers > 0 {
		result.MeanCoherence = totalCoherence / float64(snap.NumLayers)
	}
	if snap.NumLayers > 1 {
		result.MeanCrossAlignment = totalCross / float64(snap.NumLayers-1)
	}
	totalHeads := snap.NumLayers * snap.NumHeads
	if totalHeads > 0 {
		result.MeanHeadEntropy = totalEntropy / float64(totalHeads)
	}
	if totalPairs > 0 {
		result.PhaseLockScore = float64(totalPairsLocked) / float64(totalPairs)
	}

	return result
}

// cosineSim32 computes cosine similarity between two float32 slices.
func cosineSim32(a, b []float32) float64 {
	if len(a) != len(b) || len(a) == 0 {
		return 0
	}
	var dot, normA, normB float64
	for i := range a {
		ai, bi := float64(a[i]), float64(b[i])
		dot += ai * bi
		normA += ai * ai
		normB += bi * bi
	}
	denom := math.Sqrt(normA) * math.Sqrt(normB)
	if denom == 0 {
		return 0
	}
	return dot / denom
}

// meanVector computes element-wise mean across multiple float32 slices.
func meanVector(vecs [][]float32) []float32 {
	if len(vecs) == 0 {
		return nil
	}
	n := len(vecs[0])
	mean := make([]float32, n)
	for _, v := range vecs {
		for i := range v {
			if i < n {
				mean[i] += v[i]
			}
		}
	}
	scale := float32(len(vecs))
	for i := range mean {
		mean[i] /= scale
	}
	return mean
}

// headEntropy computes normalised Shannon entropy of K vector magnitudes
// across sequence positions for a single head.
func headEntropy(head []float32, seqLen, headDim int) float64 {
	if seqLen == 0 || headDim == 0 {
		return 0
	}
	// Compute magnitude per position.
	mags := make([]float64, seqLen)
	var total float64
	for pos := 0; pos < seqLen; pos++ {
		var sum float64
		start := pos * headDim
		for d := 0; d < headDim && start+d < len(head); d++ {
			v := float64(head[start+d])
			sum += v * v
		}
		mags[pos] = math.Sqrt(sum)
		total += mags[pos]
	}
	if total == 0 {
		return 0
	}
	// Normalised Shannon entropy.
	var entropy float64
	for _, m := range mags {
		p := m / total
		if p > 0 {
			entropy -= p * math.Log2(p)
		}
	}
	maxEntropy := math.Log2(float64(seqLen))
	if maxEntropy == 0 {
		return 0
	}
	return entropy / maxEntropy
}

// AttentionFeatures returns a 5D feature vector from BO metrics.
func AttentionFeatures(ar *BOResult) []float64 {
	if ar == nil {
		return make([]float64, 5)
	}
	return []float64{
		ar.MeanCoherence,
		ar.MeanCrossAlignment,
		ar.MeanHeadEntropy,
		ar.PhaseLockScore,
		math.Max(0, 1.0-float64(ar.JointCollapseCount)*0.2),
	}
}

// AttentionFeatureLabels returns the labels for the attention feature vector.
func AttentionFeatureLabels() []string {
	return []string{
		"mean_coherence",
		"cross_alignment",
		"head_entropy",
		"phase_lock",
		"joint_stability",
	}
}