4️⃣ Training a Probe

class LinearProbeTrainer:
    def __init__(self, model: HookedTransformer, args: ProbeTrainingArgs):
        self.model = model
        self.args = args
        self.linear_probe = args.setup_linear_probe(model)
def training_step(self, indices: Int[Tensor, "n_games"]) -> Float[Tensor, ""]:
        # Use indices to slice our batch of games (remember, games_id = token IDs
        # from 1 to 60, and games_square = indices of squares in board)
        indices_cpu = indices.cpu()
        games_id = board_seqs_id[indices_cpu]  # shape [batch n_moves=60]
        games_square = board_seqs_square[indices_cpu]  # shape [batch n_moves=60]
# Define seqpos slicing (note, we add n_ctx to pos_end to deal with the zero case)
        pos_start = self.args.pos_start
        pos_end = self.args.pos_end + self.model.cfg.n_ctx
# Cache resid_post from all our games (ignoring the last one)
        with t.inference_mode():
            _, cache = model.run_with_cache(
                games_id[:, :-1].to(device),
                return_type=None,
                names_filter=lambda name: name.endswith("resid_post"),
            )
# We slice from the first even index (on or after pos_start), since we're
        # just looking at predictions made after white has played a move.
        pos_start_even = pos_start + (pos_start % 2)
        seqpos_indices = np.arange(pos_start_even, pos_end, 2)
        resid_post = cache["resid_post", self.args.layer][:, seqpos_indices]
# Get probe output, i.e. probe_logits[g, p, r, c] = the 3-vector of logit
        # predictions for what color is in square [r, c] AFTER the p-th move is
        # played in game g.
        probe_logits = einops.einsum(
            resid_post,
            self.linear_probe,
            "batch pos d_model, d_model rows cols options -> batch pos rows cols options",
        )
        probe_logprobs = probe_logits.log_softmax(-1)
# Get the actual game state. The original state has {0: empty, 1: black, -1: white} and
        # we want our probe to be in the basis {0: empty, 1: theirs, 2: mine}. We're only training
        # on even moves i.e. black just played and mine = white, so we just need to map -1 -> 2.
        state = get_board_states_and_legal_moves(games_square)[0]  # shape [batch moves 8 8]
        state = state[:, seqpos_indices]  # shape [batch pos 8 8]
        state[state == -1] = 2
# Index into probe logprobs to get the logprobs for correct board state, and then
        # return loss as the mean over games & posns, and sum over rows & cols (since each
        # row & col is effectively an independent probe).
        correct_probe_logprobs = eindex(probe_logprobs, state, "game pos row col [game pos row col]")
        loss = -einops.reduce(correct_probe_logprobs, "game pos row col -> row col", "mean").sum()
if self.args.use_wandb:
            wandb.log(dict(loss=loss.item()), step=self.step)
        self.step += 1
return loss
def shuffle_training_indices(self):
        """
        Returns the tensors you'll use to index into the training data.
        """
        n_indices = self.args.num_games - (self.args.num_games % self.args.batch_size)
        full_train_indices = t.randperm(self.args.num_games)[:n_indices]
        full_train_indices = einops.rearrange(
            full_train_indices,
            "(batch_idx game_idx) -> batch_idx game_idx",
            game_idx=self.args.batch_size,
        )
        return full_train_indices
def train(self):
        self.step = 0
        if self.args.use_wandb:
            wandb.init(project=self.args.wandb_project, name=self.args.wandb_name, config=self.args)
optimizer = t.optim.AdamW(
            [self.linear_probe],
            lr=self.args.lr,
            betas=self.args.betas,
            weight_decay=self.args.weight_decay,
        )
for epoch in range(self.args.epochs):
            print(f"Epoch {epoch + 1}/{self.args.epochs}")
            full_train_indices = self.shuffle_training_indices()
            progress_bar = tqdm(full_train_indices)
            for indices in progress_bar:
                loss = self.training_step(indices)
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                progress_bar.set_description(f"Loss = {loss:.4f}")
if self.args.use_wandb:
            wandb.finish()

