Performance & quality improvements: reward caching, geometric pre-filter, sparse attention, adaptive branching, learnable schedule, pLDDT calibration

[mooreneural]

Summary

This PR adds 7 improvements across inference speed, prediction quality, and
architectural expressiveness — motivated by profiling the generation pipeline
where AF2/RF3 reward evaluation dominates wall time (~50%) and the pair
representation is the main memory bottleneck.

Tier 1 — Speed (immediate wins, no retraining required)

• Reward caching (rewards/reward_utils.py, rewards/base_reward.py):
  RewardCache LRU cache keyed by sequence bytes, attached via
  reward_model.enable_cache(max_size). In beam search, sibling branches
  share nearly identical sequences — cache hits skip AF2/RF3 entirely.
  Expected: 2–4× speedup on reward evaluation.

• Geometric energy pre-filter (rewards/energy_reward.py):
  New GeometricEnergyReward runs CA clash detection and backbone N-CA-C
  angle outlier scoring in <100ms per sample with no GPU required. Use as a
  cheap gate before expensive folding model calls to discard geometrically
  invalid candidates early. Expected: 4–8× fewer AF2/RF3 calls on bad samples.

• Adaptive branching (search/beam_search.py):
  beam_search.adaptive_branching: true config flag tapers n_branch from
  full to ⌈n_branch/2⌉ linearly over the search. Late steps (low t) have
  sharp loss landscapes and branches converge quickly — extra branching wastes
  compute. Expected: 20–30% search speedup at <1% quality cost.

Tier 2 — Quality (structural improvements)

• Sparse geometric attention (nn/modules/pair_bias_attn.py,
  nn/modules/attn_n_transition.py): build_geometric_attn_mask() restricts
  pair attention to top-K CA nearest neighbours + local radius neighbourhood.
  New GeometricMultiheadAttnAndTransition and
  GeometricSparseMultiHeadBiasedAttentionADALN_MM are drop-in replacements
  that accept an optional ca_coords argument. Expected: ~3× pair-rep memory
  reduction for n > 150 and improved interface quality from geometric bias.

• SE(3) translation centering (nn/genie2_modules/structure_net.py):
  StructureNet(center_translations=True) subtracts per-sample COM from
  frame translations before each IPA block and restores it after. Improves
  numerical stability for large complexes and makes IPA effectively
  translation-equivariant. Backward compatible — center_translations=True
  is now the default.

Tier 3 — Analysis & calibration (training improvements)

• Learnable flow schedule (flow_matching/product_space_flow_matcher.py):
  LearnableSchedule(nsteps) parameterises integration time steps as a
  softmax-cumsum over learnable logits, allowing the model to concentrate
  evaluation steps where the loss gradient is highest. Opt-in per modality
  via learnable_schedule_nsteps in the product_flowmatcher config.

• pLDDT Platt calibration (rewards/alphafold2_reward.py):
  AF2RewardModel.calibrate(plddt_vals, success_labels) fits logistic
  regression to empirical wet-lab success data. extract_results now emits
  plddt_calibrated = sigmoid(scale × pLDDT + bias) alongside raw pLDDT.
  Corrects AF2's tendency to overestimate confidence on out-of-distribution
  generative sequences.

• AE fidelity analysis (partial_autoencoder/autoencoder.py):
  AutoEncoder.analyze_reconstruction_fidelity(dataloader) reports CA RMSD,
  active latent dimensions, and a latent_z_dim sizing recommendation — makes
  it easy to diagnose whether the bottleneck or the search algorithm is the
  limiting factor.

Test plan

• Confirm compute_reward_from_samples returns identical results with cache
  enabled vs disabled on a small batch
• Run GeometricEnergyReward.score() on a known-good PDB and a clashing
  structure; verify clash rate is higher for the bad one
• Beam search with adaptive_branching: true completes without assertion
  errors for nsamples=2, beam_width=2, n_branch=4, step_checkpoints=[0,100,200]
• StructureNet(center_translations=True) produces same topology as
  center_translations=False on a single forward pass (numerics may differ
  slightly)
• LearnableSchedule(400).get_ts() is monotonic with ts[0]=0, ts[-1]=1
• AutoEncoder.analyze_reconstruction_fidelity(val_loader) returns a dict
  with mean_ca_rmsd_ang and recommendation keys