Coverage for transformer_lens/tools/analysis/interchange.py: 100%

1"""Interchange interventions on linear subspaces.

3An interchange intervention runs the model on a *base* input while setting the base's activation

4at one hook point, inside a subspace, to its value on a *source* input::

6 h' = h_base + strength * (h_source - h_base) @ B @ B.T

8``B`` is an orthonormal ``[d_model, rank]`` basis, the convention of

9:func:`~transformer_lens.tools.analysis.projection_kernel.orthonormal_subspace`. The base keeps its

10own coordinates outside the subspace. If the model then produces the output a high-level causal

11model predicts for the base with one variable taken from the source, the subspace carries that

12variable at that hook point (Geiger et al., 2021).

14The module provides the intervention, fixed-direction bases to intervene on (the difference of

15class means, a linear probe's weights, a random basis), and a norm-matched random-subspace placebo.

16:func:`neuron_basis` turns a set of MLP neurons, such as the features a sparse probe selects, into a

17basis, and :func:`random_neuron_basis` draws random neurons as its placebo.

18:func:`fit_das` learns the basis instead, by distributed alignment search (DAS; Geiger et al.,

192024): it trains an orthonormal basis so that the interchange produces the high-level model's

20predicted outputs.

22An interchange accuracy is not evidence on its own. DAS fits arbitrary label mappings and finds

23effects in randomly initialised models (Arora et al., 2024), so the module also provides the

24controls to train the same way: :func:`arbitrary_targets` for an arbitrary-label control, and

25:func:`random_init_bridge` or :func:`reinitialized` for an untrained copy of the model. Compare a

26learned subspace with both, and with the placebo and the baselines.

28Works with any model exposing the hook API (``TransformerBridge`` or a ``HookedRootModule``).

29Activations at the hook must be ``[batch, d_model]``, or ``[batch, pos, d_model]`` with one

30position index per example.

31"""

33from __future__ import annotations

35import copy

36import hashlib

37from dataclasses import dataclass

38from typing import Any, Callable, Hashable, Iterator, Optional, Sequence, Union, cast

40import numpy as np

41import torch

42import torch.nn.functional as F

43from jaxtyping import Float, Int

44from torch import nn

45from torch.nn.utils.parametrizations import orthogonal

47from transformer_lens.tools.analysis._model_state import require_eval_mode

48from transformer_lens.tools.analysis.projection_kernel import orthonormal_subspace

50TaskLogits = Callable[[torch.Tensor], Float[torch.Tensor, "batch n_classes"]]

51"""Runs the model on a batch of inputs in one forward pass and returns task logits, e.g.

52``lambda tokens: model(tokens)[:, -1, answer_token_ids]``."""

54Intervention = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]

55"""Maps (base activations, source activations), both ``[batch, d_model]``, to new base activations."""

58def random_basis(d_model: int, rank: int, seed: int) -> Float[torch.Tensor, "d_model rank"]:

59 """A seeded random orthonormal basis: the baseline any aligned basis should beat.

61 Args:

62 seed: Seed for the Gaussian matrix the basis is drawn from.

64 Returns:

65 Orthonormal basis, ``[d_model, rank]``.

66 """

67 generator = torch.Generator().manual_seed(seed)

68 return orthonormal_subspace(torch.randn(d_model, rank, generator=generator), rank=rank).basis

71def mean_difference_basis(

72 activations: Float[torch.Tensor, "batch d_model"], labels: Sequence[Hashable]

73) -> Float[torch.Tensor, "d_model rank"]:

74 """The span of the centred class means: for two classes, the difference-in-means direction.

76 Args:

77 activations: Activations, one row per example.

78 labels: The variable's value for each example; at least two distinct values.

80 Returns:

81 Orthonormal basis of rank ``n_classes - 1``, ``[d_model, n_classes - 1]``.

83 Raises:

84 ValueError: If fewer than two classes are present.

85 """

86 values = np.asarray(labels)

87 classes = np.unique(values)

88 if len(classes) < 2:

89 raise ValueError("mean_difference_basis needs at least two classes")

90 acts = activations.float().cpu()

