Teaching Large Models to Memorize Documents

# Illustrative sketch (not from the paper)
import torch
import torch.nn as nn

# Placeholder model (instruction-tuned)
model = nn.Linear(768, 768)

# ---- Inject stage ----
# Three document-generation objectives: continuation, rewrite, instruction reconstruction
def inject_training(model, docs, mix=(1,1,1)):
    # mix defines relative weighting of the three objectives (e.g., 1:1:1)
    # Here we simply loop over docs and compute a dummy loss per objective
    loss = 0.0
    for doc in docs:
        # continuation loss (placeholder)
        loss += torch.mean(model(torch.randn(1,768)))
        # rewrite loss (placeholder)
        loss += torch.mean(model(torch.randn(1,768)))
        # instruction reconstruction loss (placeholder)
        loss += torch.mean(model(torch.randn(1,768)))
    return loss

# ---- Align stage ----
# Fine‑tune on answer‑only QA pairs
def align_training(model, qa_pairs):
    loss = 0.0
    for q, a in qa_pairs:
        pred = model(torch.randn(1,768))
        loss += torch.nn.functional.mse_loss(pred, torch.randn_like(pred))
    return loss

# ---- Recover stage (post‑hoc merging) ----
# Blend domain‑adapted checkpoint with original instruction model
def recover_merge(domain_ckpt, orig_ckpt, alpha=0.5):
    # Simple linear interpolation of parameters
    merged = {}
    for name, param in domain_ckpt.items():
        merged[name] = alpha * param + (1 - alpha) * orig_ckpt[name]
    return merged

# ---- Candidate Selection ----
# Choose checkpoint that balances domain QA accuracy and general capability guardrails
def select_checkpoint(candidates, domain_metric, general_metric, frontier_ratio=0.8):
    # frontier criterion: keep checkpoints whose domain score is within frontier_ratio of the best
    best_domain = max(domain_metric.values())
    selected = []
    for ckpt in candidates:
        if domain_metric[ckpt] >= frontier_ratio * best_domain:
            selected.append(ckpt)
    # among selected, pick the one with highest general metric
    return max(selected, key=lambda ckpt: general_metric[ckpt])

// Illustrative sketch (not from the paper)
const torch = require('torch-js'); // placeholder import

// Placeholder model (instruction‑tuned)
let model = torch.nn.Linear(768, 768);

// ---- Inject stage ----
// Three document‑generation objectives: continuation, rewrite, instruction reconstruction
function injectTraining(model, docs, mix = [1, 1, 1]) {
  let loss = 0;
  docs.forEach(doc => {
    // continuation loss (placeholder)
    loss += model.forward(torch.randn([1, 768])).mean().item();
    // rewrite loss (placeholder)
    loss += model.forward(torch.randn([1, 768])).mean().item();
    // instruction reconstruction loss (placeholder)
    loss += model.forward(torch.randn([1, 768])).mean().item();
  });
  return loss;
}

// ---- Align stage ----
// Fine‑tune on answer‑only QA pairs
function alignTraining(model, qaPairs) {
  let loss = 0;
  qaPairs.forEach(([q, a]) => {
    const pred = model.forward(torch.randn([1, 768]));
    const target = torch.randnLike(pred);
    loss += torch.nn.functional.mseLoss(pred, target).item();
  });
  return loss;
}

// ---- Recover stage (post‑hoc merging) ----
// Blend domain‑adapted checkpoint with original instruction model
function recoverMerge(domainCkpt, origCkpt, alpha = 0.5) {
  const merged = {};
  for (const name in domainCkpt) {
    merged[name] = alpha * domainCkpt[name] + (1 - alpha) * origCkpt[name];
  }
  return merged;
}

// ---- Candidate Selection ----
// Choose checkpoint balancing domain QA accuracy and general capability guardrails
function selectCheckpoint(candidates, domainMetric, generalMetric, frontierRatio = 0.8) {
  const bestDomain = Math.max(...Object.values(domainMetric));
  const frontier = frontierRatio * bestDomain;
  const filtered = candidates.filter(ckpt => domainMetric[ckpt] >= frontier);
  // pick the one with highest general metric among filtered
  return filtered.reduce((best, ckpt) => (generalMetric[ckpt] > generalMetric[best] ? ckpt : best), filtered[0]);
}