Add MSE + covariance penalty term custom objective function by wagnerlmichael · Pull Request #21 · ccao-data/lightsnip

wagnerlmichael · 2026-04-10T14:24:01Z

Adding custom objective function (MSE + an additional penalty term)

The main contribution of this PR is to wire in logic to take advantage of lightgbm's support of custom objective functions to add the custom objective function that the MIT research team has developed with the aim of improving vertical equity. Based on our own internal investigation, it seems to be quite a promising result. There is EDA that show specific figures in the EI 408 output folder.

MIT has reviewed the math/ML theory portion of this PR and signed off on its correctness.

Implementation details

The main complication here is that, although supported, adding a custom objective function to lightgbm is a little bit finicky.

Here are the main complications:

lgb.train errors if it sees objective = "mse_cov" (it's not a built-in name)
lgb.train errors if it sees an mse_cov_rho parameter (it's not a LightGBM param at all, we're tunneling it through the parsnip set_engine(...) dot
lgb.train errors / behaves weirdly if both objective and obj are set

Supports model PR: ccao-data/model-res-avm#475

Other approaches tried

We tried using tailor and building the transforms into the tidymodels workflow
model res - ccao-data/model-res-avm#520.

We also tried smuggling the transforms into lightsnip itself

model res PR: lightsnip smuggle [DRAFT] model-res-avm#524
lightsnip PR: lightsnip smuggle [DRAFT] #27

Ultimately we decided against these other strategies due to a combination of increased complexity and recommendations against doing these types of transforms from the tidymodels maintainers.

wagnerlmichael · 2026-04-15T16:00:50Z

+    custom_obj <- make_obj_mse_cov(rho = rho_val, y_mean = mean(y))
+    # When `obj` is a custom callback we must NOT also set `objective` in
+    # params, otherwise lgb.train will reject the unknown name.
+    others$objective <- NULL


LightGBM takes either objective or obj args, code here, and ultimately collects one of the two as objective

More context on the comment here from the lightbgm source code. Shown here:

# extract any function objects passed for objective or metric fobj <- NULL if (is.function(params$objective)) { fobj <- params$objective params$objective <- "none" }

Checks to see if if there is a custom function supplied, saves it as fobj to be passed to C later using gradient and hessian, and the "none" value lets lightgbm know that the custom math will be supplied later, rather than it using one of its built in C objective functions.

More a clarification than anything " saves it as fobj to be passed to C later using gradient and hessian" these are the gradient and hessian as calculated by the new objection function- (yea?)

Yep, that's right

wagnerlmichael · 2026-04-20T16:49:54Z

+    small_h <- hess < zero_grad_tol
+    if (any(small_h)) hess[small_h] <- zero_grad_tol
+
+    list(grad = grad, hess = hess)


This list is passed and parsed as gpair in lightbm source code

wagnerlmichael · 2026-05-21T17:00:06Z

+  #   - whether to construct the mse_cov objective callback
+  mse_cov_rho_val <- NULL
+
+  if (!is.null(others$objective) && identical(others$objective, "mse_cov")) {


In case you're thinking "where does others$objective come from"?

In the corresponding model-res PR, we set this value within the set_engine() call in 01-train.R

objective = params$model$objective,

which grabs from our params.yaml definition

model: engine: "lightgbm" objective: "mse_cov"

objective isn't one of parsnip::boost_tree()'s modeled hyperparameters (those are things like trees, tree_depth, learn_rate, etc.) so parsnip captures it, along with anything else we pass to set_engine("lightgbm", ...), as the engine "dots." Those dots are forwarded to lightsnip::train_lightgbm(...), whose own ... is materialized into a named list on the first line of the body: others <- list(...). From that point on, others$objective == "mse_cov" and others$mse_cov_rho == 1, and those are the values the sentinel branch sniffs to decide whether to swap in the custom callback.

Even though others$objective isn't one of parsnips modeled hyperparameters, it is the canonical way to pass the objective functions, custom or supported out of the box.

The prior case with rmse

If we were to supply something the model knows like we have done historically for 'rmse', the rmse value would have floated all the way to C++ through a .h file where it basically says, if "rmse" then point to "regression" which is lands in the C++ file here

wagnerlmichael · 2026-05-21T20:08:06Z

+      y_mean = mean(y[trn_index])
+    )
+  }
+


y[train_index] will either call the full training set or just a portion of it if we are running a hyperparam search:

n <- nrow(x) if (validation > 0) { m <- min(floor(n * (1 - validation)), n - 1) if (sample_type == "random") { m <- min(floor(n * (1 - validation)), n - 1) trn_index <- sample(1:n, size = max(m, 2)) val_index <- setdiff(1:n, trn_index) }

The reason we hold out some data for hyperparam search validation is that we are evaluating the per model selection based on a small hold out set

wagnerlmichael · 2026-05-26T08:35:55Z

+  # Wire in the custom objective callback (if any) under lgb.train's `obj` arg
+  if (!is.null(custom_obj)) {
+    main_args$obj <- quote(custom_obj)
+  }


custom_obj is created on line 267 drawing from objectives.R. Then we attach this function to main_args$obj here.

Which gets handed to lgb.Booster.R here where it is saved as fobj. And then ends up being called and producing the gradient (which direction to move) and hessian (how far to move) which are then consumed by C to fit the subsequent tree

That's really clever! Just making a note (more for myself than anything else) - as to where lgbm grabs the gradient and hessian

wagnerlmichael · 2026-05-26T08:47:43Z

+    # dc/dy_pred_i = (1/n) * yc_i  (since dr_i/dy_pred_i = 1)
+    a <- yc / n
+    grad_pen <- rho * n * cov_val * a
+    hess_pen <- rho * n * (a^2)


The math here is completely borrowed from the source implementation and has been reviewed by the code authors

Very clean. Since we've empirically tested it with the diagonal approximation- no issues. Noting for future state that the (I believe) the "smooth" penalty is similar, without the diagonal assumption.

wrridgeway

Seems like the objective function itself made its way through the pipeline just fine. Thanks for all the work on this you guys. If the math has been checked by original authors of this method, that's good by me.

wrridgeway · 2026-06-24T15:07:46Z

+
+  if (!is.null(others$objective) && identical(others$objective, "mse_cov")) {
+    mse_cov_rho_val <-
+      if (is.null(mse_cov_rho)) 1e-3 else as.numeric(mse_cov_rho)


Is this default value noted anywhere?

This made me realize that the default value here doesn't really make sense, I just ported it over from the original python repo. I switched it here: e2c760d

codecov · 2026-07-09T20:18:31Z

Codecov Report

❌ Patch coverage is 20.51282% with 31 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.25%. Comparing base (b8ed4fc) to head (aa9d7b0).
⚠️ Report is 5 commits behind head on master.

Files with missing lines	Patch %	Lines
R/objectives.R	0.00%	21 Missing ⚠️
R/lightgbm.R	44.44%	10 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master      #21       +/-   ##
===========================================
- Coverage   81.20%   67.25%   -13.96%     
===========================================
  Files           3        4        +1     
  Lines         133      171       +38     
===========================================
+ Hits          108      115        +7     
- Misses         25       56       +31

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

wagnerlmichael · 2026-07-10T13:28:13Z

+  obj <- make_objective_mse_cov(rho = 0, y_mean = mean(y))
+  out <- obj(preds, dtrain)
+
+  expect_equal(out$grad, 2 * (preds - y))


This is testing the assumption that if we supply a rho of 0 it will behave exactly like MSE would. It tests that the gradients are equal to the first derivative of MSE in per-observation squared error

Since the whole covariance term project is simply appending another mathematical term to MSE, this is a check that makes sure nothing else happens along the way with the whole lightsnip lgb carry through

wagnerlmichael · 2026-07-10T14:04:24Z

+  cov_val <- mean((preds - y) * y_centered)
+  a <- y_centered / n
+
+  expect_equal(out$grad, 2 * (preds - y) + rho * n * cov_val * a)


Tests that the gradient is the MSE derivative plus the covariance term derivative

wagnerlmichael · 2026-07-10T14:25:32Z

+
+  expect_equal(nrow(pred), nrow(mtcars))
+  expect_true(all(is.finite(pred$.pred)))
+  expect_not_constant_predictions(pred$.pred)


basic model output checking

one pred per input

no NaNs/infinities

model trains correctly (no constant predictions)

wagnerlmichael · 2026-07-10T14:34:00Z

+  # The penalty weight should reach the objective: different rho values
+  # must produce different fits
+  pred_plain <- suppressWarnings(predict(fit_with_rho(0), mtcars[, -1]))
+  expect_all_preds_differ(list(pred$.pred, pred_plain$.pred))


predictions differ between a rho of 0 (plain MSE) and a non-zero rho (0.5)

wagnerlmichael · 2026-07-10T14:34:38Z

+      parsnip::fit(mpg ~ ., data = mtcars),
+    regexp = "requires `mse_cov_rho`"
+  )
+})


