Add quadratic gradient-label training for ABACUS

Author: ErjieWu

deepks/config/defaults.py

@@ -319,6 +319,7 @@ def get_default_device():
                 "cal_stress": 0,                                    # Compute stress flag.
                 "deepks_bandgap": 0,                                # DeePKS bandgap output flag.
                 "deepks_v_delta": 0,                                # DeePKS v_delta output flag.
+                "deepks_grad": 0,                                    # DeePKS grad output flag.
                 "deepks_out_labels": 1,                             # DeePKS label output flag.
                 "deepks_scf": 0,                                    # DeePKS in-SCF switch.
                 "out_wfc_lcao": 0,                                  # Output LCAO wfc flag.
@@ -367,7 +368,10 @@ def get_default_device():
         "objective": {
             "losses": [],                                           # Structured loss definitions.
             "energy_per_atom": None,                                # Energy-per-atom option.
-            "grad_penalty": None,                                   # Gradient penalty setting.
+            "force_grad_penalty": None,                             # Force-gradient penalty (requires grad_veg/eg0 data).
+            "hessian_penalty": None,                                  # Network Hessian curvature penalty weight.
+            "hessian_penalty_method": None,                           # "exact" (default) or "hutchinson".
+            "hessian_n_probes": None,                                 # Hutchinson probe count (default 1).
             "vd_divide_by_nlocal": False,                           # Normalize vd by nlocal.
             "vd_masked_loss": 0,                                    # Masked vd loss mode.
             "vd_masked_S_threshold": 1e-6,                          # S threshold for masked vd loss.
@@ -385,6 +389,7 @@ def get_default_device():
             "display_detail_test": 0,                               # Detail level for test output.
             "display_natom_loss": False,                            # Print natom-wise loss.
             "fix_embedding": False,                                 # Freeze embedding layers.
+            "trainable_patterns": None,                             # Optional shell patterns selecting trainable parameters.
             "stage_schedule": [],                                   # In-run staged training schedule.
             "optimizer": {
                 "lr": 0.01,                                         # Initial learning rate.
@@ -469,6 +474,7 @@ def get_default_backend_input(backend_name):
                 "cal_stress",
                 "deepks_bandgap",
                 "deepks_v_delta",
+                "deepks_grad",
                 "deepks_out_labels",
                 "deepks_scf",
                 "out_wfc_lcao",

deepks/config/docs.py

@@ -202,6 +202,13 @@ def _fmt_default(value):
                 "default": "optional",
                 "description": "Optimizer and scheduler settings. The current recipes use optimizer.lr, optimizer.weight_decay, scheduler.decay_steps, scheduler.decay_rate, and scheduler.stop_lr.",
             },
+            {
+                "name": "ml.train.trainable_patterns",
+                "type": "list[str] | str | null",
+                "tasks": "train, iterate",
+                "default": "null",
+                "description": "Optional shell-style model-parameter name patterns for layer-wise continuation, for example densenet.layers.3.* for a final-head-only stage.",
+            },
             {
                 "name": "ml.fit_elem",
                 "type": "bool",
@@ -311,6 +318,7 @@ def _fmt_default(value):
                     "cal_stress": "Whether to compute stress.",
                     "deepks_bandgap": "ABACUS DeePKS bandgap flag.",
                     "deepks_v_delta": "ABACUS DeePKS v_delta flag.",
+                    "deepks_grad": "ABACUS DeePKS grad output flag.",
                     "deepks_out_labels": "ABACUS DeePKS label dump flag.",
                     "deepks_scf": "ABACUS DeePKS SCF switch.",
                     "out_wfc_lcao": "ABACUS out_wfc_lcao switch.",
@@ -356,6 +364,64 @@ def render_input_parameter_doc():
                 f"| `{row['name']}` | `{row['type']}` | `{row['tasks']}` | `{row['default']}` | {description} |"
             )
 
+    lines.extend(
+        [
+            "",
+            "## Descriptor-gradient supervision",
+            "",
+            "`physics.backend.input.gradient_label` combines current-run property operators",
+            "as `M=sum_p c_p A_p^T A_p` and `b=sum_p c_p A_p^T Delta X_p`. The coefficients",
+            "are effective coefficients: DeePKS-L performs no hidden unit conversion or",
+            "loss-size normalization.",
+            "",
+            "| Parameter | Type | Default | Description |",
+            "| --- | --- | --- | --- |",
+            "| `force` | `float >= 0` | `0` | Effective force coefficient `c_F`; zero disables force. |",
+            "| `stress` | `float >= 0` | `0` | Effective stress coefficient `c_S`; nonzero stress automatically adds six local cell-strain coordinates. |",
+            "| `hr` | `float >= 0` | `1` | Effective real-space Hamiltonian coefficient `c_H`; zero disables HR. |",
+            "| `ridge` | `float >= 0` | `1e-8` | Tikhonov ridge used only when solving the direct target `g*=(M+ridge I)^-1 b`. It does not alter quadratic `M,b`. |",
+            "| `fallback` | `lstsq or pinv` | `lstsq` | Singular-solve fallback used only for the direct target. |",
+            "| `eigen_filter` | `dict or null` | `null` | Optional direct-target filter. Specify exactly one of `{rcond: value}` or `{min_eig: value}`. It does not alter quadratic `M,b`. |",
+            "",
+            "All ABACUS DeePKS energy, force, stress, Hamiltonian, and projected-Hamiltonian",
+            "labels are required to be in Hartree. Legacy Ry real-space Hamiltonian data",
+            "must be regenerated or converted once before use.",
+            "",
+            "To reproduce direct-property reductions, include them in the coefficients.",
+            "For a two-atom Si frame with direct force weight 1, stress weight 1, HR weight",
+            "0.005, six force/stress components, HR range 9, and `nlocal=26`, use:",
+            "",
+            "```yaml",
+            "gradient_label:",
+            "  force: 0.16666666666666666       # 1/6",
+            "  stress: 0.16666666666666666      # 1/6",
+            "  hr: 2.136752136752137e-5         # 0.005/(9*26)",
+            "```",
+            "",
+            "The `g_label` objective accepts exactly two loss types:",
+            "",
+            "```yaml",
+            "- name: g_label",
+            "  weight: 1.0",
+            "  loss:",
+            "    type: direct       # MSE(g, g*)",
+            "```",
+            "",
+            "or",
+            "",
+            "```yaml",
+            "- name: g_label",
+            "  weight: 1.0",
+            "  loss:",
+            "    type: quadratic    # mean_frame(g^T M g - 2 g^T b)",
+            "```",
+            "",
+            "The quadratic form is inverse-free and is the recommended property-equivalent",
+            "path. Stress-coordinate augmentation is detected from the saved label width;",
+            "there is no separate loss-side switch.",
+        ]
+    )
+
     lines.extend(
         [
             "",

deepks/config/packager.py

@@ -18,7 +18,10 @@
 _OBJECTIVE_SHARED_KEYS = {
     "losses",
     "energy_per_atom",
-    "grad_penalty",
+    "force_grad_penalty",
+    "hessian_penalty",
+    "hessian_penalty_method",
+    "hessian_n_probes",
     "vd_divide_by_nlocal",
     "vd_masked_loss",
     "vd_masked_S_threshold",
@@ -395,17 +398,8 @@ def pack_child(child_type, child_config):
             resolved_terms = _resolve_hierarchical_terms(ml)
             if resolved_terms:
                 main_train.setdefault("ml", {}).setdefault("objective", {})["terms"] = resolved_terms
-        elif use_profile_scf:
-            main_train["data"]["train"] = [
-                f"../00.scf/level.{i:02d}/data_train/*" for i in range(len(scf_profiles))
-            ]
         if data.get("test") is not None:
-            if use_profile_scf:
-                main_train["data"]["test"] = [
-                    f"../00.scf/level.{i:02d}/data_test/*" for i in range(len(scf_profiles))
-                ]
-            else:
-                main_train["data"]["test"] = "data_test/*"
+            main_train["data"]["test"] = "data_test/*"
         if isinstance(runtime.get("io"), dict):
             main_train["runtime"]["io"] = deepcopy(runtime["io"])
         if child_proj_basis:

deepks/interface/objectives/config.py

@@ -4,7 +4,10 @@
 
 
 _OBJECTIVE_OPTION_KEYS = (
-    "grad_penalty",
+    "force_grad_penalty",
+    "hessian_penalty",
+    "hessian_penalty_method",
+    "hessian_n_probes",
     "energy_per_atom",
     "vd_divide_by_nlocal",
     "vd_masked_loss",
@@ -28,6 +31,7 @@
     "density_m": ("density_m_factor", "density_m_lossfn", "density_m_occ"),
     "phi_align": ("phi_align_factor", "phi_align_lossfn", "phi_align_occ"),
     "density": ("density_factor", None, None),
+    "g_label": ("grad_factor", "grad_lossfn", None),
 }
 
 _FACTOR_KEYS = tuple(mapped[0] for mapped in _LOSS_NAME_MAP.values())
@@ -92,7 +96,8 @@ def build_descriptor_property_eval_args(objective_args, *, detailed=False):
             "energy_factor": 1.0,
             "force_factor": 0.0,
             "density_factor": 0.0,
-            "grad_penalty": 0.0,
+            "force_grad_penalty": 0.0,
+            "hessian_penalty": 0.0,
             "energy_per_atom": energy_per_atom,
         }
 

deepks/interface/objectives/descriptor_properties.py

@@ -31,7 +31,11 @@ def __init__(
         phi_align_factor=0.0,
         phi_align_occ=0,
         density_factor=0.0,
-        grad_penalty=0.0,
+        force_grad_penalty=0.0,
+        hessian_penalty=0.0,
+        grad_factor=0.0,
+        hessian_penalty_method="exact",
+        hessian_n_probes=1,
         energy_lossfn=None,
         force_lossfn=None,
         stress_lossfn=None,
@@ -43,6 +47,7 @@ def __init__(
         band_lossfn=None,
         bandgap_lossfn=None,
         density_m_lossfn=None,
+        grad_lossfn=None,
         energy_per_atom=0,
         vd_divide_by_nlocal=False,
         vd_masked_loss=0,
@@ -98,7 +103,11 @@ def __init__(
         self.density_m_factor = density_m_factor
         self.phi_align_factor = phi_align_factor
         self.d_factor = density_factor
-        self.g_penalty = grad_penalty
+        self.fg_penalty = force_grad_penalty
+        self.h_penalty = hessian_penalty
+        self.grad_factor = grad_factor
+        self.hessian_method = hessian_penalty_method
+        self.hessian_n_probes = hessian_n_probes
         self.energy_per_atom = energy_per_atom
         self.vd_divide_by_nlocal = vd_divide_by_nlocal
         self.vd_masked_loss = vd_masked_loss
@@ -137,7 +146,11 @@ def __init__(
             "phi_align_factor": phi_align_factor,
             "phi_align_occ": phi_align_occ,
             "density_factor": density_factor,
-            "grad_penalty": grad_penalty,
+            "force_grad_penalty": force_grad_penalty,
+            "hessian_penalty": hessian_penalty,
+            "grad_factor": grad_factor,
+            "hessian_penalty_method": hessian_penalty_method,
+            "hessian_n_probes": hessian_n_probes,
             "energy_lossfn": energy_lossfn,
             "force_lossfn": force_lossfn,
             "stress_lossfn": stress_lossfn,
@@ -149,6 +162,7 @@ def __init__(
             "band_lossfn": band_lossfn,
             "bandgap_lossfn": bandgap_lossfn,
             "density_m_lossfn": density_m_lossfn,
+            "grad_lossfn": grad_lossfn,
             "energy_per_atom": self.energy_per_atom,
             "vd_divide_by_nlocal": vd_divide_by_nlocal,
             "vd_masked_loss": vd_masked_loss,
@@ -179,6 +193,18 @@ def compute_losses(self, model, batch):
         if dropped:
             self._warn_dropped_properties(dropped)
         derivative_spec = self.property_engine.required_model_derivatives(requested_properties)
+        # Adapter-level uses of input_grad (grad_factor, fg_penalty, h_penalty,
+        # d_factor) bypass the property engine, so they don't appear in
+        # requested_properties and derivative_spec["input"] stays False even
+        # when they are active.  Force it True so model_input.requires_grad_
+        # is set and _compute_input_grad actually runs.
+        _adapter_needs_input_grad = (
+            self.grad_factor > 0 or self.fg_penalty > 0
+            or self.h_penalty > 0 or self.d_factor > 0
+        )
+        if _adapter_needs_input_grad and not derivative_spec.get("input"):
+            derivative_spec = dict(derivative_spec)
+            derivative_spec["input"] = True
 
         # R1 + R2: call the model in dict-in / dict-out form, then apply
         # interface-side reducers to obtain the supervision-ready primary
@@ -203,9 +229,15 @@ def compute_losses(self, model, batch):
             model_derivatives=model_derivatives,
             context=calc_context,
         )
-        if self.g_penalty > 0 and input_grad is not None and "eg0" in batch.context:
+        if self.fg_penalty > 0 and input_grad is not None and "eg0" in batch.context:
             eg_base, gveg = batch.context["eg0"], batch.context["gveg"]
             predictions["grad_total"] = torch.einsum("...apg,...ap->...g", gveg, input_grad) + eg_base
+        if self.h_penalty > 0 and input_grad is not None:
+            predictions["hessian_penalty"] = self._compute_hessian_penalty(
+                input_grad, model_input, self.hessian_method, self.hessian_n_probes
+            )
+        if self.grad_factor > 0 and input_grad is not None:
+            predictions["input_grad"] = input_grad
         if self.d_factor > 0 and input_grad is not None and "gldv" in batch.context:
             predictions["density_regularizer"] = (batch.context["gldv"] * input_grad).mean(0).sum()
 
@@ -308,6 +340,60 @@ def _compute_input_grad(primary_scalar, model_input, needs_grad):
         )
         return grad
 
+    @staticmethod
+    def _compute_hessian_penalty(input_grad, model_input, method, n_probes):
+        """Compute ||d2E/dlambda2||_F^2 for the Hessian curvature penalty.
+
+        Two methods:
+          "exact"       -- row-by-row Jacobian of input_grad; O(ndesc_per_atom) backward
+                           passes. Exploits per-atom independence: iterates over descriptor
+                           dims only (not natom*ndesc). Default.
+          "hutchinson"  -- stochastic Frobenius-norm estimate via n_probes Rademacher
+                           vectors; O(n_probes) backward passes. Cheaper for large systems.
+        """
+        if method == "hutchinson":
+            return DescriptorPropertyObjectiveAdapter._hutchinson_hessian(
+                input_grad, model_input, n_probes
+            )
+        return DescriptorPropertyObjectiveAdapter._exact_hessian(input_grad, model_input)
+
+    @staticmethod
+    def _exact_hessian(input_grad, model_input):
+        # input_grad: (batch, natom, ndesc_per_atom)  -- already create_graph=True
+        # Exploit per-atom independence: iterate over ndesc_per_atom only.
+        ndesc = input_grad.shape[-1]
+        frob_sq = input_grad.new_tensor(0.0)
+        for d in range(ndesc):
+            row = torch.autograd.grad(
+                input_grad[..., d].sum(),   # scalar -- sum over batch and atoms
+                model_input,
+                retain_graph=(d < ndesc - 1),
+                create_graph=False,
+                only_inputs=True,
+            )[0]  # (batch, natom, ndesc_per_atom) -- d-th Hessian row per atom
+            frob_sq = frob_sq + row.pow(2).mean()  # mean over batch*atoms, sum over row
+        # Both methods compute ||H||_F^2 / (batch*natom*ndesc_per_atom).
+        return frob_sq
+
+    @staticmethod
+    def _hutchinson_hessian(input_grad, model_input, n_probes):
+        # Unbiased estimator: E[||Hv||^2] = ||H||_F^2 for Rademacher v.
+        penalty = input_grad.new_tensor(0.0)
+        for _ in range(n_probes):
+            # Rademacher vector (+-1 with equal probability)
+            v = torch.randint(0, 2, input_grad.shape, dtype=input_grad.dtype,
+                              device=input_grad.device) * 2 - 1
+            Hv = torch.autograd.grad(
+                (input_grad * v).sum(),
+                model_input,
+                retain_graph=True,
+                create_graph=False,
+                only_inputs=True,
+            )[0]
+            penalty = penalty + Hv.pow(2).mean()
+        return penalty / n_probes
+
+
     def _requested_properties(self):
         requested = set()
         if self.primary_property:
@@ -338,8 +424,10 @@ def _warn_dropped_properties(self, dropped_properties):
     def print_head(self, name, data_keys, align_len=20):
         data_keys = self._normalize_field_keys(data_keys)
         info = f"{name}_energy".rjust(align_len)
-        if self.g_penalty > 0 and "eg0" in data_keys:
-            info += f"{name}_grad".rjust(align_len)
+        if self.fg_penalty > 0 and "eg0" in data_keys:
+            info += f"{name}_force_grad".rjust(align_len)
+        if self.h_penalty > 0:
+            info += f"{name}_hessian".rjust(align_len)
         if self.f_factor > 0 and "force" in data_keys:
             info += f"{name}_force".rjust(align_len)
         if self.s_factor > 0 and "stress" in data_keys:
@@ -362,6 +450,12 @@ def print_head(self, name, data_keys, align_len=20):
             info += f"{name}_phi_align".rjust(align_len)
         if self.d_factor > 0 and "gldv" in data_keys:
             info += f"{name}_density".rjust(align_len)
+        has_gradient_supervision = (
+            "g_label" in data_keys
+            or {"g_label_metric", "g_label_projection"}.issubset(data_keys)
+        )
+        if self.grad_factor > 0 and has_gradient_supervision:
+            info += f"{name}_grad".rjust(align_len)
         print(info, end="")
 
     @staticmethod
@@ -375,6 +469,7 @@ def _normalize_field_keys(data_keys):
             "lb_vd": "v_delta",
             "lb_vdr": "vdr",
             "lb_phi": "phi",
+            "lb_g": "g_label",
             "lb_band": "band",
         }
         return {aliases.get(key, key) for key in data_keys}