91 means = torch.stack([acts[torch.as_tensor(values == c)].mean(0) for c in classes])

92 return orthonormal_subspace((means - means.mean(0)).T, rank=len(classes) - 1).basis

95def probe_basis(

96 weights: Float[torch.Tensor, "n_rows d_model"],

97) -> Float[torch.Tensor, "d_model rank"]:

98 """The directions of a linear probe's weights, in activation space.

100 Undo any input scaling first: for a probe fit on standardised inputs, divide each weight by its

101 feature's standard deviation.

102

103 Args:

104 weights: Probe weights; one row for a binary probe, one row per class otherwise.

105

106 Returns:

107 Orthonormal basis: the weight direction for one row, the span of the centred rows (rank

108 ``n_rows - 1``) for several.

109 """

110 w = torch.as_tensor(weights, dtype=torch.float32)

111 if w.shape[0] == 1:

112 return orthonormal_subspace(w.T, rank=1).basis

113 return orthonormal_subspace((w - w.mean(0)).T, rank=w.shape[0] - 1).basis

114

115

116Neurons = Union[Sequence[int], np.ndarray, torch.Tensor]

117"""Integer neuron indices: a list, a NumPy array, or a tensor such as a sparse probe's

118``selected_features``."""

119

120

121def neuron_basis(d_mlp: int, neurons: Neurons) -> Float[torch.Tensor, "d_mlp k"]:

122 """The coordinate basis of a set of MLP neurons, e.g. the features a sparse probe selects.

123

124 Interchanging it sets exactly those neurons to the source's values, at a neuron hook such as

125 ``blocks.N.mlp.hook_post``. Compare it with :func:`random_neuron_basis` of the same size: a random

126 direction is the wrong placebo for a claim about neurons, because it spreads over every neuron.

127

128 Args:

129 d_mlp: Number of neurons at the hook.

130 neurons: Integer indices of the neurons, each in ``[0, d_mlp)`` and none repeated.

131

132 Returns:

133 Orthonormal basis of unit vectors, ``[d_mlp, k]``, in the order given.

134

135 Raises:

136 ValueError: If no neuron is given, or an index is not an integer, out of range or repeated.

137 """

138 index = _neuron_indices(d_mlp, neurons)

139 if index.numel() == 0:

140 raise ValueError("neuron_basis needs at least one neuron")

141 basis = torch.zeros(d_mlp, index.numel())

142 basis[index, torch.arange(index.numel())] = 1.0

143 return basis

144

145

146def random_neuron_basis(

147 d_mlp: int,

148 k: int,

149 seed: int,

150 exclude: Optional[Neurons] = None,

151) -> Float[torch.Tensor, "d_mlp k"]:

152 """The coordinate basis of ``k`` neurons drawn at random: the placebo for :func:`neuron_basis`.

153

154 Args:

155 d_mlp: Number of neurons at the hook.

156 seed: Seed for the draw; the caller's random stream is untouched.

157 exclude: Neurons the draw must avoid, usually the ones the placebo is compared with.

158

159 Returns:

160 Orthonormal basis of unit vectors, ``[d_mlp, k]``.

161

162 Raises:

163 ValueError: If fewer than ``k`` neurons remain to draw from, or ``exclude`` holds an

164 invalid index.

165 """

166 excluded = _neuron_indices(d_mlp, exclude if exclude is not None else [])

167 allowed = torch.ones(d_mlp, dtype=torch.bool)

168 allowed[excluded] = False

169 pool = allowed.nonzero().flatten()

170 if not 0 < k <= len(pool):

171 raise ValueError(

172 f"k must be in [1, {len(pool)}] with {len(excluded)} neurons excluded, got {k}"

173 )

174 generator = torch.Generator().manual_seed(seed)

175 return neuron_basis(d_mlp, pool[torch.randperm(len(pool), generator=generator)[:k]])

176

177

178def _neuron_indices(d_mlp: int, neurons: Neurons) -> Int[torch.Tensor, "k"]:

179 index = torch.as_tensor(neurons)

180 if index.numel() == 0: # an empty list arrives as a float tensor

