Debug a PyTorch Contrastive Loss
Company: Exa
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: hard
Interview Round: Technical Screen
# Debug a PyTorch Contrastive Loss
The source reports a PyTorch contrastive-loss debugging interview centered on matrix-multiplication shape mismatches. It does not provide the exact code or report a normalization-axis defect. The snippet and one-directional in-batch objective below are a practice reconstruction of the reported shape bug.
```python
import torch
import torch.nn.functional as F
def contrastive_loss(z1, z2, temperature):
logits = torch.matmul(z1, z2) / temperature
labels = torch.arange(z1.shape[0], device=z1.device)
return F.cross_entropy(logits, labels)
```
`z1` and `z2` each have shape `[batch_size, embedding_dim]`. Row `i` in `z1` is the positive partner of row `i` in `z2`; the other rows of `z2` are negatives for that anchor. Explain the failed contraction, repair the score matrix, validate the function's assumptions, and discuss numerical precision.
The required output is the mean cross-entropy for `z1` anchors scored against all `z2` rows. A symmetric average with the reverse direction is an optional extension, not part of the required repair. Likewise, L2 normalization is a modeling choice: add it only if the similarity is explicitly defined as cosine similarity.
### Constraints & Assumptions
- Both inputs are finite rank-two floating tensors on the same device, with the same shape and dtype.
- `batch_size > 0` and `embedding_dim > 0`.
- `temperature` is a finite positive Python `int` or `float` and is not a Boolean.
- Positive pairs are aligned by row index.
- For float16 or bfloat16 inputs, form logits and compute loss in float32 while preserving gradient flow. Float32 and float64 inputs may retain their dtype.
- Raise `ValueError` for invalid inputs and `FloatingPointError` if the resulting logits are non-finite.
### Clarifying Questions to Ask
- Is similarity a raw dot product or cosine similarity?
- Are negatives limited to the opposite view?
- Can one anchor have multiple positives?
- Is a one-directional or symmetric objective intended?
- What mixed-precision policy is expected?
### What a Strong Answer Covers
- Correct contraction and explicit shapes for the `[B, B]` score matrix.
- Correct diagonal labels for the stated positive-pair convention.
- Finite-value, device, dtype, shape, and temperature validation.
- Stable cross-entropy use and a defensible low-precision upcast.
- Tests that expose `B != D`, non-finite values, precision overflow, and missing gradients.
### Follow-up Questions
1. How would you add the reverse-direction loss?
2. What changes for multiple positives per anchor?
3. How would negatives from other workers enter the score matrix?
4. When would cosine normalization be preferable to raw dot products?
Quick Answer: Repair a PyTorch contrastive loss whose matrix multiplication contracts incompatible dimensions instead of producing pairwise scores. Validate aligned rank-two tensors and temperature, create the correct batch-by-batch logits, preserve gradients, and handle low-precision numerical stability.