tests that no rho supplied throws an error

jeancochrane

Awesome work! No issues with the code or the tests. I think the docs need a bit more formatting, but that should be pretty fast once you've got an environment up and running for generating the docs.

jeancochrane · 2026-07-10T23:01:07Z

+#' Custom LightGBM objective: MSE + rho * Cov(r, y)^2
+#'
+#' @description Build a custom LightGBM objective callback that minimizes a
+#' standard squared-error loss plus a soft penalty on the covariance between
+#' the per-sample residual `r = y_pred - y_true` and the (centered) labels
+#' `y_true`. The penalty pushes the model toward "vertical equity" by
+#' discouraging residuals that systematically scale with `y`.
+#'
+#' This is an R port of the `LGBCovPenalty` objective from
+#' an active collabration. (https://github.com/nicacevedo/soft-vertical-equity-constrained-mass-appraissal) # nolint
+#' It is intended to be used when the model is trained in log-space (so the
+#' "diff" residual is equivalent to a log-ratio).
+#'
+#' Penalty (using mean-centered labels y_centered = y_true - mean(y_true)):
+#' \deqn{cov = (1/n) * sum_i r_i * y_centered_i}
+#' \deqn{penalty = 0.5 * rho * n * cov^2}
+#' Diagonal Hessian approximation is used (matches the reference Python
+#' implementation).


