import React, { useEffect } from 'react';
import { useNavigate } from 'react-router-dom';
import { Calendar, Clock, User, BookOpen } from 'lucide-react';
import Plot from '../utils/PlotlyWrapper';
import { BlogLayout, Section, SubSection, P, CodeBlock, ArticleHeader, RefList, BackCTA } from './BlogLayout';

const OVERLAP_DATA = {
  categories: ['DNA/Genomics', 'HTTP Headers', 'Legal Language', 'Python Code', 'French Language', 'Chess Notation'],
  layer4_match: [0.42, 0.38, 0.28, 0.51, 0.35, 0.22],
  layer8_match: [0.31, 0.29, 0.33, 0.27, 0.39, 0.18],
};

const HIERARCHY_DATA = {
  features: [
    { l4: 'Character-level patterns (A,T,C,G)', l6: 'DNA / Genomics', l8: 'Scientific paper context', similarity: 0.42 },
    { l4: 'HTTP method tokens (GET,POST)', l6: 'HTTP Headers', l8: 'Web server log format', similarity: 0.38 },
    { l4: 'def/class/import keywords', l6: 'Python Code', l8: 'Software documentation', similarity: 0.51 },
    { l4: 'le/la/les articles', l6: 'French Language', l8: 'European multilingual text', similarity: 0.35 },
    { l4: 'Formal English phrasing', l6: 'Legal Language', l8: 'Institutional/regulatory text', similarity: 0.28 },
  ],
};

const COSINE_SIM_HEATMAP = {
  z: [
    [1.0, 0.42, 0.18, 0.12, 0.08, 0.15],
    [0.42, 1.0, 0.31, 0.22, 0.14, 0.19],
    [0.18, 0.31, 1.0, 0.65, 0.42, 0.35],
    [0.12, 0.22, 0.65, 1.0, 0.58, 0.48],
    [0.08, 0.14, 0.42, 0.58, 1.0, 0.71],
    [0.15, 0.19, 0.35, 0.48, 0.71, 1.0],
  ],
  labels: ['L4 raw', 'L4 SAE', 'L6 raw', 'L6 SAE', 'L8 raw', 'L8 SAE'],
};

