Models From First Principles 02: EBT — From One Score to Q, V, Policy and Advantage
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.82is 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.