[Nitpick, required] These docs don't seem to be rendering super well. A few issues:

The math in the title is neither LaTeX-formatted nor code-formatted, so it looks awkward

Link text is not getting formatted as a hyperlink

The math in the last paragraph is not being formatted properly

I tested the docs by running renv::install(".", dependencies = TRUE) and then running pkgdown::build_site(preview = TRUE). Let me know if you want help figuring out how to fix these formatting issues!

wagnerlmichael · 2026-07-13T17:39:15Z

 \name{make_objective_mse_cov}
 \alias{make_objective_mse_cov}
-\title{Custom LightGBM objective: MSE + rho * Cov(r, y)^2}
+\title{Custom LightGBM objective: MSE with a squared covariance penalty}


Seems like title doesn't support LaTeX so I just used english

wagnerlmichael · 2026-07-13T17:52:46Z

-RoxygenNote: 7.2.3
+Roxygen: list(markdown = TRUE)
 Config/testthat/edition: 3
+Config/roxygen2/version: 8.0.0


Used roxygen2::roxygenise() locally which auto-added this

wagnerlmichael · 2026-07-13T17:53:13Z

    workflows,
    yardstick
-RoxygenNote: 7.2.3
+Roxygen: list(markdown = TRUE)


This fixed the hyperlink among other things

jeancochrane

Nice work fixing up the docs!

jeancochrane · 2026-07-13T19:05:22Z

+#' @param verbose Integer. Verbosity level: < 0 = Fatal, 0 = Error (Warning),
+#'   1 = Info, > 1 = Debug.


[Question, non-blocking] What's the motivation for this change?

The previous > was attempted to be parsed by roxygen as markdown before being translated to Rd format, but it errored out because apparently Rd files have no blockquote equivalent. It was because the previous line started with >

https://github.com/r-lib/roxygen2/blob/aaacc1322cffadfaa556968f4fb5a363482b086d/R/markdown.R#L109

…475) The main contribution of this PR is to wire in logic to take advantage of lightgbm's support of custom objective functions to add the custom objective function option that the MIT research team has developed with the aim of improving vertical equity. Based on our own internal investigation, it seems to be quite a promising specification. There is EDA that show specific figures in the EI 408 output folder. This PR adds `params.yaml` file edits which include 1. The ability to select the custom obj function from `params.yaml` (`mse_cov`) 2. A toggle to set the model to log the training data. This is necessary for our implementation of the custom objective function. Adding the ability to log the outcome variable is also a generally good option to have in terms of ML modeling strategies. With `select * from model.metadata where run_id = '2026-05-20-infallible-dan'` you can see the persistence of `log_transform_enable` With `SELECT * FROM "model"."parameter_final" where run_id = '2026-05-20-infallible-dan'` you can see `objective = mse_cov` and `mse_cov_rho=1` Supported by lightbgm PR: ccao-data/lightsnip#21 ## Other approaches tried We tried using tailor and building the transforms into the tidymodels workflow model res - #520. We also tried smuggling the transforms into lightsnip itself - model res PR: #524 - lightsnip PR: ccao-data/lightsnip#27 Ultimately we decided against these other strategies due to a combination of increased complexity and recommendations against doing these types of transforms from the tidymodels maintainers. --------- Co-authored-by: William Ridgeway <10358980+wrridgeway@users.noreply.github.com> Co-authored-by: TimCookCountyDS <Timothy.Sparer@cookcountyil.gov>

wagnerlmichael added 3 commits April 10, 2026 09:19

Add draft of cov penalty

61da860

Add roxygen rd file

