Spaces:
Running
Running
File size: 2,754 Bytes
beea5e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """Primitive order-1/n rules for the NTK-mean two-point function.
The data model is deliberately small: an outer cubic interaction exposes an
internal insertion slot, and an order-1/n correction vertex can fill it only
when its field, derivative order, and external NTK-line count match. These are
the rules stated in Section 4 and applied in Section 5.1 of arXiv:2508.11522v4.
"""
from dataclasses import asdict, dataclass
from fractions import Fraction
@dataclass(frozen=True)
class OuterVertex:
name: str
insertion_field: str
derivative_orders: tuple[int, ...]
@dataclass(frozen=True)
class CorrectionVertex:
name: str
field: str
derivative_order: int
external_ntk_lines: int
coefficient: Fraction
algebraic_term: str
OUTER_VERTICES = (
OuterVertex("sigma_prime_pair", "ntk", (0, 2)),
OuterVertex("delta_omega", "preactivation", (2, 4)),
)
CORRECTION_VERTICES = (
CorrectionVertex(
"Theta1", "ntk", 0, 2, Fraction(1),
"Cw*Theta1*E[sigma1_prime*sigma2_prime]",
),
CorrectionVertex(
"K1", "preactivation", 2, 0, Fraction(1, 2),
"sum(K1[b1,b2]*E[d_b1_b2 DeltaOmega12])",
),
CorrectionVertex(
"V", "preactivation", 4, 0, Fraction(1, 8),
"sum(V[(b1,b2),(b3,b4)]*E[d_b1_b2_b3_b4 DeltaOmega12])",
),
CorrectionVertex(
"D", "ntk", 2, 2, Fraction(1, 2),
"Cw*sum(E[d_b1_b2 sigma1_prime*sigma2_prime]*D[b1,b2,1,2])",
),
CorrectionVertex(
"F", "ntk", 2, 2, Fraction(1),
"Cw*sum(E[d_b1_b2 sigma1_prime*sigma2_prime]*F[b1,1,b2,2])",
),
)
def enumerate_diagrams(drop_vertex: str | None = None) -> list[dict]:
"""Enumerate every compatible diagram from primitive slot constraints."""
diagrams = []
for outer in OUTER_VERTICES:
for correction in CORRECTION_VERTICES:
if correction.name == drop_vertex:
continue
if correction.field != outer.insertion_field:
continue
if correction.derivative_order not in outer.derivative_orders:
continue
if correction.field == "ntk" and correction.external_ntk_lines != 2:
continue
if correction.field == "preactivation" and correction.external_ntk_lines != 0:
continue
diagrams.append(
{
"id": f"{outer.name}:{correction.name}",
"outer_vertex": asdict(outer),
"correction_vertex": {
**asdict(correction),
"coefficient": str(correction.coefficient),
},
}
)
return sorted(diagrams, key=lambda item: item["id"])
|