Train one TokenRouter per checkpoint layer via binary cross-entropy.
(
hidden_states: Dict,
labels: Dict[int, torch.Tensor],
config: TIDEConfig,
epochs: int = 100,
lr: float = 1e-3,
device: str = "cpu",
)
| 104 | |
| 105 | |
| 106 | def train_routers( |
| 107 | hidden_states: Dict, |
| 108 | labels: Dict[int, torch.Tensor], |
| 109 | config: TIDEConfig, |
| 110 | epochs: int = 100, |
| 111 | lr: float = 1e-3, |
| 112 | device: str = "cpu", |
| 113 | ) -> Dict[int, TokenRouter]: |
| 114 | """Train one TokenRouter per checkpoint layer via binary cross-entropy.""" |
| 115 | routers = {} |
| 116 | hidden_dim = hidden_states["final"].shape[-1] |
| 117 | |
| 118 | for layer_idx in sorted(labels.keys()): |
| 119 | h = hidden_states[layer_idx].to(device).float() |
| 120 | y = labels[layer_idx].to(device).float() |
| 121 | n = min(h.shape[0], y.shape[0]) |
| 122 | h, y = h[:n], y[:n] |
| 123 | |
| 124 | router = TokenRouter(hidden_dim, config.router_bottleneck_dim).to(device) |
| 125 | optimizer = torch.optim.Adam(router.parameters(), lr=lr) |
| 126 | criterion = nn.BCELoss() |
| 127 | |
| 128 | best_loss = float("inf") |
| 129 | for epoch in range(epochs): |
| 130 | optimizer.zero_grad() |
| 131 | pred = router(h) |
| 132 | loss = criterion(pred, y) |
| 133 | loss.backward() |
| 134 | optimizer.step() |
| 135 | |
| 136 | if loss.item() < best_loss: |
| 137 | best_loss = loss.item() |
| 138 | |
| 139 | if (epoch + 1) % 25 == 0: |
| 140 | acc = ((pred > 0.5) == (y > 0.5)).float().mean().item() |
| 141 | logger.info(f" Layer {layer_idx} epoch {epoch+1}: loss={loss.item():.4f} acc={acc:.3f}") |
| 142 | |
| 143 | routers[layer_idx] = router.cpu() |
| 144 | logger.info(f" Layer {layer_idx} final loss: {best_loss:.4f}") |
| 145 | |
| 146 | return routers |
| 147 | |
| 148 | |
| 149 | def calibrate( |