promote metric0.3

bridge2vec.py

@@ -242,6 +242,17 @@ def schedule(step, warmup, total):
     return 0.5 * (1 + math.cos(math.pi * (step - warmup) / max(1, total - warmup)))
 
 
+def metric_loss(z, value, temperature, target_temperature):
+    """Make hands that play alike close: soft targets from the distance between expected tables."""
+    z = z.float().flatten(0, 1)
+    value = value.detach().float().flatten(0, 1).flatten(1)
+    self_pair = torch.eye(len(z), dtype=torch.bool, device=z.device)
+    distance = torch.cdist(value, value, p=1) / value.shape[1]
+    target = (-distance / target_temperature).masked_fill(self_pair, -1e9).softmax(1)
+    logits = (z @ z.T / temperature).masked_fill(self_pair, -1e9)
+    return -(target * logits.log_softmax(1)).sum(1).mean()
+
+
 def push_result(repo, name, result, revision=None):
     """Upload a result as results/<name>.json in the model repo."""
     HfApi().upload_file(
@@ -277,12 +288,14 @@ def train(args):
         batch = torch.randint(len(cards), (args.batch_size,), device=device)
         hands, tricks = permute_suits(cards[batch], dd[batch])
         with torch.autocast(device, dtype=torch.bfloat16, enabled=device == "cuda"):
-            _, logits, value, shape = model(hands)
+            z, logits, value, shape = model(hands)
         expected = torch.stack([tricks.roll(-d, 2) for d in range(4)], 1).float()
         table_loss = functional.cross_entropy(logits.float().flatten(0, 2), tricks.flatten())
         value_loss = functional.mse_loss(value.float(), expected)
         shape_loss = functional.mse_loss(shape.float(), profile(hands) / scale)
-        loss = table_loss + args.value_weight * value_loss + args.aux_weight * shape_loss
+        metric = metric_loss(z, value, args.temperature, args.target_temperature)
+        loss = (table_loss + args.value_weight * value_loss + args.aux_weight * shape_loss
+                + args.metric_weight * metric)
         optimizer.zero_grad(set_to_none=True)
         loss.backward()
         torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
@@ -295,6 +308,7 @@ def train(args):
                 "table": round(table_loss.item(), 4),
                 "value": round(value_loss.item(), 4),
                 "shape": round(shape_loss.item(), 4),
+                "metric": round(metric.item(), 4),
                 "exact": round((logits.argmax(-1) == tricks).float().mean().item(), 3),
                 "minutes": round((time.time() - start) / 60, 1),
             }
@@ -331,6 +345,9 @@ def train(args):
                 "learning_rate": args.lr,
                 "value_weight": args.value_weight,
                 "aux_weight": args.aux_weight,
+                "metric_weight": args.metric_weight,
+                "temperature": args.temperature,
+                "target_temperature": args.target_temperature,
                 "dim": args.dim,
                 "depth": args.depth,
                 "embed_dim": args.embed_dim,
@@ -560,6 +577,10 @@ def main():
     train_parser.add_argument("--warmup", type=int, default=1000)
     train_parser.add_argument("--value-weight", type=float, default=0.1)
     train_parser.add_argument("--aux-weight", type=float, default=0.1)
+    train_parser.add_argument("--metric-weight", type=float, default=0.0,
+                              help="weight of the loss that ties embedding to table distance")
+    train_parser.add_argument("--temperature", type=float, default=0.1)
+    train_parser.add_argument("--target-temperature", type=float, default=0.1)
     train_parser.add_argument("--dim", type=int, default=256)
     train_parser.add_argument("--depth", type=int, default=4)
     train_parser.add_argument("--embed-dim", type=int, default=128)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e7d5c6b469fe2e36cc221144fad60a221da1ca3e3f0d21eab41068607180cb90
+oid sha256:c788114fe205ccd973888f9530318853b73360f3f9443bd53ae176972702b3ea
 size 19656236

results/eval.json

@@ -2,36 +2,36 @@
  "model": "jgalego/bridge2vec",
  "tables": {
   "deals": 32000,
-  "mae": 0.366,
-  "exact": 0.654,
-  "within_one": 0.984,
-  "table_exact": 0.016,
+  "mae": 0.398,
+  "exact": 0.628,
+  "within_one": 0.98,
+  "table_exact": 0.012,
   "mae_by_strain": {
-   "S": 0.34,
-   "H": 0.343,
-   "D": 0.344,
-   "C": 0.342,
-   "N": 0.46
+   "S": 0.368,
+   "H": 0.371,
+   "D": 0.369,
+   "C": 0.368,
+   "N": 0.513
   }
  },
  "hands": {
   "hands": 1000,
   "deals_per_hand": 32.0,
-  "expected_mae": 0.345,
+  "expected_mae": 0.346,
   "constant_mae": 1.251,
   "retrieval": {
-   "embedding": 0.836,
-   "hcp_shape": 0.613,
-   "random": 1.752
+   "embedding": 0.55,
+   "hcp_shape": 0.616,
+   "random": 1.781
   },
   "hard_pairs": {
    "triplets": 98216,
-   "accuracy": 0.769
+   "accuracy": 0.901
   }
  },
  "deal_retrieval": {
-  "embedding": 1.788,
-  "hcp_shape": 1.411,
-  "random": 3.131
+  "embedding": 1.277,
+  "hcp_shape": 1.412,
+  "random": 3.135
  }
 }

results/train.json

@@ -1,11 +1,12 @@
 {
  "step": 49999,
- "loss": 1.2938,
- "table": 0.7635,
- "value": 5.3025,
- "shape": 0.0006,
- "exact": 0.667,
- "minutes": 80.6,
+ "loss": 2.7616,
+ "table": 0.8088,
+ "value": 5.3021,
+ "shape": 0.0011,
+ "metric": 4.7414,
+ "exact": 0.648,
+ "minutes": 117.4,
  "data": "jgalego/bridge2vec-deals",
  "deals": 200000,
  "steps": 50000,
@@ -13,11 +14,14 @@
  "learning_rate": 0.0003,
  "value_weight": 0.1,
  "aux_weight": 0.1,
+ "metric_weight": 0.3,
+ "temperature": 0.1,
+ "target_temperature": 0.1,
  "dim": 256,
  "depth": 4,
  "embed_dim": 128,
  "hidden": 1024,
  "parameters": 4912479,
- "runtime_s": 4841,
+ "runtime_s": 7050,
  "device": "NVIDIA A10G"
 }