181 return torch.empty(0, dtype=torch.long)

182 if index.dtype == torch.bool or index.is_floating_point() or index.is_complex():

183 raise ValueError(f"neuron indices must be integers, got dtype {index.dtype}")

184 if index.dim() != 1:

185 raise ValueError(f"neuron indices must be one-dimensional, got shape {tuple(index.shape)}")

186 index = index.long().cpu()

187 if ((index < 0) | (index >= d_mlp)).any():

188 raise ValueError(f"neuron indices must be in [0, {d_mlp}), got {index.tolist()}")

189 if index.unique().numel() != index.numel():

190 raise ValueError(f"neuron indices must not repeat, got {index.tolist()}")

191 return index

192

193

194class LearnedSubspace(nn.Module):

195 """A trainable rank-``rank`` subspace of ``d_model`` whose basis stays orthonormal.

196

197 Args:

198 rank: Number of basis directions, in ``[1, d_model]``.

199 """

200

201 def __init__(self, d_model: int, rank: int):

202 super().__init__()

203 if not 0 < rank <= d_model:

204 raise ValueError(f"rank must be in [1, {d_model}], got {rank}")

205 self.proj = orthogonal(nn.Linear(rank, d_model, bias=False))

206

207 @property

208 def basis(self) -> Float[torch.Tensor, "d_model rank"]:

209 return cast(torch.Tensor, self.proj.weight)

210

211

212Basis = Union[torch.Tensor, LearnedSubspace]

213

214

215def _basis(basis: Basis, like: torch.Tensor) -> torch.Tensor:

216 b = basis.basis if isinstance(basis, LearnedSubspace) else basis

217 return b.to(like.device, like.dtype)

218

219

220@dataclass(frozen=True)

221class SubspaceInterchange:

222 """Set the base's coordinates inside ``basis`` to the source's.

223

224 Attributes:

225 basis: Orthonormal basis, ``[d_model, rank]``, or a :class:`LearnedSubspace`.

226 strength: 1 interchanges the coordinates; other values scale the move, 0 leaves the base.

227 """

228

229 basis: Basis

230 strength: float = 1.0

231

232 def __call__(self, base: torch.Tensor, source: torch.Tensor) -> torch.Tensor:

233 B = _basis(self.basis, base)

234 return base + self.strength * ((source - base) @ B) @ B.T

235

236

237@dataclass(frozen=True)

238class NormMatchedPlacebo:

239 """Move each base exactly as far as ``reference`` does, but inside a random subspace.

240

241 An effect that the placebo also produces comes from displacing the activation by that much,

242 not from the aligned subspace.

243

244 Attributes:

245 reference: The intervention whose displacement is matched.

246 basis: Orthonormal basis of the placebo subspace, usually from :func:`random_basis` with

247 the reference's rank, or from :func:`random_neuron_basis` for a reference built by

248 :func:`neuron_basis`.

249 """

250

251 reference: SubspaceInterchange

252 basis: torch.Tensor

253

254 def __call__(self, base: torch.Tensor, source: torch.Tensor) -> torch.Tensor:

255 distance = (self.reference(base, source) - base).norm(dim=-1, keepdim=True)

256 B = self.basis.to(base.device, base.dtype)

257 move = ((source - base) @ B) @ B.T

258 return base + move * (distance / move.norm(dim=-1, keepdim=True).clamp_min(1e-8))

259

260

261def _select(act: torch.Tensor, positions: Optional[torch.Tensor]) -> torch.Tensor:

262 if positions is None:

263 if act.ndim != 2:

264 raise ValueError(

265 f"activation has shape {tuple(act.shape)}; pass one position index per example"

266 )

267 return act

268 if act.ndim != 3:

269 raise ValueError(

270 f"with positions, the activation must be [batch, pos, d_model]; "

271 f"got shape {tuple(act.shape)}"

272 )

273 if positions.shape != (act.shape[0],):

274 raise ValueError(

275 f"positions must hold one index per example, shape ({act.shape[0]},); "

276 f"got {tuple(positions.shape)}"

277 )

