`nx_ltr_lib.nx` source

↩ module page · 225 lines · 9990 B
1// nx_ltr_lib.nx -- LEARNING-TO-RANK reranker as a library so a gate drives it in-process on a
2// synthetic fixture (search rung R0, 2026-09-14). COORDINATE ASCENT (Metzler-Croft) that directly
3// maximises training-fold nDCG@10 one weight at a time by step-halving line search, INITIALISED at
4// the BM25-only solution so it can never rank worse than BM25 on the folds it trains on; evaluated by
5// K-FOLD CROSS-VALIDATION so no query's own label ever trains the weights that score it (the
6// contamination-free path when only a test split exists). Integer/fixed-point end to end: features are
7// min-max normalised to LTR_SCALE per query, weights are i64, ranking is an integer dot product, so
8// the result is bit-reproducible. The harness (nx_beir_eval) stores per-candidate features + gold +
9// idcg during its BM25 pass and calls ltr_cv over the SAME candidate pool, so the learned arm is the
10// same controlled comparison as the semppmi and rrf arms. license_tier: ORIGINAL
11
12import "nx_syscalls.nx"
13
14const LTR_NFEAT: i64 = 10         // bm25 ppmi_maxsim cover covidf sumtf doclen dense_maxsim_centered dense_maxsim minwin bigram
15const LTR_SCALE: i64 = 1000       // per-query per-feature min-max target range (0..LTR_SCALE)
16const LTR_FOLDS: i64 = 5          // 5-fold cross-validation
17const LTR_TOPK: i64 = 10          // nDCG@10
18const LTR_CA_INIT: i64 = 1024     // initial bm25 (feature 0) weight -- ranking starts AS bm25
19const LTR_CA_STEP0: i64 = 1024    // line-search first step, halved to 0; comparable to the init so a feature can reach bm25 scale
20const LTR_CA_PASSES: i64 = 12     // coordinate sweeps; ends early when a whole sweep improves nothing
21
22// nDCG log2 discount, 1-based rank -> round(1000/log2(rank+1)); the metric definition, not a tunable.
23func ltr_disc(r: i64) -> i64 {
  if r == 1 { return 1000 }
  if r == 2 { return 631 }
  if r == 3 { return 500 }
  if r == 4 { return 431 }
  if r == 5 { return 387 }
  if r == 6 { return 356 }
  if r == 7 { return 333 }
  if r == 8 { return 315 }
  if r == 9 { return 301 }
  if r == 10 { return 289 }
  return 0
35}
36
37// per-query per-feature min-max normalise feat -> fn (0..LTR_SCALE). feat and fn are flat arrays
38// indexed (q*maxcand + c)*LTR_NFEAT + k. A feature constant within a query normalises to 0 (no vote),
39// so a slot the harness leaves unfilled contributes nothing rather than a bias. Normalisation is
40// per-query and uses only that query's own features, never a label and never a cross-query statistic,
41// so a held-out query leaks nothing.
42func ltr_norm(feat: *i64, fn: *i64, nq: i64, ncand: *i64, maxcand: i64) -> i64 {
  var q: i64 = 0
  while q < nq {
      let nc: i64 = ncand[q]
      var k: i64 = 0
      while k < LTR_NFEAT {
          var mn: i64 = 0
          var mx: i64 = 0
          var first: i64 = 1
          var c: i64 = 0
          while c < nc {
              let v: i64 = feat[(q * maxcand + c) * LTR_NFEAT + k]
              if first == 1 { mn = v; mx = v; first = 0 } else { if v < mn { mn = v } if v > mx { mx = v } }
              c = c + 1
          }
          let rng: i64 = mx - mn
          c = 0
          while c < nc {
              let idx: i64 = (q * maxcand + c) * LTR_NFEAT + k
              if rng > 0 { fn[idx] = (LTR_SCALE * (feat[idx] - mn)) / rng } else { fn[idx] = 0 }
              c = c + 1
          }
          k = k + 1
      }
      q = q + 1
  }
  return 0
69}
70
71func ltr_dot(w: *i64, fn: *i64, base: i64) -> i64 {
  var s: i64 = 0
  var k: i64 = 0
  while k < LTR_NFEAT { s = s + w[k] * fn[base + k]; k = k + 1 }
  return s
76}
77
78// mean permil nDCG@10 over the queries selected by sel: sel==1 scores the held-out fold (fold==holdout),
79// sel==0 scores the training folds (fold!=holdout). cnt_out[0] = queries scored. holdout<0 with sel==1
80// scores EVERY query (the whole-set control path).
81func ltr_ndcg_sel(w: *i64, fn: *i64, gold: *i64, fold: *i64, nq: i64, ncand: *i64, maxcand: i64, idcg: *i64, holdout: i64, sel: i64, cnt_out: *i64) -> i64 {
  let used: *u8 = sys_mmap(maxcand)
  var sumperm: i64 = 0
  var cnt: i64 = 0
  var q: i64 = 0
  while q < nq {
      var take: i64 = 0
      if sel == 1 { if holdout < 0 { take = 1 } else { if fold[q] == holdout { take = 1 } } }
      else { if fold[q] != holdout { take = 1 } }
      if take == 1 {
          let nc: i64 = ncand[q]
          var c: i64 = 0
          while c < nc { used[c] = 0 as u8; c = c + 1 }
          var dcg: i64 = 0
          var r: i64 = 0
          while r < LTR_TOPK {
              var best: i64 = 0 - 1
              var c2: i64 = 0
              while c2 < nc {
                  if used[c2] == (0 as u8) {
                      if best < 0 { best = c2 } else {
                          if ltr_dot(w, fn, (q * maxcand + c2) * LTR_NFEAT) > ltr_dot(w, fn, (q * maxcand + best) * LTR_NFEAT) { best = c2 }
                      }
                  }
                  c2 = c2 + 1
              }
              if best >= 0 { used[best] = 1 as u8; dcg = dcg + gold[q * maxcand + best] * ltr_disc(r + 1) }
              r = r + 1
          }
          if idcg[q] > 0 { sumperm = sumperm + (dcg * 1000) / idcg[q]; cnt = cnt + 1 }
      }
      q = q + 1
  }
  sys_munmap(used, maxcand)
  cnt_out[0] = cnt
  var m: i64 = 0
  if cnt > 0 { m = sumperm / cnt }
  return m
119}
120
121// COORDINATE ASCENT on training-fold nDCG@10, initialised at bm25-only. For each feature, a step-halving
122// line search moves its weight in whichever direction raises training nDCG, accepting only improvements,
123// so the training-fold score is monotone non-decreasing from the bm25 baseline. w_out gets the weights.
124func ltr_train_ca(fn: *i64, gold: *i64, fold: *i64, nq: i64, ncand: *i64, maxcand: i64, idcg: *i64, holdout: i64, w_out: *i64) -> i64 {
  let w: *i64 = sys_mmap(LTR_NFEAT * 8) as *i64
  let cb: *i64 = sys_mmap(8) as *i64
  var k: i64 = 0
  while k < LTR_NFEAT { w[k] = 0; k = k + 1 }
  w[0] = LTR_CA_INIT
  var cur: i64 = ltr_ndcg_sel(w, fn, gold, fold, nq, ncand, maxcand, idcg, holdout, 0, cb)
  var pass: i64 = 0
  var going: i64 = 1
  while going == 1 {
      var anyimp: i64 = 0
      var kk: i64 = 0
      while kk < LTR_NFEAT {
          var step: i64 = LTR_CA_STEP0
          while step >= 1 {
              let orig: i64 = w[kk]
              w[kk] = orig + step
              var s: i64 = ltr_ndcg_sel(w, fn, gold, fold, nq, ncand, maxcand, idcg, holdout, 0, cb)
              if s > cur { cur = s; anyimp = 1 } else {
                  w[kk] = orig - step
                  s = ltr_ndcg_sel(w, fn, gold, fold, nq, ncand, maxcand, idcg, holdout, 0, cb)
                  if s > cur { cur = s; anyimp = 1 } else { w[kk] = orig; step = step / 2 }
              }
          }
          kk = kk + 1
      }
      pass = pass + 1
      if anyimp == 0 { going = 0 }
      if pass >= LTR_CA_PASSES { going = 0 }
  }
  k = 0
  while k < LTR_NFEAT { w_out[k] = w[k]; k = k + 1 }
  sys_munmap(w as *u8, LTR_NFEAT * 8)
  sys_munmap(cb as *u8, 8)
  return 0
159}
160
161// full k-fold CV. fold[q] must be set by the caller (q mod LTR_FOLDS is the standard, written by the
162// harness). Returns the mean permil nDCG@10 over every query (each held out exactly once).
163// cnt_out[0] = queries scored. This is nr_rerank's measured arm.
164func ltr_cv(feat: *i64, gold: *i64, fold: *i64, nq: i64, ncand: *i64, maxcand: i64, idcg: *i64, cnt_out: *i64) -> i64 {
  let fn: *i64 = sys_mmap(nq * maxcand * LTR_NFEAT * 8) as *i64
  ltr_norm(feat, fn, nq, ncand, maxcand)
  let w: *i64 = sys_mmap(LTR_NFEAT * 8) as *i64
  let c1: *i64 = sys_mmap(8) as *i64
  var totalperm: i64 = 0
  var totalcnt: i64 = 0
  var f: i64 = 0
  while f < LTR_FOLDS {
      ltr_train_ca(fn, gold, fold, nq, ncand, maxcand, idcg, f, w)
      // held-out sum: ltr_ndcg_sel returns a MEAN, so multiply back by its count to sum across folds
      let holdmean: i64 = ltr_ndcg_sel(w, fn, gold, fold, nq, ncand, maxcand, idcg, f, 1, c1)
      totalperm = totalperm + holdmean * c1[0]
      totalcnt = totalcnt + c1[0]
      f = f + 1
  }
  sys_munmap(fn as *u8, nq * maxcand * LTR_NFEAT * 8)
  sys_munmap(w as *u8, LTR_NFEAT * 8)
  sys_munmap(c1 as *u8, 8)
  cnt_out[0] = totalcnt
  var mean: i64 = 0
  if totalcnt > 0 { mean = totalperm / totalcnt }
  return mean
187}
188
189// CONTROL: rank each query by a SINGLE raw feature (no training) and return the mean permil nDCG@10.
190// Called with featk=0 (bm25) it must reproduce the harness's own BM25 arm -- the plumbing check that
191// separates a trainer that learns bad weights from a feature-capture bug. cnt_out[0] = queries scored.
192func ltr_ndcg_single(feat: *i64, gold: *i64, nq: i64, ncand: *i64, maxcand: i64, idcg: *i64, featk: i64, cnt_out: *i64) -> i64 {
  let used: *u8 = sys_mmap(maxcand)
  var sumperm: i64 = 0
  var cnt: i64 = 0
  var q: i64 = 0
  while q < nq {
      let nc: i64 = ncand[q]
      var c: i64 = 0
      while c < nc { used[c] = 0 as u8; c = c + 1 }
      var dcg: i64 = 0
      var r: i64 = 0
      while r < LTR_TOPK {
          var best: i64 = 0 - 1
          var c2: i64 = 0
          while c2 < nc {
              if used[c2] == (0 as u8) {
                  if best < 0 { best = c2 } else {
                      if feat[(q * maxcand + c2) * LTR_NFEAT + featk] > feat[(q * maxcand + best) * LTR_NFEAT + featk] { best = c2 }
                  }
              }
              c2 = c2 + 1
          }
          if best >= 0 { used[best] = 1 as u8; dcg = dcg + gold[q * maxcand + best] * ltr_disc(r + 1) }
          r = r + 1
      }
      if idcg[q] > 0 { sumperm = sumperm + (dcg * 1000) / idcg[q]; cnt = cnt + 1 }
      q = q + 1
  }
  sys_munmap(used, maxcand)
  cnt_out[0] = cnt
  var m: i64 = 0
  if cnt > 0 { m = sumperm / cnt }
  return m
225}
nx_ltr_lib.nx source

`nx_ltr_lib.nx` source
