# EBT — From One Score to Q, V, Policy and Advantage
# Book: Models From First Principles
# Chapter URL: https://programmer.ie/books/models-from-first-principles/03-chapter/
# Prose fingerprint (sha256, normalized): 6f676789d07d77d508303e2e78b1ace2b7d6ca5287798ca4d03adc7c17cf812e
# Snapshot generated at build time by Hugo. This is the complete rendered prose
# of the chapter; it is not a transcript of the source file.

EBT — From One Score to Q, V, Policy and Advantage In the previous post we built MR.Q: a small model that takes a context embedding and a response embedding, combines them, and predicts one scalar.
That architecture is useful because it is brutally simple:
context embedding + response embedding ↓ encoder ↓ representation z ↓ predictor ↓ Q value But one scalar eventually becomes restrictive.
Suppose the model predicts that a candidate response has quality 0.82.
That number tells us something about the candidate.
It does not tell us:
whether 0.82 is good relative to what was expected for this context; whether the model believes a different action should be taken; whether several downstream actions are plausible; whether the candidate is merely acceptable or genuinely better than baseline; whether the representation contains enough information to support more than one decision. This is the architectural pressure that leads to EBT.
Instead of asking one question of the representation, we ask several.
context + response ↓ encoder ↓ z ┌───────────┼───────────┐ ↓ ↓ ↓ Q V Policy │ │ │ │ │ action logits │ │ ↓ │ │ action probs │ │ └──── Q - V ┘ ↓ advantage This looks like a much more sophisticated model.
It is not one mysterious object.
It is still a collection of small neural networks sharing one representation.
That is the central idea of this series.
1. Start with the shared representation We begin exactly where MR.Q began.
Assume we already have two embeddings:
context_emb.shape == [B, D] response_emb.shape == [B, D] For example:
B = 32 D = 1024 so:
context_emb [32, 1024] response_emb [32, 1024] The simplest way to combine them is concatenation:
combined = torch.cat([context_emb, response_emb], dim=-1) Now:
combined [32, 2048] Then we compress that pair into a smaller latent representation.
import torch from torch import nn class PairEncoder(nn.Module): def __init__(self, embedding_dim=1024, hidden_dim=256): super().__init__() self.net = nn.Sequential( nn.Linear(embedding_dim * 2, hidden_dim), nn.ReLU(), nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, hidden_dim), ) def forward(self, context_emb, response_emb): pair = torch.cat([context_emb, response_emb], dim=-1) return self.net(pair) The contract is:
[B, D] + [B, D] ↓ concatenate ↓ [B, 2D] ↓ encoder ↓ [B, H] If D = 1024 and H = 256:
[B, 1024] + [B, 1024] ↓ [B, 2048] ↓ [B, 256] Everything that follows will read the same [B, 256] representation.
That is important.
We are not training three completely separate models from the raw inputs.
We are learning one shared representation and attaching several small decision surfaces to it.
2. The Q head The Q head asks approximately:
How good is this context-response pair?
Architecturally it can be tiny.
class QHead(nn.Module): def __init__(self, hidden_dim=256): super().__init__() self.net = nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, 1), ) def forward(self, z): return self.net(z).squeeze(-1) Shape flow:
[B, 256] ↓ Linear(256 → 128) ↓ [B, 128] ↓ ReLU ↓ Linear(128 → 1) ↓ [B, 1] ↓ squeeze(-1) ↓ [B] Notice what happened.
We called it a Q head.
But mechanically it is only:
Linear ↓ ReLU ↓ Linear The name does not create the semantics.
The training objective does.
That distinction remains fundamental.
3. The V head Now we add a second model.
The V head has nearly the same architecture:
class VHead(nn.Module): def __init__(self, hidden_dim=256): super().__init__() self.net = nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, 1), ) def forward(self, z): return self.net(z).squeeze(-1) Mechanically it looks almost identical to Q.
So why have both?
Because they are intended to learn different targets.
A useful conceptual distinction is:
Q = value of this particular candidate/action V = expected value for the state/context Then:
Q - V asks:
Is this candidate better or worse than what we should normally expect here?
That difference is the advantage.
4. Advantage is not another neural network This is worth slowing down for.
A model architecture diagram can become intimidating because every named output sounds like another learned model.
Advantage is not necessarily another learned network.
It can simply be arithmetic:
advantage = q_value - state_value If:
Q = 0.82 V = 0.60 then:
A = +0.22 The candidate is above the estimated baseline.
If:
Q = 0.43 V = 0.61 then:
A = -0.18 The candidate is below baseline.
This is a recurring pattern in neural architecture:
learn several useful quantities, then derive additional signals algebraically.
Not every output needs another network.
5. The policy head The third learned component is the policy head.
Instead of returning one scalar, it returns one logit per action.
Suppose our action space is:
0 = reject 1 = accept 2 = revise Then:
class PolicyHead(nn.Module): def __init__(self, hidden_dim=256, num_actions=3): super().__init__() self.net = nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, num_actions), ) def forward(self, z): return self.net(z) Shape flow:
[B, 256] ↓ Linear(256 → 128) ↓ [B, 128] ↓ ReLU ↓ Linear(128 → 3) ↓ [B, 3] Those three numbers are logits.
They are not yet probabilities.
action_logits = policy_head(z) action_probs = torch.softmax(action_logits, dim=-1) Now each row sums to approximately one.
print(action_probs[0]) might produce:
tensor([0.08, 0.71, 0.21]) Meaning:
reject 8% accept 71% revise 21% Again, the architecture is not magical.
It is another two-layer MLP reading the same latent representation.
6. Put the pieces together Now we can build the whole model.
class EBTModel(nn.Module): def __init__( self, embedding_dim=1024, hidden_dim=256, num_actions=3, ): super().__init__() self.encoder = PairEncoder( embedding_dim=embedding_dim, hidden_dim=hidden_dim, ) self.q_head = QHead(hidden_dim) self.v_head = VHead(hidden_dim) self.policy_head = PolicyHead(hidden_dim, num_actions) def forward(self, context_emb, response_emb): z = self.encoder(context_emb, response_emb) q_value = self.q_head(z) state_value = self.v_head(z) action_logits = self.policy_head(z) action_probs = torch.softmax(action_logits, dim=-1) advantage = q_value - state_value return { "z": z, "q_value": q_value, "state_value": state_value, "advantage": advantage, "action_logits": action_logits, "action_probs": action_probs, } The complete model is therefore:
EBTModel │ ├── PairEncoder │ ├── Linear │ ├── ReLU │ ├── LayerNorm │ └── Linear │ ├── QHead │ ├── Linear │ ├── ReLU │ └── Linear │ ├── VHead │ ├── Linear │ ├── ReLU │ └── Linear │ └── PolicyHead ├── Linear ├── ReLU └── Linear And then one subtraction:
advantage = Q - V That is the model inside the model.
7. Prove the tensor contracts Before training anything, test the shapes.
B = 8 D = 1024 context = torch.randn(B, D) response = torch.randn(B, D) model = EBTModel( embedding_dim=D, hidden_dim=256, num_actions=3, ) out = model(context, response) for key, value in out.items(): print(key, tuple(value.shape)) Expected:
z (8, 256) q_value (8,) state_value (8,) advantage (8,) action_logits (8, 3) action_probs (8, 3) Then assert it.
assert out["z"].shape == (B, 256) assert out["q_value"].shape == (B,) assert out["state_value"].shape == (B,) assert out["advantage"].shape == (B,) assert out["action_logits"].shape == (B, 3) assert out["action_probs"].shape == (B, 3) A model diagram is a hypothesis.
The runtime tensor shapes are evidence.
8. Why shared representation matters We could build three independent encoders:
context + response → Q encoder → Q context + response → V encoder → V context + response → policy encoder → policy But that triples much of the representation work.
Instead EBT uses:
context + response ↓ shared encoder ↓ z ┌─────┼─────┐ ↓ ↓ ↓ Q V Policy This creates two potential advantages.
First, computation is shared.
Second, the representation receives training signal from several objectives.
But shared representation is not automatically better.
The tasks can interfere with one another.
If the policy objective pushes z in one direction while the Q objective pushes it in another, multi-task learning can hurt.
So this is not a free architectural improvement.
It is an empirical question.
9. Count the parameters Let us make the cost concrete.
def count_parameters(module): return sum(p.numel() for p in module.parameters() if p.requires_grad) print("encoder", count_parameters(model.encoder)) print("Q", count_parameters(model.q_head)) print("V", count_parameters(model.v_head)) print("policy", count_parameters(model.policy_head)) print("total", count_parameters(model)) With:
embedding_dim = 1024 hidden_dim = 256 num_actions = 3 the first encoder layer dominates because it maps:
2048 → 256 The heads are comparatively cheap.
That tells us something architectural:
Once you have paid for a useful shared representation, adding a few small prediction heads can be inexpensive.
This is one reason multi-head architectures are attractive.
10. Q and V need different targets A common implementation mistake is to build separate heads but train them on effectively the same target.
Then you have created two networks with different names but no meaningful semantic separation.
Suppose every training example has a scalar observed return r.
A naive implementation might do:
q_loss = F.mse_loss(q_value, r) v_loss = F.mse_loss(state_value, r) That may make Q and V collapse toward the same function.
Then:
Q ≈ V and therefore:
advantage ≈ 0 The architecture contains Q and V.
The learning problem does not.
That distinction matters more than the class definitions.
11. One way to train Q Suppose each context-response pair has an observed scalar target:
reward = torch.tensor([ 0.9, 0.2, 0.7, 0.4, ]) Then Q can be trained by regression:
q_loss = F.mse_loss(out["q_value"], reward) Or, if the target is binary acceptance:
q_loss = F.binary_cross_entropy_with_logits( out["q_value"], accepted.float(), ) Or Q could be trained pairwise.
q_pref = model(ctx, preferred)["q_value"] q_rej = model(ctx, rejected)["q_value"] q_loss = F.softplus(-(q_pref - q_rej)).mean() Same head.
Different meaning.
12. One way to train V: expectile regression A useful way to train a value surface is expectile regression.
The point is not to explain an entire offline-RL literature here.
The important mechanism is simple.
For residual:
delta = target - prediction we weight positive and negative residuals asymmetrically.
def expectile_loss(pred, target, tau=0.7): delta = target - pred weight = torch.where( delta > 0, tau, 1.0 - tau, ) return (weight * delta.square()).mean() If:
tau = 0.5 this behaves symmetrically like ordinary squared error up to a constant weighting.
If:
tau > 0.5 under-prediction receives more weight than over-prediction.
This gives us a mechanism for fitting a particular conditional expectile rather than simply the mean.
For example:
v_loss = expectile_loss( out["state_value"], q_target.detach(), tau=0.7, ) The detach() is deliberate if we do not want the V loss to change the target-producing Q graph.
Whether that is the correct objective depends on the learning design.
Again: architecture and objective are separate decisions.
13. Train the policy head Suppose the desired action for each example is an integer class:
0 reject 1 accept 2 revise Then:
policy_target = torch.tensor([1, 0, 2, 1]) and:
policy_loss = F.cross_entropy( out["action_logits"], policy_target, ) Do not apply softmax before cross_entropy.
Use raw logits.
The probabilities are useful for reporting and decisions:
action_probs = torch.softmax(out["action_logits"], dim=-1) but the loss wants logits.
14. Multi-task training Now the complete loss might look like:
loss = ( q_weight * q_loss + v_weight * v_loss + policy_weight * policy_loss ) For example:
loss = ( 1.0 * q_loss + 0.5 * v_loss + 0.25 * policy_loss ) Those weights are not decoration.
They determine how strongly each task reshapes the shared encoder.
If the policy loss is numerically ten times larger than the Q loss, even equal explicit weights may not produce equal influence.
So inspect the actual values.
print({ "q": q_loss.item(), "v": v_loss.item(), "policy": policy_loss.item(), }) Better still, inspect gradient contributions.
15. Which task is training the encoder? Because all heads share z, each objective can send gradients into the encoder.
We can measure this.
def grad_norm(module): total = 0.0 for p in module.parameters(): if p.grad is None: continue total += p.grad.detach().pow(2).sum().item() return total ** 0.5 Run one objective at a time.
optimizer.zero_grad() q_loss.backward(retain_graph=True) print("encoder grad from Q:", grad_norm(model.encoder)) Then repeat for V and policy with fresh forward passes.
This tells you something the architecture diagram cannot:
which objectives are actually shaping the shared representation.
16. Gradient conflict Suppose Q and policy both update the encoder.
Their gradients may point in similar directions.
Or opposing directions.
We can inspect the cosine similarity between flattened gradient vectors.
def collect_grad_vector(module): parts = [] for p in module.parameters(): if p.grad is None: continue parts.append(p.grad.detach().flatten()) if not parts: return None return torch.cat(parts) Then compare task gradients.
import torch.nn.functional as F g_q = collect_grad_vector(model.encoder) With separately computed vectors:
cos = F.cosine_similarity( g_q.unsqueeze(0), g_policy.unsqueeze(0), ).item() Interpretation:
+1 strongly aligned 0 roughly orthogonal -1 strongly opposed This does not automatically tell us what to change.
But it turns vague “multi-task interference” into something observable.
17. Advantage should have a distribution Once Q and V exist, inspect the advantage distribution.
adv = out["advantage"].detach() print({ "mean": adv.mean().item(), "std": adv.std(unbiased=False).item(), "min": adv.min().item(), "max": adv.max().item(), }) If advantage is always nearly zero, possible explanations include:
Q and V learned the same target; one head is copying the other through shared features; the data has little variation; the objective is wrong; the model is underfit; the intended semantics of Q and V were never represented in the labels. Do not simply celebrate because the model returns an advantage key.
Check whether the signal contains information.
18. Correlate advantage with outcomes Suppose positive advantage is intended to mean “better than baseline.”
Then test it.
For a binary observed outcome:
better_than_baseline = (observed_return > baseline_return).float() You can compare:
advantage > 0 against that target.
pred = (advantage > 0).float() accuracy = (pred == better_than_baseline).float().mean() That is crude, but it demonstrates the right scientific habit:
derive an observable implication from the name you gave the signal.
If the implication fails, the model does not get credit because the variable was called advantage.
19. Policy confidence and entropy The policy head gives a distribution.
That means we can measure how concentrated it is.
def normalized_entropy(probs, eps=1e-8): n_actions = probs.size(-1) entropy = -(probs * (probs + eps).log()).sum(dim=-1) return entropy / torch.log( torch.tensor(float(n_actions), device=probs.device) ) Then:
entropy = normalized_entropy(out["action_probs"]) Interpretation:
entropy near 0 → concentrated policy entropy near 1 → diffuse policy But low entropy does not mean correctness.
A model can be confidently wrong.
So evaluate calibration and action accuracy separately.
20. Calibration is not ranking This distinction mattered for MR.Q and matters even more here.
Suppose the Q head correctly ranks candidates:
A > B > C but outputs:
A = 8.7 B = 7.3 C = 6.5 That may be perfectly useful for ranking.
It is not a calibrated probability.
Likewise:
torch.sigmoid(q_value) puts the number into [0,1].
That alone does not make it calibrated.
A monotonic transformation preserves order but changes scale.
Always ask:
Do I need ordering? Do I need a threshold? Do I need a probability? Do I need an expected return? Those are different evaluation problems.
21. A learnable score scale Some implementations introduce a learned scalar controlling score scale.
For example:
class ScaledQ(nn.Module): def __init__(self): super().__init__() self.scale_logit = nn.Parameter(torch.tensor(0.0)) def forward(self, q): scale = torch.sigmoid(self.scale_logit) return q * scale This looks harmless.
But there is a subtle PyTorch trap.
Do not do this during the differentiable path:
scale = torch.sigmoid(self.scale_logit).item() .item() converts the tensor into a Python number and breaks gradient flow through the scale.
Use:
scale = torch.sigmoid(self.scale_logit) if you actually want to train it.
This is exactly the kind of bug the previous PyTorch series prepared us to notice.
22. Build a complete synthetic experiment Let us create a toy problem where the correct structure is known.
We will generate context and response embeddings and define a hidden quality rule.
import torch import torch.nn.functional as F torch.manual_seed(7) N = 4096 D = 32 context = torch.randn(N, D) response = torch.randn(N, D) # Hidden compatibility signal. compat = (context * response).mean(dim=-1) # Context-specific baseline. baseline = 0.25 * context[:, :4].mean(dim=-1) # Observed return. reward = compat + baseline + 0.05 * torch.randn(N) # Three actions derived from the relative result. adv_true = reward - baseline action = torch.where( adv_true > 0.10, torch.tensor(1), # accept torch.where( adv_true < -0.10, torch.tensor(0), # reject torch.tensor(2), # revise ), ) Now the model has related but non-identical tasks.
Q can predict reward.
V can predict baseline.
Policy can predict the action.
This is much cleaner than giving Q and V the same label and hoping semantics emerge.
23. Train the model model = EBTModel( embedding_dim=D, hidden_dim=64, num_actions=3, ) optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3) batch_size = 128 for step in range(1000): idx = torch.randint(0, N, (batch_size,)) out = model(context[idx], response[idx]) q_loss = F.mse_loss(out["q_value"], reward[idx]) v_loss = F.mse_loss(out["state_value"], baseline[idx]) policy_loss = F.cross_entropy(out["action_logits"], action[idx]) loss = q_loss + v_loss + 0.5 * policy_loss optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 == 0: with torch.no_grad(): pred_action = out["action_logits"].argmax(dim=-1) action_acc = (pred_action == action[idx]).float().mean() print( step, "loss", round(loss.item(), 4), "q", round(q_loss.item(), 4), "v", round(v_loss.item(), 4), "policy_acc", round(action_acc.item(), 3), ) The important point is not the exact loss values.
The important point is that each head has an explicit target corresponding to its intended meaning.
24. Evaluate each head separately After training:
with torch.no_grad(): out = model(context, response) Q error:
q_mse = F.mse_loss(out["q_value"], reward) V error:
v_mse = F.mse_loss(out["state_value"], baseline) Policy accuracy:
policy_acc = ( out["action_logits"].argmax(dim=-1) == action ).float().mean() Advantage error:
adv_mse = F.mse_loss( out["advantage"], adv_true, ) Print all of them.
print({ "q_mse": q_mse.item(), "v_mse": v_mse.item(), "policy_accuracy": policy_acc.item(), "advantage_mse": adv_mse.item(), }) A multi-head model should not be summarized by one total loss if you care about the semantics of each head.
25. Compare against MR.Q Now we can run the first important ablation.
Train an MR.Q-style model only on Q.
Then compare:
MR.Q Q performance vs EBT Q performance V performance policy performance advantage performance Possible outcomes:
Outcome A EBT improves Q too.
Maybe the auxiliary tasks regularize the shared representation.
Outcome B Q stays the same, but EBT adds useful extra signals.
Still valuable.
Outcome C Q gets worse while the extra heads work.
Now we have a trade-off.
Outcome D Nothing improves.
Then the extra complexity may not be justified.
This is how architectural progression should be evaluated.
Not:
newer model = automatically better but:
new component ↓ new measurable capability? ↓ benefit greater than cost? 26. Head ablations Remove one head at a time.
Q only Q + V Q + Policy Q + V + Policy Then compare:
Q error V error policy accuracy advantage quality parameter count training time A simple experiment table might look like:
model q_mse v_mse policy_acc params Q only ... — — ... Q + V ... ... — ... Q + Policy ... — ... ... Q + V + Policy ... ... ... ... Now we can say what each head actually buys.
27. Shared encoder ablation Another experiment:
Shared pair → encoder → z → Q/V/Policy Separate pair → Q encoder → Q pair → V encoder → V pair → P encoder → Policy The separate version has more parameters and more compute.
But it removes task interference.
There is no universal winner.
This experiment teaches a major architectural trade-off:
representation sharing saves capacity and can transfer useful structure, but it also couples objectives.
28. Freeze the encoder We can probe whether the heads themselves are the bottleneck.
First train the shared encoder and heads.
Then freeze the encoder:
for p in model.encoder.parameters(): p.requires_grad = False Train only the heads.
If performance remains strong, the representation already contains what the heads need.
If performance collapses when adapting to a new task, the encoder may need task-specific adjustment.
This is one way to separate:
representation problem from:
head problem 29. Probe the latent representation directly A model with multiple heads encourages us to inspect z.
z = out["z"].detach() Basic checks:
print({ "mean": z.mean().item(), "std": z.std(unbiased=False).item(), "norm": z.norm(dim=-1).mean().item(), }) Look for:
collapse toward a constant vector; exploding norms; dead dimensions; highly unstable scale; train/eval differences. You can also train cheap linear probes on frozen z.
If a linear probe can predict the policy action from z, that tells us action-relevant information exists in the representation.
It does not prove the representation is optimal.
But it gives evidence about what information is accessible.
30. Candidate reranking with Q and advantage Suppose one context has many candidate responses.
context.shape == [D] candidates.shape == [K, D] Repeat the context:
ctx_batch = context.unsqueeze(0).expand(K, -1) Then:
out = model(ctx_batch, candidates) You can rank by Q:
order_q = torch.argsort(out["q_value"], descending=True) or by advantage:
order_a = torch.argsort(out["advantage"], descending=True) If V is identical for every candidate under the same context, those rankings may be identical.
That observation matters.
If V is computed from the same pair representation as Q, however, it is not necessarily a pure context-only baseline.
That raises a conceptual question.
31. Is your V really V(s)? Look carefully at the data flow:
context + response ↓ z ↓ V If z contains the response, then V has access to the action/candidate too.
Strictly speaking, that may not behave like a traditional state-only value function V(s).
It may instead be another pair-conditioned scalar head.
That does not make the architecture useless.
But the notation can overclaim the semantics.
If we truly want state-only V, we could build:
context ↓ state encoder ↓ V while Q receives:
context + response ↓ pair encoder ↓ Q Now the structural distinction is explicit.
32. A stricter state/action architecture For example:
class StrictValueModel(nn.Module): def __init__(self, embedding_dim=1024, hidden_dim=256): super().__init__() self.state_encoder = nn.Sequential( nn.Linear(embedding_dim, hidden_dim), nn.ReLU(), ) self.pair_encoder = nn.Sequential( nn.Linear(embedding_dim * 2, hidden_dim), nn.ReLU(), ) self.v_head = nn.Linear(hidden_dim, 1) self.q_head = nn.Linear(hidden_dim, 1) def forward(self, context, response): zs = self.state_encoder(context) zsa = self.pair_encoder(torch.cat([context, response], dim=-1)) v = self.v_head(zs).squeeze(-1) q = self.q_head(zsa).squeeze(-1) return { "q": q, "v": v, "advantage": q - v, } This is more expensive because the representations are partly separate.
But the semantics are cleaner.
This is exactly the kind of trade-off the series is trying to expose.
33. What does the policy actually condition on? The same question applies to policy.
If the policy sees a representation built from both context and candidate response, then it is answering something like:
Given this context and this candidate, what action should I take?
That can make perfect sense for actions such as:
accept reject revise But it is different from a policy that generates the candidate itself.
Again, the architecture may be valid.
We just need to describe the semantics accurately.
34. Prediction head vs decision rule Suppose the policy probabilities are:
reject 0.20 accept 0.55 revise 0.25 The neural model produced a distribution.
The decision system still needs a rule.
Simplest:
action = action_probs.argmax(dim=-1) But a production system might use thresholds:
accept only if P(accept) > 0.8 otherwise revise or combine signals:
accept if: policy says accept AND Q > threshold AND advantage > 0 The neural model and the decision policy are not necessarily the same thing.
Keeping them separate often makes systems easier to debug.
35. Contradictory heads are useful telemetry Suppose:
Q = high V = low advantage = strongly positive policy = reject That disagreement may indicate:
policy miscalibration; conflicting supervision; representation interference; distribution shift; a bug; genuinely different task definitions. If we collapse everything into one scalar, we cannot see this contradiction.
Multiple heads create more failure modes.
They also create more observability.
That is an underrated benefit.
36. Build a disagreement diagnostic We can define a crude consistency rule.
For example:
accept_prob = out["action_probs"][:, 1] q_prob = torch.sigmoid(out["q_value"]) head_gap = (accept_prob - q_prob).abs() Then inspect the largest gaps.
worst = torch.topk(head_gap, k=10).indices Those samples are often much more informative than random validation examples.
Multi-head models give us internal disagreements to investigate.
37. Do not detach advantage blindly Sometimes code computes:
advantage = (q_value - state_value).detach() That may be correct if advantage is only telemetry or is used as a fixed target/weight elsewhere.
But detaching means downstream losses using advantage will not backpropagate into Q or V through that path.
Compare:
advantage = q_value - state_value with:
advantage = (q_value - state_value).detach() Those are not interchangeable.
The shape is identical.
The autograd graph is not.
This is another example of why understanding the previous PyTorch series matters when reading real architectures.
38. Verify the graph A quick experiment:
adv = out["q_value"] - out["state_value"] print(adv.requires_grad) Expected:
True After detach:
adv_detached = adv.detach() print(adv_detached.requires_grad) Expected:
False If advantage is used only for logging, detach is sensible.
If it participates in a differentiable objective, detach may silently remove learning signal.
39. Test actual parameter updates For a multi-head model, verify every intended component moves.
def snapshot(module): return { name: p.detach().clone() for name, p in module.named_parameters() } before = snapshot(model) Run one training step.
Then:
for name, p in model.named_parameters(): delta = (p.detach() - before[name]).abs().max().item() print(name, delta) If a head never changes:
it may not be in the optimizer; its loss may not be included; gradients may be detached; its target may be constant; its gradients may be numerically tiny. Do not assume because the head exists that it is learning.
40. Tiny-batch overfit test Before a large training run, take perhaps 16 examples.
Train until the model should memorize them.
If:
Q cannot fit V cannot fit policy cannot fit then do not begin hyperparameter sweeps.
Find the bug.
A multi-head tiny-batch test might track:
q_mse v_mse policy_accuracy advantage_mse If Q and V fit but policy does not, you have narrowed the problem dramatically.
41. Save the model contract with the checkpoint Do not save only:
torch.save(model.state_dict(), "ebt.pt") Save the assumptions.
checkpoint = { "state_dict": model.state_dict(), "embedding_dim": D, "hidden_dim": 256, "num_actions": 3, "action_names": ["reject", "accept", "revise"], "q_objective": "reward_regression", "v_objective": "baseline_regression", "embedding_model": "your-embedding-version", } torch.save(checkpoint, "ebt.pt") A checkpoint without its semantic contract is easy to misuse.
42. The model can be replaced piece by piece Because the architecture is compositional, we can change components independently.
Replace:
PairEncoder with a richer interaction encoder.
Keep:
QHead VHead PolicyHead Or widen only Q.
Or remove policy.
Or give V a context-only encoder.
This modularity is not merely clean code.
It makes controlled experiments possible.
43. A richer pair encoder Concatenation forces the network to learn interactions from scratch.
We can expose useful pair relationships directly.
class InteractionEncoder(nn.Module): def __init__(self, embedding_dim, hidden_dim): super().__init__() feature_dim = embedding_dim * 4 self.net = nn.Sequential( nn.Linear(feature_dim, hidden_dim), nn.ReLU(), nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, hidden_dim), ) def forward(self, c, r): features = torch.cat([ c, r, c - r, c * r, ], dim=-1) return self.net(features) Now:
[c, r, c-r, c*r] makes several relations explicit.
Does it help?
Benchmark it.
44. Evaluation matrix At this point we can define a reusable comparison table.
architecture Q V policy advantage params latency MR.Q ✓ — — — ... ... EBT concat ✓ ✓ ✓ ✓ ... ... EBT interactions ✓ ✓ ✓ ✓ ... ... EBT separate V encoder ✓ ✓ ✓ ✓ ... ... And metrics:
Q MSE / ranking accuracy V MSE policy accuracy policy calibration advantage correlation parameter count training time inference latency Now “EBT is better than MR.Q” becomes a testable claim rather than a naming convention.
45. What EBT adds conceptually MR.Q taught us:
pair → representation → scalar EBT adds:
pair ↓ shared representation ↓ multiple specialized views Specifically:
Q candidate quality/value V baseline/state value Policy action preference A Q - V The important innovation is not that any one head is sophisticated.
It is that several cheap models can interrogate one shared representation for different purposes.
46. What EBT does not prove The architecture does not prove:
Q estimates a meaningful return; V is truly state-only; advantage predicts improvement; policy probabilities are calibrated; multi-task learning improves the encoder; more heads outperform MR.Q; the representation captures causal structure; the output is robust out of distribution. Those require evidence.
That is a feature of the methodology, not a criticism of the architecture.
We want to know exactly where the model ends and the empirical claim begins.
47. The next pressure: make the components explicit Our current EBT implementation has the right conceptual pieces, but they are still mostly embedded inside one model definition.
The next architectural move is to make the components first-class:
QHead VHead PolicyHead Encoder with explicit interfaces and independently inspectable behavior.
That sounds like a software-engineering change.
It is also a modelling change.
Once the heads are explicit, we can:
initialize them differently; checkpoint them independently; compare their parameter statistics; train or freeze them separately; replace one without rewriting the others; inspect their gradient flows independently; assign more precise semantics to each interface. That is where SICQL takes us next.
What we learned EBT looks substantially more sophisticated than MR.Q.
But after decomposition it is still straightforward:
PairEncoder ↓ z ┌─┼─────────┐ ↓ ↓ ↓ Q V Policy │ │ │ └─┴→ A └→ probabilities The architecture teaches several reusable ideas:
Shared representations make adding heads cheap. A named head gets its meaning from its objective, not its class name. Derived quantities such as advantage may be arithmetic rather than learned networks. Multi-task learning can help or interfere; measure it. A V head that sees the action is not automatically a pure state value function. Multiple heads provide internal disagreement signals that a scalar model cannot expose. Every extra component should earn its complexity through measurable benefit. And most importantly:
A more advanced model is often just a useful representation plus several smaller models asking different questions of it.
That is the model inside the model.
In the next post we will make those components explicit and build SICQL from first principles: QHead, VHead, PolicyHead and the composition of models into a decision architecture.