278 return act[torch.arange(act.shape[0], device=act.device), positions.to(act.device)]

279

280

281def _write(act: torch.Tensor, positions: Optional[torch.Tensor], new: torch.Tensor) -> torch.Tensor:

282 if positions is None:

283 return new

284 out = act.clone()

285 out[torch.arange(act.shape[0], device=act.device), positions.to(act.device)] = new

286 return out

287

288

289def _chunks(n: int, size: int) -> Iterator[slice]:

290 for start in range(0, n, size):

291 yield slice(start, min(start + size, n))

292

293

294def _at(

295 positions: Optional[torch.Tensor], idx: Union[slice, torch.Tensor]

296) -> Optional[torch.Tensor]:

297 return None if positions is None else positions[idx]

298

299

300@torch.no_grad()

301def collect_activations(

302 model: Any,

303 hook_name: str,

304 inputs: torch.Tensor,

305 task_logits: TaskLogits,

306 *,

307 positions: Optional[Int[torch.Tensor, "batch"]] = None,

308 batch_size: int = 256,

309) -> Float[torch.Tensor, "batch d_model"]:

310 """Activations at ``hook_name``, at one position per example if given, returned on the CPU.

311

312 Args:

313 model: A model exposing the hook API.

314 inputs: Model inputs, one row per example.

315 task_logits: Runs the model on a batch of inputs.

316 positions: For ``[batch, pos, d_model]`` activations, the position to read per example.

317 batch_size: Rows per forward pass.

318

319 Returns:

320 Activations, ``[batch, d_model]``.

321

322 Raises:

323 ValueError: If any part of the model is in training mode, or the activation does not match

324 the positions given (or not given).

325 RuntimeError: If the hook does not fire exactly once per forward pass.

326 """

327 require_eval_mode(model, operation="collect_activations")

328 out = []

329 for idx in _chunks(len(inputs), batch_size):

330 store: list[torch.Tensor] = []

331

332 def grab(

333 act: torch.Tensor, hook: Any, _positions: Optional[torch.Tensor] = _at(positions, idx)

334 ) -> None:

335 store.append(_select(act, _positions).detach().cpu())

336

337 with model.hooks(fwd_hooks=[(hook_name, grab)]):

338 task_logits(inputs[idx])

339 if len(store) != 1:

340 raise RuntimeError(f"{hook_name} fired {len(store)} times in one forward pass")

341 out.append(store[0])

342 return torch.cat(out)

343

344

345def run_with_intervention(

346 model: Any,

347 hook_name: str,

348 intervention: Intervention,

349 base_inputs: torch.Tensor,

350 source_activations: Float[torch.Tensor, "batch d_model"],

351 task_logits: TaskLogits,

352 *,

353 positions: Optional[Int[torch.Tensor, "batch"]] = None,

354) -> Float[torch.Tensor, "batch n_classes"]:

355 """Task logits for ``base_inputs`` with ``intervention`` applied at ``hook_name``, in one pass.

356

357 Differentiable with respect to any parameters of ``intervention``, for training an

358 intervention. To evaluate one over many pairs in bounded memory, use :func:`intervened_logits`.

359

360 Args:

361 model: A model exposing the hook API.

362 intervention: Maps base and source activations to new base activations, e.g. a

363 :class:`SubspaceInterchange`.

364 source_activations: Activations from the source inputs at the same hook point, one row per

365 base, e.g. from :func:`collect_activations`.

366 task_logits: Runs the model on a batch of inputs.

367 positions: For ``[batch, pos, d_model]`` activations, the position to intervene on per base.

368

369 Returns:

370 Task logits under the intervention.

371

372 Raises:

373 ValueError: If the activation does not match the positions given (or not given).

374 RuntimeError: If the hook does not fire exactly once in the forward pass.

375 """

376 calls = 0

377

378 def hook(act: torch.Tensor, hook: Any) -> torch.Tensor:

379 nonlocal calls

380 calls += 1

381 selected = _select(act, positions)

382 source = source_activations.to(act.device, act.dtype)

383 return _write(act, positions, intervention(selected, source))