@dataclass
class MultiProbeTrainingArgs(ProbeTrainingArgs):
    modes: int = 3  # even, odd, both (i.e. the data we train on)
def setup_linear_probe(self, model: HookedTransformer):
        linear_probe = t.randn(
            self.modes,
            model.cfg.d_model,
            self.rows,
            self.cols,
            self.options,
            device=device,
        ) / np.sqrt(model.cfg.d_model)
        linear_probe.requires_grad = True
        return linear_probe
class LinearMultiProbeTrainer(LinearProbeTrainer):
    def training_step(self, indices: Int[Tensor, "n_games"]) -> Float[Tensor, ""]:
        indices_cpu = indices.cpu()
        games_id = board_seqs_id[indices_cpu]  # shape [batch n_moves=60]
        games_square = board_seqs_square[indices_cpu]  # shape [batch n_moves=60]
pos_start = self.args.pos_start
        pos_end = self.args.pos_end + self.model.cfg.n_ctx
# Cache resid_post from all our games (ignoring the last one)
        with t.inference_mode():
            _, cache = model.run_with_cache(
                games_id[:, :-1].to(device),
                return_type=None,
                names_filter=lambda name: name.endswith("resid_post"),
            )
# We're training on all modes, so we slice all resid values in our range.
        resid_post = cache["resid_post", self.args.layer][:, pos_start:pos_end]
# Get probe output, i.e. probe_logits[m, g, p, r, c] = the 3-vector of logit predictions from
        # mode-m probe, for what color is in square [r, c] AFTER the p-th move is played in game g.
        probe_logits = einops.einsum(
            resid_post,
            self.linear_probe,
            "game pos d_model, mode d_model rows cols options -> mode game pos rows cols options",
        )
        probe_logprobs = probe_logits.log_softmax(-1)
# Get the actual game state. The original state has {0: empty, 1: black, -1: white} and
        # we want our probe to be in the basis {0: empty, 1: theirs, 2: mine}. For even moves,
        # mine = white, so we map -1 -> 2. For odd moves, mine = black, so we map {1, -1} -> {2, 1}.
        state = get_board_states_and_legal_moves(games_square)[0]
        state[:, ::2][state[:, ::2] == -1] = 2
        state[:, 1::2][state[:, 1::2] == 1] = 2
        state[:, 1::2][state[:, 1::2] == -1] = 1
        state = state[:, pos_start:pos_end]
# Index into the probe logprobs with the correct indices (note, each of our 3 probe modes
        # gives us a different tensor of logprobs).
        correct_probe_logprobs = eindex(
            probe_logprobs,
            state,
            "mode game pos row col [game pos row col]",  # -> shape [mode game pos row col]
        )
        # Get the logprobs we'll be using for our 3 different probes. Remember that for the even
        # and odd probes we need to take only the even/odd moves respectively (and also that we've
        # already sliced logprobs from pos_start: pos_end).
        pos_start_even, pos_start_odd = (0, 1) if pos_start % 2 == 0 else (1, 0)
        even_probe_logprobs = correct_probe_logprobs[0, pos_start_even::2]
        odd_probe_logprobs = correct_probe_logprobs[1, pos_start_odd::2]
        both_probe_logprobs = correct_probe_logprobs[2]
        # Get our 3 different loss functions
        loss_even = -einops.reduce(even_probe_logprobs, "game pos row col -> row col", "mean").sum()
        loss_odd = -einops.reduce(odd_probe_logprobs, "game pos row col -> row col", "mean").sum()
        loss_both = -einops.reduce(both_probe_logprobs, "game pos row col -> row col", "mean").sum()
        # We backprop on the sum of all 3 losses
        loss = loss_even + loss_odd + loss_both
if self.args.use_wandb:
            wandb.log(
                dict(
                    loss=loss.item(),
                    loss_even=loss_even.item(),
                    loss_odd=loss_odd.item(),
                    loss_both=loss_both.item(),
                ),
                step=self.step,
            )
        self.step += 1
return loss