2. Constitutional-AI-Style Classifier Head on Frozen Embeddings

Problem

A Constitutional-AI-style review step needs to classify a candidate response against a fixed set of written principles, using embeddings from a large frozen encoder that is far too expensive to fine-tune for this one task. The standard move is to freeze the encoder entirely and train only a small classifier head on top of its output embeddings.

The head must never let gradients leak back into the frozen encoder's output, even if a caller forgets to detach the embeddings before passing them in — that guarantee belongs inside the head, not the caller.

Source: src/2_constitutional_classifier_head.py

class ConstitutionalClassifierHead(torch.nn.Module):
    def __init__(self, embed_dim: int, num_principles: int): ...
    def forward(self, embeddings: torch.Tensor) -> torch.Tensor: ...

>>> torch.manual_seed(0)
>>> head = ConstitutionalClassifierHead(embed_dim=4, num_principles=3)
>>> out = head(torch.randn(2, 4))
>>> out.shape
torch.Size([2, 3])
>>> torch.allclose(out.sum(dim=-1), torch.ones(2))
True

Step-by-Step Approach

1. In __init__, build one small trainable layer: nn.Linear(embed_dim, num_principles). This is the only part of the model that gets gradients.
2. In forward, call embeddings.detach() before doing anything else with the input — this severs the input from whatever graph produced it, regardless of what the caller did.
3. Pass the detached embeddings through the linear layer to get raw logits of shape (batch, num_principles).
4. Apply softmax along the last dimension so each row is a valid probability distribution over principles.
5. Sanity check: after a backward pass through the output, the original (pre-detach) embeddings tensor should have grad is None, while the linear layer's weight should have a populated gradient.

The key insight is that "frozen" is a guarantee the module itself must enforce with .detach(), not a convention the caller is trusted to follow — relying on the caller to remember torch.no_grad() or pre-detach the embeddings is exactly the kind of bug that silently wastes compute fine-tuning a model that was supposed to stay frozen.

Reference solution

import torch
import torch.nn as nn
import torch.nn.functional as F


class ConstitutionalClassifierHead(nn.Module):
    def __init__(self, embed_dim: int, num_principles: int):
        super().__init__()
        self.linear = nn.Linear(embed_dim, num_principles)

    def forward(self, embeddings: torch.Tensor) -> torch.Tensor:
        # .detach() guarantees no gradient reaches the (frozen) input
        # embeddings even if the caller forgot to detach them upstream --
        # the head owns this guarantee, not its caller
        logits = self.linear(embeddings.detach())  # (batch, num_principles)
        return F.softmax(logits, dim=-1)

Key Functions & Tricks

• tensor.detach() — returns a new tensor sharing storage but severed from the autograd graph; anything computed from it never propagates gradients backward past this point.
• nn.Linear(embed_dim, num_principles) — the only trainable component; its weight and bias receive gradients normally since they were never detached.
• F.softmax(logits, dim=-1) — converts raw scores into a probability distribution over the last (principles) axis.
• tensor.grad is None — the way to verify in a test that no gradient reached a given leaf tensor after .backward().

How to Recognize This Pattern