384

385 with model.hooks(fwd_hooks=[(hook_name, hook)]):

386 logits = task_logits(base_inputs)

387 # A second firing would reuse source activations collected for the first

388 if calls != 1:

389 raise RuntimeError(f"{hook_name} fired {calls} times in one forward pass")

390 return logits

391

392

393@torch.no_grad()

394def intervened_logits(

395 model: Any,

396 hook_name: str,

397 intervention: Intervention,

398 base_inputs: torch.Tensor,

399 source_activations: Float[torch.Tensor, "batch d_model"],

400 task_logits: TaskLogits,

401 *,

402 positions: Optional[Int[torch.Tensor, "batch"]] = None,

403 batch_size: int = 256,

404) -> Float[torch.Tensor, "batch n_classes"]:

405 """Task logits under ``intervention``, in batches without gradients, returned on the CPU.

406

407 The evaluation counterpart of :func:`run_with_intervention`, with the same arguments and

408 checks: memory is bounded by ``batch_size`` rather than by the number of pairs.

409

410 Args:

411 model: A model exposing the hook API.

412 intervention: Maps base and source activations to new base activations.

413 source_activations: Activations from the source inputs at the same hook point, one row per

414 base.

415 task_logits: Runs the model on a batch of inputs.

416 positions: For ``[batch, pos, d_model]`` activations, the position to intervene on per base.

417 batch_size: Rows per forward pass.

418

419 Returns:

420 Task logits under the intervention, ``[batch, n_classes]``.

421

422 Raises:

423 ValueError: If any part of the model is in training mode, or the inputs do not pair up.

424 """

425 require_eval_mode(model, operation="intervened_logits")

426 if len(source_activations) != len(base_inputs):

427 raise ValueError("source_activations must have one row per base input")

428 out = []

429 for idx in _chunks(len(base_inputs), batch_size):

430 logits = run_with_intervention(

431 model,

432 hook_name,

433 intervention,

434 base_inputs[idx],

435 source_activations[idx],

436 task_logits,

437 positions=_at(positions, idx),

438 )

439 out.append(logits.detach().cpu())

440 return torch.cat(out)

441

442

443def _model_device(model: Any) -> torch.device:

444 parameter = next(iter(model.parameters()), None)

445 return parameter.device if parameter is not None else torch.device("cpu")

446

447

448@dataclass(frozen=True)

449class DASResult:

450 """A subspace learned by :func:`fit_das`.

451

452 Attributes:

453 subspace: The learned subspace; pass it to :class:`SubspaceInterchange`.

454 losses: Mean training loss per epoch.

455 hook_name: The hook point it was trained at.

456 """

457

458 subspace: LearnedSubspace

459 losses: tuple[float, ...]

460 hook_name: str

461 rank: int

462

463 @property

464 def basis(self) -> Float[torch.Tensor, "d_model rank"]:

465 return self.subspace.basis.detach()

466

467

468def fit_das(

469 model: Any,

470 hook_name: str,

471 base_inputs: torch.Tensor,

472 source_inputs: torch.Tensor,

473 targets: Int[torch.Tensor, "batch"],

474 task_logits: TaskLogits,

475 *,

476 rank: int,

477 base_positions: Optional[Int[torch.Tensor, "batch"]] = None,

478 source_positions: Optional[Int[torch.Tensor, "batch"]] = None,

479 epochs: int = 20,

480 batch_size: int = 128,

481 lr: float = 1e-2,

482 strength: float = 1.0,

483 seed: int = 0,

484 device: Optional[Union[str, torch.device]] = None,

485 micro_batch_size: Optional[int] = None,

486) -> DASResult:

