Models From First Principles 01: MR.Q — Building a Neural Quality Model From Two Embeddings
MR.Q — Building a Neural Quality Model From Two Embeddings
In the previous post, we established the core idea behind this series:
A complicated model becomes understandable when you recursively decompose it into smaller models, blocks, layers and tensor operations.
Now we build the first real model.
Not a transformer.
Not a giant language model.
Not an agent.
A scorer.
We will take two embeddings:
context embedding
response embedding
combine them, encode the relationship between them, and predict one scalar:
Q(context, response)
That is the entire model.
And that simplicity is exactly why it is a useful place to start.
The architecture is inspired by the MR.Q model used in Stephanie, but everything in this article is standalone PyTorch. You do not need Stephanie, and you do not need any of the later models in this series.
The point is to understand what a quality model actually is before we add value heads, policy heads, recurrence, uncertainty, sparse representations or specialized optimizers.
1. The problem: score a response in context
Suppose we have a prompt:
Explain why gradient accumulation is useful.
and a candidate response:
Gradient accumulation allows several smaller micro-batches to contribute
before an optimizer step, approximating a larger effective batch when memory
is limited.
We want a model that produces something like:
0.83
or perhaps a raw score such as:
2.41
The important point is that the model consumes two things and returns one thing:
(context, response) -> scalar
This is not yet a language model.
It does not generate text.
It evaluates a pair.
That makes it useful for:
- ranking candidate responses;
- preference modelling;
- reranking retrieval results;
- reward modelling;
- filtering generated outputs;
- critic models;
- search over candidate plans;
- model comparison;
- selecting among several LLM generations.
The architecture can be tiny because another system may already have converted the text into embeddings.
2. The model in one diagram
Here is the complete idea:
context text response text
| |
v v
context embedding response embedding
| |
+----------------+-------------------+
|
v
pair encoder
|
v
joint representation
|
v
predictor
|
v
Q value
If each embedding has dimension D, then the simplest implementation is:
[B, D] + [B, D]
|
v
concatenate
|
v
[B, 2D]
|
v
MLP
|
v
[B, hidden]
|
v
scalar head
|
v
[B]
That is already enough to train a useful pair scorer.
3. Start with fake embeddings
We are deliberately separating text encoding from pair scoring.
For now, imagine another model has already turned the context and response into vectors.
import torch
B = 4
D = 8
context_emb = torch.randn(B, D)
response_emb = torch.randn(B, D)
print(context_emb.shape)
print(response_emb.shape)
Output:
torch.Size([4, 8])
torch.Size([4, 8])
Our contract is:
context_emb: [B, D]
response_emb: [B, D]
Same batch size.
Same embedding dimension.
That is the first invariant worth asserting.
def assert_pair_contract(context_emb, response_emb):
assert context_emb.ndim == 2
assert response_emb.ndim == 2
assert context_emb.shape == response_emb.shape
Then:
assert_pair_contract(context_emb, response_emb)
A large fraction of neural-network bugs are easier when the expected shape is written down before the architecture.
4. The smallest possible pair model
Before building an encoder, let us create an almost embarrassingly simple model.
Concatenate the two embeddings and apply one linear layer.
from torch import nn
class LinearPairScorer(nn.Module):
def __init__(self, embedding_dim):
super().__init__()
self.scorer = nn.Linear(embedding_dim * 2, 1)
def forward(self, context_emb, response_emb):
pair = torch.cat([context_emb, response_emb], dim=-1)
return self.scorer(pair).squeeze(-1)
Use it:
model = LinearPairScorer(D)
score = model(context_emb, response_emb)
print(score.shape)
Output:
torch.Size([4])
This is a real model.
It has trainable parameters.
It can learn a scalar function over pairs of embeddings.
But it has an important limitation.
It is linear.
The model cannot easily learn richer interactions between dimensions of the context and dimensions of the response.
So we insert an encoder.
5. The model inside the model
Now the architecture becomes:
PairQualityModel
|
+-- PairEncoder
|
+-- Predictor
That is the first example of the central theme of this series.
The model is made out of models.
Let us implement the pair encoder independently.
class PairEncoder(nn.Module):
def __init__(self, embedding_dim, hidden_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(embedding_dim * 2, hidden_dim),
nn.ReLU(),
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)
Test it:
encoder = PairEncoder(
embedding_dim=8,
hidden_dim=16,
)
z = encoder(context_emb, response_emb)
print(z.shape)
Output:
torch.Size([4, 16])
The two inputs have become one joint representation:
z = encoder(context, response)
This representation is the interesting part of the model.
The predictor that follows it can now be extremely small.
6. The predictor is another model
class ScalarPredictor(nn.Module):
def __init__(self, hidden_dim):
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)
Again:
Predictor
|
+-- Linear
+-- ReLU
+-- Linear
There is no magic left.
The predictor is just a function from:
[B, hidden]
to:
[B]
7. Compose the full MR.Q-style model
class MRQStyleModel(nn.Module):
def __init__(self, embedding_dim=1024, hidden_dim=256):
super().__init__()
self.encoder = PairEncoder(
embedding_dim=embedding_dim,
hidden_dim=hidden_dim,
)
self.predictor = ScalarPredictor(
hidden_dim=hidden_dim,
)
def forward(self, context_emb, response_emb):
z = self.encoder(context_emb, response_emb)
q_value = self.predictor(z)
return q_value
That is the complete architecture.
context [B,D] ----+
|
+--> PairEncoder --> z [B,H] --> Predictor --> q [B]
|
response [B,D] ---+
Use it:
model = MRQStyleModel(
embedding_dim=8,
hidden_dim=16,
)
q = model(context_emb, response_emb)
print(q)
print(q.shape)
8. What exactly is Q?
This is where terminology can become misleading.
The architecture outputs a scalar called q_value.
That name does not magically make it reinforcement learning.
The meaning of the scalar depends on the training objective and the data.
For example, the same architecture could learn:
human preference probability
or:
expected reward
or:
quality score
or:
relevance
or:
probability of acceptance
or even:
estimated downstream utility
The network architecture only says:
pair -> scalar
The objective gives that scalar meaning.
This distinction matters throughout this entire series.
A head named uncertainty_head does not become uncertainty because we named it that.
A head named reasoning_score does not prove reasoning.
A scalar becomes meaningful because of:
data + labels/targets + objective + evaluation
9. Option A: regression
Suppose each pair has a target quality score between 0 and 1.
quality = torch.tensor([
0.9,
0.2,
0.7,
0.4,
])
We can train with mean-squared error.
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
pred = model(context_emb, response_emb)
loss = torch.nn.functional.mse_loss(pred, quality)
optimizer.zero_grad()
loss.backward()
optimizer.step()
Now q_value is being trained as a regression estimate of the quality target.
10. Option B: binary preference or acceptance
Suppose instead we have labels:
0 = rejected
1 = accepted
Then we can treat the raw scalar as a logit.
labels = torch.tensor([
1.0,
0.0,
1.0,
0.0,
])
logits = model(context_emb, response_emb)
loss = torch.nn.functional.binary_cross_entropy_with_logits(
logits,
labels,
)
At inference time:
probability = torch.sigmoid(logits)
The same architecture now has a different interpretation.
This is why we should avoid baking sigmoid() into the model too early.
Raw logits are generally more flexible.
11. Option C: pairwise ranking
This is particularly interesting for candidate selection.
Suppose for the same context we have:
response A = preferred
response B = rejected
We compute:
q_a = model(context_emb, response_a_emb)
q_b = model(context_emb, response_b_emb)
We want:
q_a > q_b
A simple ranking loss is:
loss = -torch.log(torch.sigmoid(q_a - q_b)).mean()
Equivalent intuition:
preferred score - rejected score
should be positive.
Now the model is not being trained to recover an absolute number.
It is being trained to order candidates.
That can be exactly what we need.
12. A reusable pairwise ranking loss
def pairwise_logistic_loss(preferred_score, rejected_score):
margin = preferred_score - rejected_score
return torch.nn.functional.softplus(-margin).mean()
Why softplus(-margin)?
Because:
softplus(-x) = -log(sigmoid(x))
numerically stably.
Use it:
q_preferred = model(context_emb, preferred_emb)
q_rejected = model(context_emb, rejected_emb)
loss = pairwise_logistic_loss(
q_preferred,
q_rejected,
)
This is already enough to build a basic preference model.
13. Create a synthetic training problem
Let us make the model learn something we can verify.
We will define synthetic quality as similarity between the context and response embeddings.
import torch
import torch.nn.functional as F
def make_dataset(n=5000, d=32):
context = torch.randn(n, d)
response = torch.randn(n, d)
target = F.cosine_similarity(
context,
response,
dim=-1,
)
return context, response, target
This target lies approximately in:
[-1, 1]
Create data:
context, response, target = make_dataset()
print(context.shape)
print(response.shape)
print(target.shape)
14. Train the model
torch.manual_seed(7)
model = MRQStyleModel(
embedding_dim=32,
hidden_dim=64,
)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
)
for step in range(1000):
idx = torch.randint(0, len(context), (128,))
c = context[idx]
r = response[idx]
y = target[idx]
pred = model(c, r)
loss = F.mse_loss(pred, y)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
if step % 100 == 0:
print(step, float(loss))
This is now a complete learning system.
Nothing in it requires a language model at training time.
If embeddings have already been computed, the scorer can be tiny and cheap.
15. Validate it properly
Never evaluate on the same examples used for training.
train_c = context[:4000]
train_r = response[:4000]
train_y = target[:4000]
val_c = context[4000:]
val_r = response[4000:]
val_y = target[4000:]
Then:
model.eval()
with torch.no_grad():
val_pred = model(val_c, val_r)
val_loss = F.mse_loss(val_pred, val_y)
print("val loss:", float(val_loss))
The purpose of the validation set is not ceremony.
It answers:
Did the model learn the relationship, or did it merely fit the examples it saw?
16. Ranking accuracy can be more useful than MSE
If the real application is candidate ranking, MSE may not be the most informative metric.
Suppose we compare two responses for each context.
We can test whether the model gives the better response the higher score.
def ranking_accuracy(q_good, q_bad):
return (q_good > q_bad).float().mean()
For a ranking system, this may matter more than absolute calibration.
That leads to an important rule:
Evaluate the model according to the decision it is actually used to make.
17. Why concatenate the embeddings?
The simplest pair encoder uses:
pair = torch.cat([
context_emb,
response_emb,
], dim=-1)
If both are [B, D], the result is:
[B, 2D]
Concatenation preserves both vectors separately.
The following linear layers can learn how dimensions from one side interact with dimensions from the other.
But concatenation is not the only possible pair representation.
18. Difference features
We could add:
difference = response_emb - context_emb
Then:
pair = torch.cat([
context_emb,
response_emb,
response_emb - context_emb,
], dim=-1)
Shape:
[B, 3D]
This gives the network an explicit directional difference feature.
19. Elementwise product
Another common interaction feature is:
product = context_emb * response_emb
Then:
pair = torch.cat([
context_emb,
response_emb,
context_emb - response_emb,
context_emb * response_emb,
], dim=-1)
Shape:
[B, 4D]
The product exposes dimension-wise agreement directly.
20. An interaction-rich encoder
class InteractionPairEncoder(nn.Module):
def __init__(self, embedding_dim, hidden_dim):
super().__init__()
input_dim = embedding_dim * 4
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.LayerNorm(hidden_dim),
nn.Linear(hidden_dim, hidden_dim),
)
def forward(self, context_emb, response_emb):
features = torch.cat([
context_emb,
response_emb,
context_emb - response_emb,
context_emb * response_emb,
], dim=-1)
return self.net(features)
Now we have changed the architecture while keeping the model interface the same.
That is powerful.
(context_emb, response_emb) -> z
The rest of the system does not need to know how z was created.
21. Stable interfaces let us replace components
Suppose we define:
class QualityModel(nn.Module):
def __init__(self, encoder, predictor):
super().__init__()
self.encoder = encoder
self.predictor = predictor
def forward(self, context_emb, response_emb):
z = self.encoder(context_emb, response_emb)
return self.predictor(z)
Now the architecture becomes composable.
model = QualityModel(
encoder=InteractionPairEncoder(32, 64),
predictor=ScalarPredictor(64),
)
Later we could swap the encoder for:
- a deeper MLP;
- cross-attention;
- bilinear interaction;
- recurrent fusion;
- a transformer;
- a learned projection over multiple candidate features.
The predictor does not care.
This is architectural decomposition in practice.
22. Count the parameters
A model can sound sophisticated while being tiny.
Let us prove how many parameters we actually have.
def count_parameters(model):
return sum(
p.numel()
for p in model.parameters()
if p.requires_grad
)
Example:
model = MRQStyleModel(
embedding_dim=1024,
hidden_dim=256,
)
print(count_parameters(model))
The exact number depends on the chosen predictor dimensions, but the important point is that this is still a small network compared with the embedding model that may have produced the inputs.
The scorer can therefore be trained, copied, evaluated and iterated much more cheaply than a full LLM.
23. Where are the expensive parameters?
For the first encoder layer:
Linear(2D -> H)
parameter count is approximately:
2D * H + H
For:
D = 1024
H = 256
that layer alone contains roughly:
2048 * 256
= 524,288 weights
plus bias.
This immediately tells us something useful:
The first projection dominates a large part of the model’s parameter budget.
When reading architecture code, parameter arithmetic often reveals more than the class names.
24. Parameter count helper by module
def parameter_report(model):
rows = []
for name, module in model.named_modules():
own_params = sum(
p.numel()
for p in module.parameters(recurse=False)
)
if own_params:
rows.append((name, own_params))
return rows
Then:
for name, n in parameter_report(model):
print(f"{name:30s} {n:>10,d}")
This is one of the simplest ways to understand where a model actually spends capacity.
25. Inspect the representation
The model’s output is a scalar.
But the encoder’s latent representation z is often more interesting.
with torch.no_grad():
z = model.encoder(context_emb, response_emb)
print(z.shape)
print(z.mean().item())
print(z.std().item())
Why inspect it?
Because a bad predictor may be easy to replace.
A collapsed representation is a deeper problem.
If z is nearly constant for all pairs, the predictor cannot recover useful distinctions.
26. Measure representation variance
def representation_stats(z):
return {
"mean": float(z.mean()),
"std": float(z.std()),
"min": float(z.min()),
"max": float(z.max()),
"mean_feature_std": float(z.std(dim=0).mean()),
}
Use it:
print(representation_stats(z))
Low variance does not automatically prove collapse, but it is a useful diagnostic clue.
27. Prove gradients reach both components
After backward:
loss.backward()
we can inspect which parts of the model received gradients.
def gradient_report(model):
for name, p in model.named_parameters():
if p.grad is None:
print(name, "NO GRAD")
else:
print(
name,
"grad_norm=",
float(p.grad.norm()),
)
A pair scorer that runs successfully but never updates the encoder is not learning the model we think it is learning.
28. Prove parameters actually change
This is a pattern worth carrying over from the PyTorch debugging series.
before = {
name: p.detach().clone()
for name, p in model.named_parameters()
}
loss.backward()
optimizer.step()
for name, p in model.named_parameters():
changed = not torch.equal(before[name], p.detach())
print(name, changed)
Generated training code may look correct while a parameter group, detached tensor or optimizer construction error prevents real updates.
The parameter delta is evidence.
29. A very common bug: sigmoid too early
Suppose the model does this:
return torch.sigmoid(self.predictor(z))
and the training code uses:
binary_cross_entropy_with_logits(...)
That is wrong.
binary_cross_entropy_with_logits() expects raw logits and applies the appropriate sigmoid internally in a numerically stable way.
Correct:
logits = model(context_emb, response_emb)
loss = F.binary_cross_entropy_with_logits(
logits,
labels,
)
Inference:
prob = torch.sigmoid(logits)
Keep model output semantics explicit.
30. Another common bug: .squeeze() removes the batch dimension
This looks harmless:
q = self.predictor(z).squeeze()
But when batch size is one:
[B, 1] = [1, 1]
plain .squeeze() can return a scalar tensor:
[]
instead of:
[1]
Safer:
q = self.predictor(z).squeeze(-1)
Now only the final singleton dimension is removed.
31. Shape test for batch size one
def test_batch_one_shape():
model = MRQStyleModel(embedding_dim=32, hidden_dim=64)
c = torch.randn(1, 32)
r = torch.randn(1, 32)
q = model(c, r)
assert q.shape == (1,)
This tiny test can prevent surprisingly annoying downstream bugs.
32. Device contract
If the model is on CUDA but the embeddings are still on CPU:
Expected all tensors to be on the same device
A clean training loop moves data explicitly:
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
model = model.to(device)
context_batch = context_batch.to(device)
response_batch = response_batch.to(device)
target_batch = target_batch.to(device)
Avoid hiding excessive device movement inside low-level model components unless you have a strong reason.
The caller usually knows where tensors should live.
33. Dtype contract
Embeddings are normally floating point.
assert context_emb.dtype.is_floating_point
assert response_emb.dtype.is_floating_point
If one input unexpectedly arrives as integer token IDs, that is not a small dtype problem.
It means the wrong representation entered the model.
This is why shape and semantic contracts matter.
34. Build a batch scorer API
@torch.no_grad()
def score_pairs(model, context_emb, response_emb):
model.eval()
return model(context_emb, response_emb)
If the training objective used logits and we want probabilities:
@torch.no_grad()
def score_probabilities(model, context_emb, response_emb):
logits = model(context_emb, response_emb)
return torch.sigmoid(logits)
Keep the raw model prediction separate from presentation semantics.
35. Ranking many candidates for one context
This is a common real use case.
One context:
[D]
Many candidate responses:
[N, D]
Repeat the context across candidates:
def rank_candidates(model, context_emb, candidate_embs):
if context_emb.ndim == 1:
context_emb = context_emb.unsqueeze(0)
context_batch = context_emb.expand(
candidate_embs.shape[0],
-1,
)
with torch.no_grad():
scores = model(
context_batch,
candidate_embs,
)
return torch.argsort(
scores,
descending=True,
), scores
This turns the quality model into a reranker.
36. Top-k selection
order, scores = rank_candidates(
model,
context_emb,
candidate_embs,
)
top_k = order[:3]
Now a large generator can produce ten candidates and a tiny scorer can choose three.
This is a fundamentally different architecture from asking one giant model to do everything.
37. Why use embeddings at all?
Why not feed raw text into the scorer?
Because embeddings give us separation of responsibilities.
text model
|
v
semantic representation
|
v
small task-specific scorer
This can provide:
- cheaper retraining;
- smaller datasets;
- faster iteration;
- lower inference cost;
- easier experimentation;
- independent versioning of representation and scoring layers.
But it also creates a dependency.
The scorer can only learn from information preserved by the embedding model.
If the embedding destroys a distinction important for quality, the scorer cannot reconstruct it.
38. Frozen representation vs end-to-end training
There are two broad designs.
Frozen embedding model
text
↓
fixed encoder
↓
embedding
↓
trainable scorer
Advantages:
- cheap;
- stable;
- reproducible;
- easy to cache.
End-to-end model
text
↓
trainable encoder
↓
embedding
↓
trainable scorer
Advantages:
- representation can adapt to the scoring task.
Costs:
- more parameters;
- more memory;
- more training complexity;
- easier overfitting;
- harder deployment.
Neither is universally correct.
39. Cached embeddings change the economics
Suppose we have one million training pairs.
If embeddings are fixed, we can compute them once and cache them.
Then training the scorer repeatedly becomes a small matrix-learning problem rather than repeatedly invoking an expensive text encoder.
This changes experimentation dramatically.
You can try:
- different hidden sizes;
- different losses;
- different pair features;
- different optimizers;
- calibration strategies;
- uncertainty heads;
- ranking objectives;
without recomputing the expensive representation every time.
40. But beware embedding version drift
If embeddings are cached, their identity becomes part of the dataset.
You need to know:
embedding model
embedding model version
normalization
pooling method
embedding dimension
preprocessing
A scorer trained on one embedding space cannot safely be fed arbitrary vectors from another embedding model merely because the shape is also [1024].
Same shape does not mean same semantics.
41. Save the scorer
Because the model is an ordinary nn.Module:
torch.save(
model.state_dict(),
"mrq_style.pt",
)
Load it:
model = MRQStyleModel(
embedding_dim=32,
hidden_dim=64,
)
state = torch.load(
"mrq_style.pt",
map_location="cpu",
)
model.load_state_dict(state)
But the checkpoint alone is not enough.
You also need the architecture configuration and embedding contract.
42. Save configuration with the checkpoint
checkpoint = {
"model_state": model.state_dict(),
"embedding_dim": 32,
"hidden_dim": 64,
"embedding_model": "example-encoder-v1",
"objective": "pairwise_logistic",
}
torch.save(checkpoint, "quality_model.pt")
Now the artifact carries more of its own meaning.
43. Calibration is not ranking
A model can rank responses correctly but produce poorly calibrated probabilities.
For example:
predicted probability = 0.95
should ideally correspond to an event that occurs about 95% of the time under the relevant distribution.
Ranking asks:
A > B ?
Calibration asks:
Does 0.8 actually behave like 80%?
These are different properties.
Later models in this series introduce explicit calibration-related machinery, but it is useful to understand the distinction now.
44. A baseline can beat a neural network
Before celebrating the model, compare it against simpler alternatives.
If embeddings are normalized, cosine similarity itself may be a strong baseline.
baseline = F.cosine_similarity(
context_emb,
response_emb,
dim=-1,
)
If the neural model does not outperform a trivial baseline on the real task, architectural sophistication has not bought us anything.
That is an important discipline for this entire series.
45. Baseline ladder
For a pair scorer, I would compare at least:
1. random
2. cosine similarity
3. linear scorer
4. one-hidden-layer MLP
5. deeper pair encoder
6. richer interaction features
Only move upward when evidence justifies it.
This turns architecture into an experiment rather than an aesthetic preference.
46. Ablate the interaction features
Suppose the richer encoder uses:
context
response
difference
product
Train variants:
A: context + response
B: context + response + difference
C: context + response + product
D: all four
Then compare validation ranking accuracy.
If the extra features do not help, remove them.
The model should earn its complexity.
47. Hidden width is a hypothesis
Try:
32
64
128
256
512
Do not assume bigger is better.
For a fixed embedding representation and modest dataset, a smaller pair model may generalize better.
Track:
parameter count
training loss
validation loss
ranking accuracy
latency
Now hidden width is a measurable trade-off.
48. A complete small implementation
Here is a compact standalone version.
import torch
import torch.nn as nn
import torch.nn.functional as F
class PairEncoder(nn.Module):
def __init__(self, embedding_dim, hidden_dim):
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),
nn.ReLU(),
)
def forward(self, context_emb, response_emb):
assert context_emb.shape == response_emb.shape
pair = torch.cat([context_emb, response_emb], dim=-1)
return self.net(pair)
class ScalarPredictor(nn.Module):
def __init__(self, hidden_dim):
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)
class MRQStyleModel(nn.Module):
def __init__(self, embedding_dim=1024, hidden_dim=256):
super().__init__()
self.encoder = PairEncoder(
embedding_dim,
hidden_dim,
)
self.predictor = ScalarPredictor(
hidden_dim,
)
def forward(self, context_emb, response_emb):
z = self.encoder(context_emb, response_emb)
q = self.predictor(z)
return q
Nothing hidden.
Nothing pretrained inside the scorer.
Two vectors enter.
One number leaves.
49. Complete pairwise training example
def pairwise_loss(q_preferred, q_rejected):
return F.softplus(
-(q_preferred - q_rejected)
).mean()
model = MRQStyleModel(
embedding_dim=32,
hidden_dim=64,
)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
)
for step in range(1000):
context = torch.randn(128, 32)
preferred = context + 0.2 * torch.randn(128, 32)
rejected = torch.randn(128, 32)
q_good = model(context, preferred)
q_bad = model(context, rejected)
loss = pairwise_loss(q_good, q_bad)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
if step % 100 == 0:
with torch.no_grad():
acc = (q_good > q_bad).float().mean()
print(
step,
"loss=",
float(loss),
"ranking_acc=",
float(acc),
)
This toy problem has a known rule:
The preferred response embedding is constructed to resemble the context more closely.
The model should learn that ordering.
50. Tiny-batch overfit test
Before training on millions of examples, prove the model can fit eight.
context = torch.randn(8, 32)
preferred = context + 0.05 * torch.randn(8, 32)
rejected = torch.randn(8, 32)
Train repeatedly on those same eight pairs.
If the model cannot drive training ranking accuracy close to 100%, something may be wrong with:
- the objective;
- gradient flow;
- optimizer membership;
- target construction;
- architecture;
- learning rate.
This test is extremely cheap and extremely informative.
51. Diagnostic function
@torch.no_grad()
def diagnose_pair_model(model, context_emb, response_emb):
model.eval()
assert context_emb.ndim == 2
assert response_emb.ndim == 2
assert context_emb.shape == response_emb.shape
z = model.encoder(context_emb, response_emb)
q = model.predictor(z)
print("context:", tuple(context_emb.shape))
print("response:", tuple(response_emb.shape))
print("latent:", tuple(z.shape))
print("q:", tuple(q.shape))
print("latent mean:", float(z.mean()))
print("latent std:", float(z.std()))
print("q mean:", float(q.mean()))
print("q std:", float(q.std()))
Use it before blaming the optimizer.
52. What MR.Q does not give us
This model is intentionally limited.
It gives us one scalar.
That means it does not inherently distinguish:
How good is this action?
from:
How good is this state generally?
It does not produce an explicit policy distribution.
It does not estimate disagreement.
It does not expose uncertainty.
It does not reason recursively.
It does not maintain recurrent state.
It does not know when to halt.
Those limitations are not defects in this post.
They are the reason the next architectures exist.
53. The architectural pressure that leads to EBT
Imagine we have a scalar Q score:
Q(context, response) = 0.82
That tells us this pair is good according to the learned objective.
But suppose we also want to know:
What quality would we expect from this state regardless of this action?
That suggests a second quantity:
V(state)
Then:
Advantage = Q - V
Now we can ask:
Is this candidate merely good because the situation is easy, or is it better than what we normally expect here?
And perhaps we also want a distribution over actions:
Policy(action | state)
Now our one-headed scorer naturally wants to become a multi-headed model.
That is where EBT enters.
54. The evolution in one picture
MR.Q-style scorer:
pair
|
v
encoder
|
v
Q
The next architecture:
shared representation
|
+-----------+-----------+
| | |
v v v
Q V Policy
| |
+-----+-----+
|
v
Advantage
Notice what happened.
We did not throw away the original idea.
We kept the shared representation and expanded what we ask of it.
That pattern will repeat throughout this series.
55. What to remember
If you remember only a few things from this post, remember these.
A quality model can be tiny
If text has already been embedded, pair scoring may require only a small task-specific network.
The encoder is often more important than the scalar head
The predictor can only work with the representation it receives.
Architecture does not define semantics by itself
Q means whatever the training objective and data make it mean.
Ranking and calibration are different problems
A good ordering model is not automatically a good probability model.
Compare against trivial baselines
Cosine similarity might already solve more of the task than expected.
Preserve stable interfaces
If the encoder always maps:
(context, response) -> z
we can replace it without rewriting the whole system.
Prove complexity earns its place
Use ablations, parameter counts, validation metrics and runtime measurements.
56. The deeper lesson
We started with something that sounds abstract:
MR.Q quality model
Then we decomposed it:
MR.Q-style model
|
+-- pair encoder
| |
| +-- concatenate
| +-- Linear
| +-- ReLU
| +-- LayerNorm
| +-- Linear
|
+-- scalar predictor
|
+-- Linear
+-- ReLU
+-- Linear
Then even those layers reduce to tensor operations we already understand from the PyTorch series.
This is the method we will keep using.
When a model becomes intimidating, descend one level.
If that level is still intimidating, descend again.
Eventually you reach:
matrix multiplication
addition
normalization
activation
state update
loss
At that level there is very little magic left.
Next: EBT — From One Score to Q, V, Policy and Advantage
The next post starts from the exact limitation we have reached here.
One scalar is useful.
But one scalar cannot tell us everything we want to know about a decision.
We will keep the shared representation and add three independent prediction surfaces:
Q
V
Policy
Then we will derive:
Advantage = Q - V
And once again, we will build each component independently before composing the complete model.