799914f

Update pkgdown yml

f228237

wagnerlmichael commented Apr 15, 2026

View reviewed changes

wagnerlmichael added 3 commits April 15, 2026 12:45

Lint

f9b0565

Add nolint

2f91001

Edit nolint

f8828a4

wagnerlmichael mentioned this pull request Apr 20, 2026

Add covariance penalty option and configurable logging of sale price ccao-data/model-res-avm#475

Merged

wagnerlmichael commented Apr 20, 2026

View reviewed changes

wagnerlmichael added 8 commits May 12, 2026 13:15

Attempt to fix data leakage

ee68585

Edit comments

6b0dd91

Add comments

6735cf8

Add spacing

7c11ea8

Fix spacing

76472b5

Lint

ed3d09e

Decrease line length

b44164f

Remove zero grad lines

be61335

wagnerlmichael commented May 21, 2026

View reviewed changes

Remove zero_grad code

0c94694

wagnerlmichael commented May 21, 2026

View reviewed changes

wagnerlmichael changed the title ~~[WIP] Test custom obj function~~ [WIP] Add MSE + covariance penalty term custom objective function May 26, 2026

wagnerlmichael commented May 26, 2026

View reviewed changes

wagnerlmichael and others added 2 commits May 26, 2026 04:02

Remove type=raw spec

61613fd

Merge branch 'master' into test-new-obj-function

55cc691

wagnerlmichael marked this pull request as ready for review May 26, 2026 12:51

wagnerlmichael requested review from jeancochrane and wrridgeway as code owners May 26, 2026 12:51

wagnerlmichael changed the title ~~[WIP] Add MSE + covariance penalty term custom objective function~~ Add MSE + covariance penalty term custom objective function May 26, 2026

wagnerlmichael requested a review from TimCookCountyDS May 26, 2026 12:51

wrridgeway approved these changes Jun 24, 2026

View reviewed changes

jeancochrane mentioned this pull request Jun 30, 2026

MIT regressivity penalty implementation: R lgbm modifications, Finalize the addition of the basic covariance penalty to lightsnip + lgbm and test in pipeline ccao-data/model-res-avm#480

Closed

wagnerlmichael and others added 9 commits July 8, 2026 15:45

Swap yc to y_centered

67e1f4a

Swap in for ==

446e4cb

Merge branch 'master' into test-new-obj-function

1a883ab

Point to actions branch

3d87b06

Point back to main

a4ac299

Switch default value check to null check

e2c760d

Remove nolint tag

5c624fa

Add y mean check

490ba70

Version bump

aa9d7b0

Add new tests

5caaf70

wagnerlmichael commented Jul 10, 2026

View reviewed changes

Remove hessian test

55f1408

wagnerlmichael commented Jul 10, 2026

View reviewed changes

Remove hessian tset

7c2590d

wagnerlmichael commented Jul 10, 2026

View reviewed changes

jeancochrane approved these changes Jul 10, 2026

View reviewed changes

Update formatting

bf8f04e

wagnerlmichael commented Jul 13, 2026

View reviewed changes

jeancochrane approved these changes Jul 13, 2026

View reviewed changes

wagnerlmichael merged commit 855d364 into master Jul 14, 2026
11 checks passed

wagnerlmichael deleted the test-new-obj-function branch July 14, 2026 17:46

jeancochrane mentioned this pull request Jul 15, 2026

lightsnip smuggle [DRAFT] #27

Closed

		#' @param verbose Integer. Verbosity level: < 0 = Fatal, 0 = Error (Warning),
		#' 1 = Info, > 1 = Debug.

Uh oh!

Conversation

wagnerlmichael commented Apr 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Adding custom objective function (MSE + an additional penalty term)

Implementation details

Other approaches tried

Uh oh!

wagnerlmichael Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wagnerlmichael May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

In case you're thinking "where does others$objective come from"?

The prior case with rmse

Uh oh!

wagnerlmichael May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wagnerlmichael May 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wrridgeway left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented Jul 9, 2026

Codecov Report

Uh oh!

wagnerlmichael Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jeancochrane left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wagnerlmichael Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jeancochrane left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

wagnerlmichael commented Apr 10, 2026 •

edited

Loading

wagnerlmichael Apr 15, 2026 •

edited

Loading

wagnerlmichael May 21, 2026 •

edited

Loading

In case you're thinking "where does `others$objective` come from"?

The prior case with `rmse`

wagnerlmichael May 21, 2026 •

edited

Loading

wagnerlmichael May 26, 2026 •

edited

Loading

wagnerlmichael Jul 10, 2026 •

edited

Loading

wagnerlmichael Jul 13, 2026 •

edited

Loading