487 """Learn the subspace whose interchange makes the model output ``targets`` (DAS).

488

489 ``targets[i]`` is the class the high-level causal model predicts for base ``i`` with the

490 aligned variable set to its value in source ``i``. The basis is trained with cross-entropy

491 on the task logits under the interchange. Model parameters are frozen during the call and

492 their ``requires_grad`` flags restored afterwards.

493

494 Train on pairs where several variables differ between base and source, labelled by the

495 high-level model's interchange, as well as on pairs that differ only in the aligned variable:

496 otherwise the subspace can learn whatever the pairs happen to move together.

497

498 Args:

499 model: A model exposing the hook API.

500 hook_name: The hook point to align.

501 base_inputs: Base inputs, one row per pair.

502 source_inputs: Source inputs, one row per pair.

503 targets: The high-level model's output class for each pair under the interchange.

504 task_logits: Runs the model on a batch of inputs and returns task logits.

505 rank: Dimension of the subspace.

506 base_positions: For ``[batch, pos, d_model]`` activations, the base position per pair.

507 source_positions: Likewise for the sources.

508 epochs: Passes over the pairs.

509 batch_size: Pairs per optimisation step.

510 lr: Adam learning rate.

511 strength: Interchange strength during training.

512 seed: Seed for the basis initialisation and the batch order.

513 device: Device for the subspace's parameters; by default the device of the model's

514 parameters, or the CPU for a model without any.

515 micro_batch_size: Split each batch into parts this size and accumulate gradients: the same

516 gradients, up to floating-point rounding, in less memory.

517

518 Returns:

519 The learned subspace with its training losses.

520

521 Raises:

522 ValueError: If any part of the model is in training mode, or the inputs and targets differ

523 in length.

524 """

525 require_eval_mode(model, operation="fit_das")

526 if not (len(base_inputs) == len(source_inputs) == len(targets)):

527 raise ValueError("base_inputs, source_inputs and targets must have the same length")

528 rng = np.random.default_rng(seed)

529 micro = micro_batch_size or batch_size

530 source_acts = collect_activations(

531 model, hook_name, source_inputs, task_logits, positions=source_positions, batch_size=micro

532 )

533 # The subspace is built on the CPU, so seed only the CPU generator, inside a fork that restores

534 # it; torch.manual_seed would also reseed every CUDA, MPS and XPU generator

535 with torch.random.fork_rng(devices=[]):

536 torch.default_generator.manual_seed(seed)

537 subspace = LearnedSubspace(source_acts.shape[1], rank)

538 subspace = subspace.to(device if device is not None else _model_device(model))

539 step = SubspaceInterchange(subspace, strength)

540 optimizer = torch.optim.Adam(subspace.parameters(), lr=lr)

541 labels = torch.as_tensor(targets, dtype=torch.long)

542

543 flags = [p.requires_grad for p in model.parameters()]

544 model.requires_grad_(False)

545 losses = []

546 try:

547 for _ in range(epochs):

548 order = torch.as_tensor(rng.permutation(len(labels)))

549 total = 0.0

550 for chunk in _chunks(len(order), batch_size):

551 batch = order[chunk]

552 optimizer.zero_grad()

553 for part in _chunks(len(batch), micro):

554 idx = batch[part]

555 logits = run_with_intervention(

556 model,

557 hook_name,

558 step,

559 base_inputs[idx],

560 source_acts[idx],

561 task_logits,

562 positions=_at(base_positions, idx),

563 )

564 # Summed, then divided by the whole batch, so split batches give equal gradients

565 loss = F.cross_entropy(

566 logits, labels[idx].to(logits.device), reduction="sum"

567 ) / len(batch)

568 loss.backward()

569 total += loss.item() * len(batch)

570 optimizer.step()

571 losses.append(total / len(order))

572 finally:

573 for p, flag in zip(model.parameters(), flags):

574 p.requires_grad_(flag)

575 return DASResult(subspace.eval(), tuple(losses), hook_name, rank)

576

577

578def arbitrary_targets(

579 keys: Sequence[Hashable], n_classes: int, seed: int

580) -> Int[torch.Tensor, "batch"]:

581 """A fixed pseudo-random class per key: targets for the arbitrary-label control.

582

583 Train :func:`fit_das` on these instead of the high-level model's targets, with the same pairs,

584 rank and budget. The interchange accuracy it reaches shows how much accuracy the method's

585 expressivity alone produces; a learned subspace should clearly exceed it.

586

587 Args:

588 keys: One key per pair identifying its source, e.g. the source text.

589

590 Returns:

591 A class in ``[0, n_classes)`` per key; equal keys always get equal classes.

592 """