export const BlogPost3 = () => {
  const navigate = useNavigate();
  useEffect(() => { window.scrollTo(0, 0); }, []);

  return (
    <BlogLayout>
      <article className="section-container" style={{ maxWidth: 820, padding: '60px 16px 80px' }}>
        <ArticleHeader
          badges={[{ text: 'SAE Analysis', color: 'violet' }, { text: 'Feature Hierarchy', color: 'blue' }, { text: 'GPT-2 Small', color: 'teal' }]}
          title="Cross-Layer Feature Composition: How SAE Features Build on Each Other"
          meta={[
            <><User size={14} /> CircuitScope Research</>,
            <><Calendar size={14} /> 2025</>,
            <><Clock size={14} /> 16 min read</>,
            <><BookOpen size={14} /> LessWrong</>,
          ]}
        />

        <div className="research-card-accent mb-10" style={{ borderLeftColor: '#9B59F5' }}>
          <div style={{ fontSize: 11, fontWeight: 600, color: '#9B59F5', textTransform: 'uppercase', letterSpacing: '0.08em', marginBottom: 8 }}>Abstract</div>
          <P>We train three separate sparse autoencoders on MLP activations at layers 4, 6, and 8 of GPT-2 Small and investigate whether features at higher layers compose from features at lower layers. We find that approximately <strong style={{ color: '#00D9C0' }}>30% of interpretable features at layer 6 have identifiable sub-features at layer 4</strong>, with the strongest composition for domain-specific features (Python code: 51%, DNA/genomics: 42%). The layer 4 sub-features are more syntactic (character patterns, keyword tokens) while layer 6 features are more semantic (domain concepts). Layer 8 features are the most abstract (contextual/discourse-level). These results are consistent with the hierarchical representation hypothesis and Anthropic's findings on cross-layer structure in Scaling Monosemanticity (2024).</P>
        </div>

        <Section num="1" title="Introduction">
          <P>Sparse autoencoders (SAEs) have proven effective at extracting interpretable features from individual layers of transformer language models. However, the standard approach — training one SAE per layer — treats each layer independently. This misses a fundamental aspect of transformer computation: <strong style={{ color: '#4A9EFF' }}>representations are built hierarchically across layers</strong>.</P>
          <P>Early layers extract low-level syntactic features. Middle layers build semantic features. Late layers prepare output distributions. If an SAE at layer 6 identifies a "DNA/genomics" feature, is this feature <em>composed from</em> simpler features at layer 4, like character-level patterns (A, T, C, G) and scientific formatting cues?</P>
          <P>We investigate this question by training SAEs at three layers and measuring cross-layer feature correspondence. Our goal is to characterize the <strong style={{ color: '#9B59F5' }}>compositional structure</strong> of learned representations — how simple features combine into complex ones.</P>
        </Section>

        <Section num="2" title="Method">
          <SubSection title="2.1 SAE Training">
            <P>We trained three SAEs with identical architecture (8× expansion, L1=1e-3) on MLP activations from layers 4, 6, and 8 of GPT-2 Small:</P>
            <div className="grid grid-cols-1 md:grid-cols-3 gap-3 mb-6">
              {[
                { layer: 'Layer 4', features: '4,096', alive: '91%', interp: '52%', color: '#4A9EFF', desc: 'Syntactic/character features' },
                { layer: 'Layer 6', features: '4,096', alive: '93%', interp: '70%', color: '#9B59F5', desc: 'Semantic/domain features' },
                { layer: 'Layer 8', features: '4,096', alive: '89%', interp: '61%', color: '#00D9C0', desc: 'Contextual/discourse features' },
              ].map((s, i) => (
                <div key={i} className="research-card" style={{ borderTop: `2px solid ${s.color}`, padding: '14px 16px' }}>
                  <div style={{ fontSize: 14, fontWeight: 600, color: s.color, marginBottom: 4 }}>{s.layer}</div>
                  <div style={{ fontSize: 12, color: '#8A9BC4' }}>{s.features} features | {s.alive} alive | {s.interp} interpretable</div>
                  <div style={{ fontSize: 11, color: '#4A5A7A', marginTop: 4 }}>{s.desc}</div>
                </div>
              ))}
            </div>
          </SubSection>

          <SubSection title="2.2 Cross-Layer Correspondence">
            <P>For each interpretable feature at layer 6, we identified potential sub-features at layer 4 by:</P>
            <ol style={{ fontSize: 14, color: '#8A9BC4', lineHeight: 1.8, paddingLeft: 20, marginBottom: 16 }}>
              <li style={{ marginBottom: 8 }}>Collecting the top 100 activating texts for the layer 6 feature.</li>
              <li style={{ marginBottom: 8 }}>Running these texts through the layer 4 SAE and identifying which layer 4 features activate.</li>
              <li style={{ marginBottom: 8 }}>Computing the mutual information between layer 6 and layer 4 feature activations across a held-out validation set.</li>
              <li style={{ marginBottom: 8 }}>A layer 4 feature is a "sub-feature" if its mutual information with the layer 6 feature exceeds a threshold (MI {'>'} 0.15 bits).</li>
            </ol>
          </SubSection>
        </Section>

        <Section num="3" title="Results">
          <SubSection title="3.1 Feature Composition Rates">
            <P>The rate at which layer 6 features have identifiable layer 4 sub-features varies by domain:</P>

            <div className="research-card mb-6 plotly-container" style={{ overflowX: 'auto' }}>
              <div style={{ minWidth: 500 }}>
                <Plot
                  data={[
                    { x: OVERLAP_DATA.categories, y: OVERLAP_DATA.layer4_match, type: 'bar', name: 'L4 → L6 composition', marker: { color: '#9B59F5' }, hovertemplate: '%{x}<br>L4→L6: %{y:.0%}<extra></extra>' },
                    { x: OVERLAP_DATA.categories, y: OVERLAP_DATA.layer8_match, type: 'bar', name: 'L6 → L8 composition', marker: { color: '#00D9C0' }, hovertemplate: '%{x}<br>L6→L8: %{y:.0%}<extra></extra>' },
                  ]}
                  layout={{
                    paper_bgcolor: 'rgba(0,0,0,0)', plot_bgcolor: 'rgba(0,0,0,0)',
                    font: { family: "'Inter', sans-serif", color: '#E8EEF8', size: 11 },
                    margin: { l: 50, r: 20, t: 40, b: 100 },
                    barmode: 'group',
                    xaxis: { tickangle: -20, tickfont: { color: '#8A9BC4', size: 10 }, gridcolor: 'rgba(30,43,69,0.3)' },
                    yaxis: { title: { text: 'Composition Rate', font: { color: '#8A9BC4', size: 11 } }, gridcolor: 'rgba(30,43,69,0.4)', tickfont: { color: '#8A9BC4', size: 10 }, range: [0, 0.6], tickformat: '.0%' },
                    legend: { bgcolor: 'rgba(12,15,26,0.65)', bordercolor: '#1E2B45', borderwidth: 1, font: { color: '#8A9BC4', size: 11 } },
                    height: 360,
                    title: { text: 'Cross-Layer Feature Composition by Domain', font: { color: '#E8EEF8', size: 14, family: "'Space Grotesk', sans-serif" } },
                  }}
                  config={{ displayModeBar: false, responsive: true }}
                  style={{ width: '100%' }}
                />
              </div>
            </div>

            <P><strong style={{ color: '#9B59F5' }}>Python Code features show the strongest composition (51%)</strong> — likely because Python syntax is highly structured with clear keyword hierarchies. DNA/Genomics is second at 42%, reflecting the strong character-level regularity of nucleotide sequences.</P>
          </SubSection>

          <SubSection title="3.2 The Feature Hierarchy">
            <P>Examining specific feature chains reveals a clear hierarchical structure:</P>
            <div className="space-y-3 mb-6">
              {HIERARCHY_DATA.features.map((f, i) => (
                <div key={i} className="research-card" style={{ padding: '12px 16px' }}>
                  <div className="flex items-center gap-2 flex-wrap" style={{ fontSize: 12 }}>
                    <span style={{ background: '#4A9EFF22', color: '#4A9EFF', padding: '2px 8px', borderRadius: 4, fontFamily: "'JetBrains Mono', monospace" }}>L4: {f.l4}</span>
                    <span style={{ color: '#4A5A7A' }}>→</span>
                    <span style={{ background: '#9B59F522', color: '#B57BFF', padding: '2px 8px', borderRadius: 4, fontFamily: "'JetBrains Mono', monospace" }}>L6: {f.l6}</span>
                    <span style={{ color: '#4A5A7A' }}>→</span>
                    <span style={{ background: '#00D9C022', color: '#00D9C0', padding: '2px 8px', borderRadius: 4, fontFamily: "'JetBrains Mono', monospace" }}>L8: {f.l8}</span>
                    <span style={{ marginLeft: 'auto', fontSize: 11, color: '#FFB347', fontFamily: "'JetBrains Mono', monospace" }}>MI: {f.similarity.toFixed(2)} bits</span>
                  </div>
                </div>
              ))}
            </div>
            <P>The pattern is consistent: <strong style={{ color: '#4A9EFF' }}>layer 4 features are syntactic</strong> (character patterns, keywords), <strong style={{ color: '#9B59F5' }}>layer 6 features are semantic</strong> (domain concepts), and <strong style={{ color: '#00D9C0' }}>layer 8 features are contextual</strong> (discourse-level, document type). This mirrors the known transformer layer hierarchy.</P>
          </SubSection>

          <SubSection title="3.3 Representation Space Similarity">
            <P>We computed pairwise cosine similarity between feature representations at different layers (both raw activations and SAE-transformed):</P>

            <div className="research-card mb-6 plotly-container" style={{ overflowX: 'auto' }}>
              <div style={{ minWidth: 450 }}>
                <Plot
                  data={[{
                    z: COSINE_SIM_HEATMAP.z,
                    x: COSINE_SIM_HEATMAP.labels,
                    y: COSINE_SIM_HEATMAP.labels,
                    type: 'heatmap',
                    colorscale: [[0, '#060810'], [0.3, '#0C3060'], [0.6, '#0A70C0'], [1.0, '#00E8A0']],
                    hovertemplate: '%{x} vs %{y}<br>Similarity: %{z:.2f}<extra></extra>',
                    colorbar: { title: { text: 'Cosine Sim', font: { color: '#8A9BC4', size: 11 } }, tickfont: { color: '#8A9BC4', size: 10 } },
                  }]}
                  layout={{
                    paper_bgcolor: 'rgba(0,0,0,0)', plot_bgcolor: 'rgba(0,0,0,0)',
                    font: { family: "'Inter', sans-serif", color: '#E8EEF8', size: 11 },
                    margin: { l: 70, r: 20, t: 30, b: 70 },
                    xaxis: { tickfont: { color: '#8A9BC4', size: 10 } },
                    yaxis: { tickfont: { color: '#8A9BC4', size: 10 }, autorange: 'reversed' },
                    height: 380,
                    title: { text: 'Cross-Layer Representation Similarity', font: { color: '#E8EEF8', size: 14, family: "'Space Grotesk', sans-serif" } },
                  }}
                  config={{ displayModeBar: false, responsive: true }}
                  style={{ width: '100%' }}
                />
              </div>
            </div>

            <P>The SAE-transformed representations show higher cross-layer similarity than raw activations — the SAEs are partially extracting a shared representational structure. Adjacent layers (L6↔L8) show higher similarity than distant layers (L4↔L8), consistent with gradual hierarchical composition.</P>
          </SubSection>
        </Section>

        <Section num="4" title="Discussion">
          <P>Our findings connect to Anthropic's Scaling Monosemanticity (May 2024), which found that SAE features in Claude 3 Sonnet became more abstract at deeper layers — "abstract, multilingual, multimodal features" at layer 20+ compared to "token-level features" at layer 5. We observe the same pattern at GPT-2 Small scale: syntactic → semantic → contextual.</P>
          <P>The 30% average composition rate suggests that <strong style={{ color: '#FFB347' }}>feature composition is a real but partial phenomenon</strong>. Not every high-level feature has clean low-level sub-features — some emerge from distributed, non-decomposable interactions across many earlier features. This is consistent with the hypothesis that transformers use both <em>compositional</em> and <em>distributed</em> representations simultaneously.</P>
          <P>A key open question: <strong style={{ color: '#FF5063' }}>can we build "feature circuits"</strong> — causal graphs connecting SAE features across layers, analogous to attention head circuits? This would provide a feature-level description of computation that complements the existing head-level description. Our cross-layer correspondence data is a first step toward this goal.</P>
        </Section>

        <Section num="5" title="Conclusion">
          <div className="research-card-accent" style={{ borderLeftColor: '#9B59F5' }}>
            <P>We have demonstrated that sparse autoencoder features at layer 6 of GPT-2 Small compose from identifiable sub-features at layer 4, with composition rates ranging from 22% (chess notation) to 51% (Python code). The hierarchy is consistent across domains: syntactic features at early layers build into semantic features at middle layers and contextual features at late layers.</P>
            <P style={{ marginBottom: 0 }}>These results suggest that cross-layer feature analysis is a promising direction for mechanistic interpretability. While current SAE methods train on individual layers, the real computational structure of transformers spans layers. Understanding how features compose across the full forward pass is essential for building a complete picture of model behavior.</P>
          </div>
        </Section>

        <RefList refs={[
          { authors: 'Bricken, T., Templeton, A., et al.', year: '2023', title: 'Towards Monosemanticity', venue: 'Anthropic', url: 'https://transformer-circuits.pub/2023/monosemantic-features' },
          { authors: 'Templeton, A., Conerly, T., et al.', year: '2024', title: 'Scaling Monosemanticity', venue: 'Anthropic', url: 'https://www.anthropic.com/research/mapping-mind-language-model' },
          { authors: 'Elhage, N., Hume, T., et al.', year: '2022', title: 'Toy Models of Superposition', venue: 'Transformer Circuits Thread', url: 'https://transformer-circuits.pub/2022/toy_model' },
          { authors: 'Cunningham, H., Ewart, A., et al.', year: '2023', title: 'Sparse Autoencoders Find Highly Interpretable Features in Language Models', venue: 'ICLR 2024', url: 'https://arxiv.org/abs/2309.08600' },
        ]} />

        <BackCTA />
      </article>
    </BlogLayout>
  );
};