Probing code

[orgadhadas]

Hi,

Thank you for this interesting work.

Can you share the code for your probing experiment?
Specifically, what are the labels that the probe is trained on? And how did you convert the probe's performance to figure 4?

[da03]

Thanks for your interest! I haven't cleaned up the code, but a messy version can be found in branch interpret: https://github.com/da03/Internalize_CoT_Step_by_Step/blob/interpret/train.sh. The labels are CoT tokens. Figure 4 is drawn using tikz after applying a threshold of 0.95:

\begin{figure*}[!t]
    \centering
    \resizebox{0.8\textwidth}{!}{
  \begin{tikzpicture}[node distance=0.1cm, auto]

% Define colors
\definecolor{digitcolor}{RGB}{70,130,180} % SteelBlue
\definecolor{boxcolor}{RGB}{205,92,92} % IndianRed
\definecolor{connectorcolor}{RGB}{34,139,34} % ForestGreen
\tikzstyle{digitnode} = [draw, circle, fill=digitcolor!20, minimum size=0.7cm, inner sep=0pt]

% First row: Input
\node (d0) at (0,0) {Input:};

\node[digitnode, circle, fill=digitcolor!20, right=of d0] (d3) {$a_1$};
\node[digitnode, circle, fill=digitcolor!20, right=of d3] (d2) {$a_2$};
\node[digitnode, circle, fill=digitcolor!20, right=of d2] (d1) {$a_3$};
\node[digitnode, circle, fill=digitcolor!20, right=of d1] (mult) {$\times$};
\node[digitnode, circle, fill=digitcolor!20, right=of mult] (d6) {$b_1$};
\node[digitnode, circle, fill=digitcolor!20, right=of d6] (d5) {$b_2$};
\node[digitnode, circle, fill=digitcolor!20, right=of d5] (d4) {$b_3$};

% Second row: CoT (Chain of Thought)
\node [below=1.5cm of d0] (cot_label) {CoT:};

% Nodes for '837'
\node[digitnode, circle, fill=digitcolor!20, below=1.5cm of d3] (i837_1) {$p_{11}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i837_1] (i837_2) {$p_{12}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i837_2] (i837_3) {$p_{13}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i837_3] (i837_4) {$p_{14}$};

% Bracket and label above '837'
\draw[decorate, decoration={brace, amplitude=11pt}, yshift=6pt] (i837_1.north west) -- (i837_4.north east) node[midway, above=6pt] {$p_1 = b_1 \times a_3a_2a_1$};

% Nodes for '516'
\node[digitnode, circle, fill=digitcolor!20, right=0.2cm of i837_4] (i516_1) {$p_{21}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i516_1] (i516_2) {$p_{22}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i516_2] (i516_3) {$p_{23}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i516_3] (i516_4) {$p_{24}$};


% Bracket and label above '516'
\draw[decorate, decoration={brace, amplitude=11pt}, yshift=6pt] (i516_1.north west) -- (i516_4.north east) node[midway, above=6pt] {$p_2 = b_2 \times a_3a_2a_1$};

% Nodes for '8886'
\node[digitnode, circle, fill=digitcolor!20, right=0.2cm of i516_4] (i8886_1) {$s_1$};
\node[digitnode, circle, fill=digitcolor!20, right=of i8886_1] (i8886_2) {$s_2$};
\node[digitnode, circle, fill=digitcolor!20, right=of i8886_2] (i8886_3) {$s_3$};
\node[digitnode, circle, fill=digitcolor!20, right=of i8886_3] (i8886_4) {$s_4$};
\node[digitnode, circle, fill=digitcolor!20, right=of i8886_4] (i8886_5) {$s_5$};


% Label above '8886'
%\node[above=6pt of i8886_2] (p1p2label) {$s=p_1+10p_2$};
\draw[decorate, decoration={brace, amplitude=11pt}, yshift=6pt] (i8886_1.north west) -- (i8886_5.north east) node[midway, above=6pt] {$s=p_1+10p_2$};

% Nodes for '294'
\node[digitnode, circle, fill=digitcolor!20, right=0.2cm of i8886_5] (i294_1) {$p_{31}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i294_1] (i294_2) {$p_{32}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i294_2] (i294_3) {$p_{33}$};
\node[digitnode, circle, fill=digitcolor!20, right=of i294_3] (i294_4) {$p_{34}$};

\draw[decorate, decoration={brace, amplitude=11pt}, yshift=6pt] (i294_1.north west) -- (i294_4.north east) node[midway, above=6pt] {$p_3=b_3\times a_3a_2a_1$};


% Collect all CoT nodes into an array
\coordinate (cot_nodes_start) at (i837_1.west);
\coordinate (cot_nodes_end) at (i294_4.east);
\def\cotnodes{i837_1,i837_2,i837_3,i837_4,i516_1,i516_2,i516_3,i516_4,i8886_1,i8886_2,i8886_3,i8886_4,i8886_5,i294_1,i294_2,i294_3,i294_4}

% Third row: Model States
\node [below = 3cm of cot_label, text width=1cm,align=center] (model_label) {Model States:};

\node[digitnode, circle, fill=digitcolor!20, below=6.5cm of d3] (t3) {$a_1$};
\node[digitnode, circle, fill=digitcolor!20, right=of t3] (t2) {$a_2$};
\node[digitnode, circle, fill=digitcolor!20, right=of t2] (t1) {$a_3$};
\node[digitnode, circle, fill=digitcolor!20, right=of t1] (tmult) {$\times$};
\node[digitnode, circle, fill=digitcolor!20, right=of tmult] (t6) {$b_1$};
\node[digitnode, circle, fill=digitcolor!20, right=of t6] (t5) {$b_2$};
\node[digitnode, circle, fill=digitcolor!20, right=of t5] (t4) {$b_3$};
\node[digitnode, circle, fill=digitcolor!20, right=of t4] (teq) {=};

% Boxes above each circle
\foreach \i/\n in {1/t3, 2/t2, 3/t1, 4/tmult, 5/t6, 6/t5, 7/t4, 8/teq} {
    \node[draw, rectangle, fill=boxcolor!20, minimum width=0.6cm, minimum height=1.2cm, above=0.5cm of \n] (h\i) {};
}

% Collect all Model State boxes into an array
\def\statenodes{h1,h2,h3,h4,h5,h6,h7,h8}

% Convert the lists into arrays for easier access
\renewcommand{\split}[1]{\expandafter\splitaux#1,\relax}
\newcommand{\splitaux}[1]{%
    \ifx#1\relax\else
        \xdef\listarray{\listarray{#1}}
        \expandafter\splitaux
    \fi
}

% Prepare arrays of nodes
\def\listarray{}
\split{\cotnodes,}
\def\cotnodearray{\listarray}
\def\listarray{}
\split{\statenodes,}
\def\statenodearray{\listarray}

% Define the weight matrix
% We'll store the weights in a macro that maps a pair of indices to a weight
\makeatletter
\newcommand{\weight}[2]{%
    \ifcase#1 % row 0
        \ifcase#2 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\fi% row 0
    \or % row 1
        \ifcase#2 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\fi% row 1
    \or % row 2
        \ifcase#2 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\fi% row 2
    \or % row 3
        \ifcase#2 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\fi% row 3
    \or % row 4
        \ifcase#2 1.00\or 1.00\or 1.00\or 1.00\or 0\or 0\or 0\or 0\or 1.00\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\fi% row 4
    \or % row 5
        \ifcase#2 0\or 0\or 0\or 0\or 1.00\or 1.00\or 0.98\or 1.00\or 0\or 0\or 0\or 0\or 1.00\or 0\or 0\or 0\or 0\fi% row 5
    \or % row 6
        \ifcase#2 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 1.00\or 1.00\or 0.99\or 1.00\fi% row 6
    \or % row 7
        \ifcase#2 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\or 0\fi% row 7
    \fi
}
\makeatother

% Draw connections from every CoT node to every Model State box
% The line widths are determined by the weight macro
\foreach \i [count=\row from 0] in \statenodes {
    \foreach \j [count=\col from 0] in \cotnodes {
        % Get the weight from the weight macro
        \pgfmathsetmacro{\linewidth}{\weight{\row}{\col}}
        % Draw the connection if the weight is greater than zero
        \ifdim\linewidth pt>0pt
            \draw[connectorcolor, line width=\linewidth mm] (\j.south) -- (\i.north);
        \fi
    }
}
\end{tikzpicture}}
\end{figure*}