593

594 def draw(key: Hashable) -> int:

595 digest = hashlib.sha256(f"{seed}:{key!r}".encode()).digest()

596 return int.from_bytes(digest[:8], "big") % n_classes

597

598 return torch.tensor([draw(k) for k in keys], dtype=torch.long)

599

600

601def reinitialized(model: nn.Module, seed: int) -> nn.Module:

602 """A copy of ``model`` with freshly initialised parameters: the random-init control.

603

604 For models built from standard ``nn`` layers, such as a ``HookedRootModule``: every module with

605 ``reset_parameters`` is reset, so every parameter must belong to one. Buffers, such as fixed

606 architectural matrices, are kept. For HuggingFace models use :func:`random_init_bridge`, which

607 applies the architecture's own initialisation scheme.

608

609 Args:

610 model: The trained model; it is not modified.

611

612 Returns:

613 The untrained copy, on the device of the model's parameters, in evaluation mode.

614

615 Raises:

616 ValueError: If a parameter belongs to no module with ``reset_parameters``, so it would keep

617 its trained value.

618 """

619 covered = {

620 id(parameter)

621 for module in model.modules()

622 if callable(getattr(module, "reset_parameters", None))

623 for parameter in module.parameters(recurse=False)

624 }

625 missing = [name for name, parameter in model.named_parameters() if id(parameter) not in covered]

626 if missing:

627 more = f" and {len(missing) - 5} more" if len(missing) > 5 else ""

628 raise ValueError(

629 f"reinitialized cannot reset {', '.join(missing[:5])}{more}: no module's "

630 "reset_parameters covers them, so they would keep their trained values. Reset them "

631 "yourself, or use random_init_bridge for a HuggingFace model."

632 )

633 device = _model_device(model)

634 clone = copy.deepcopy(model).cpu()

635 # Reset on the CPU, seeding only the CPU generator inside a fork that restores it;

636 # torch.manual_seed would also reseed every CUDA, MPS and XPU generator

637 with torch.random.fork_rng(devices=[]):

638 torch.default_generator.manual_seed(seed)

639 for module in clone.modules():

640 reset = getattr(module, "reset_parameters", None)

641 if callable(reset):

642 reset()

643 return clone.to(device).eval()

644

645

646def random_init_bridge(

647 model_name: str,

648 *,

649 seed: int,

650 dtype: torch.dtype = torch.float32,

651 device: Optional[Union[str, torch.device]] = None,

652 **boot_kwargs: Any,

653) -> Any:

654 """A ``TransformerBridge`` with the architecture's HuggingFace initialisation and no training.

655

656 The random-init control: train :func:`fit_das` on it with the same pairs, rank and budget as

657 on the trained model. Effects DAS finds here come from the method and the architecture, not

658 from anything the model learned. Built with ``load_weights=False``, so the model class and

659 prepared configuration match the trained bridge.

660

661 Args:

662 model_name: Model name or alias, as for ``TransformerBridge.boot_transformers``.

663 device: Device for the bridge; by default the device ``boot_transformers`` would choose.

664 **boot_kwargs: Passed on to ``TransformerBridge.boot_transformers``, e.g. ``revision``.

665

666 Returns:

667 The untrained ``TransformerBridge``, in evaluation mode.

668 """

669 from transformer_lens.model_bridge import TransformerBridge

670 from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params

671

672 # The model is built on the CPU, so seed only the CPU generator, inside a fork that restores it

673 with torch.random.fork_rng(devices=[]):

674 torch.default_generator.manual_seed(seed)

675 bridge = TransformerBridge.boot_transformers(

676 model_name, load_weights=False, dtype=dtype, device=device, **boot_kwargs

677 )

678 # Building from a config follows the config's dtype and leaves the weights on the CPU,

679 # whatever device boot recorded; HuggingFace also leaves the model in training mode

680 maybe_cast_floating_params(bridge.original_model, dtype)

681 return bridge.to(bridge.cfg.device).eval()