Add layer_stride parameter for PromptDiff optimization

- Add layer_stride parameter to control which layers are captured
- Default to 1 (all layers) for AttentionExplorer
- PromptDiff can use layer_stride=2 for every other layer
- Reduces matrix count from 20 to 10 for better visualization fit

🤖 Generated with Claude Code (https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

backend/model_service.py

@@ -47,6 +47,7 @@ class GenerationRequest(BaseModel):
     top_p: Optional[float] = None
     extract_traces: bool = True
     sampling_rate: float = 0.005
+    layer_stride: int = 1  # 1 = all layers, 2 = every other layer, etc.
 
 class AblatedGenerationRequest(BaseModel):
     prompt: str
@@ -500,7 +501,8 @@ class ModelManager:
         temperature: float = 0.7,
         top_k: Optional[int] = None,
         top_p: Optional[float] = None,
-        sampling_rate: float = 0.005
+        sampling_rate: float = 0.005,
+        layer_stride: int = 1  # 1 = all layers, 2 = every other layer, etc.
     ) -> Dict[str, Any]:
         """Generate text with trace extraction"""
         if not self.model or not self.tokenizer:
@@ -643,8 +645,8 @@ class ModelManager:
                         # Clear previous partial traces and add complete ones
                         traces = []  # Reset traces to only include complete attention patterns
                         
-                        # Capture ALL layers for complete visualization
-                        for layer_idx in range(num_layers):
+                        # Capture layers based on stride (1 = all, 2 = every other, etc.)
+                        for layer_idx in range(0, num_layers, layer_stride):
                             try:
                                 # Get all token IDs (prompt + generated)
                                 all_token_ids = inputs["input_ids"][0].tolist()
@@ -892,7 +894,8 @@ async def generate(request: GenerationRequest, authenticated: bool = Depends(ver
         temperature=request.temperature,
         top_k=request.top_k,
         top_p=request.top_p,
-        sampling_rate=request.sampling_rate if request.extract_traces else 0
+        sampling_rate=request.sampling_rate if request.extract_traces else 0,
+        layer_stride=request.layer_stride
     )
     return result
 

components/ui/tooltip.tsx

@@ -0,0 +1,30 @@
+"use client"
+
+import * as React from "react"
+import * as TooltipPrimitive from "@radix-ui/react-tooltip"
+
+import { cn } from "@/lib/utils"
+
+const TooltipProvider = TooltipPrimitive.Provider
+
+const Tooltip = TooltipPrimitive.Root
+
+const TooltipTrigger = TooltipPrimitive.Trigger
+
+const TooltipContent = React.forwardRef<
+  React.ElementRef<typeof TooltipPrimitive.Content>,
+  React.ComponentPropsWithoutRef<typeof TooltipPrimitive.Content>
+>(({ className, sideOffset = 4, ...props }, ref) => (
+  <TooltipPrimitive.Content
+    ref={ref}
+    sideOffset={sideOffset}
+    className={cn(
+      "z-50 overflow-hidden rounded-md bg-gray-900 px-3 py-1.5 text-sm text-gray-100 shadow-md animate-in fade-in-0 zoom-in-95 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:zoom-out-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 border border-gray-800",
+      className
+    )}
+    {...props}
+  />
+))
+TooltipContent.displayName = TooltipPrimitive.Content.displayName
+
+export { Tooltip, TooltipTrigger, TooltipContent, TooltipProvider }