Behaviour of `base_margin` during prediction

[bblodfon]

Hi,

I wanted to better undestand the behavior of using a base_margin during training and prediction of an xgboost model. I read the very nice tutorial and tried to extend it in the case where we have a simple train/test split. As you will see in the below code, I have two models that are exactly the same (one trained with base_score = 0.5 and one with base_margin = 0) - which makes sense according to the tutorial - but upon prediction, using (or not using) the base_margin results in different predictions. Is that expected behavior? Maybe I have misundestood something?

library(xgboost)
library(testthat)

# Load data, split train/test
data(agaricus.train, package = "xgboost")
X = agaricus.train$data
y = agaricus.train$label

set.seed(42)
train_ids = sample(x = 1:nrow(X), size = 4000)
test_ids = setdiff(1:nrow(X), train_ids)

xgb_data_train = xgboost::xgb.DMatrix(data = X[train_ids, ], label = y[train_ids])
xgb_data_test  = xgboost::xgb.DMatrix(data = X[test_ids, ], label = y[test_ids])

# m1 model with no base_margin (default base_score = 0.5)
m1 = xgboost::xgb.train(
  data = xgb_data_train,
  params = list(base_score = 0.5, objective = "binary:logistic"),
  nrounds = 1)

# TRAIN ----
# does offset change trained model? => YES (seems also independent of nrounds)
# try base_margin = 1 => NOT the same as base_score = 0.5
xgboost::setinfo(xgb_data_train, "base_margin", rep(1, length(train_ids)))
#> [1] TRUE
m2 = xgboost::xgb.train(
  data = xgb_data_train,
  params = list(objective = "binary:logistic"),
  nrounds = 1)
testthat::expect_equal(xgb.dump(m1), xgb.dump(m2)) # fails
#> Error: Expected `xgb.dump(m1)` to equal `xgb.dump(m2)`.
#> Differences:
#> 8/16 mismatches
#> x[4]: "3:leaf=0.541463435"
#> y[4]: "3:leaf=0.360770881"
#> 
#> x[7]: "11:leaf=-0.41538465"
#> y[7]: "11:leaf=-0.712710977"
#> 
#> x[8]: "12:leaf=0.495652169"
#> y[8]: "12:leaf=0.323709249"
#> 
#> x[9]: "8:leaf=-0.594997048"
#> y[9]: "8:leaf=-1.10756671"
#> 
#> x[12]: "9:leaf=0.582978785"
#> y[12]: "9:leaf=0.395674318"

# try model m3 with base_margin = 0 => same as base_score = 0.5!!!
xgboost::setinfo(xgb_data_train, "base_margin", rep(0, length(train_ids)))
#> [1] TRUE

m3 = xgboost::xgb.train(
  data = xgb_data_train,
  # base_score doesn't matter
  params = list(base_score = 0.2, objective = "binary:logistic"),
  nrounds = 1)
testthat::expect_equal(xgb.dump(m1), xgb.dump(m3)) # models ARE THE SAME!

# PREDICT ----
# predict with no base_margin
p1 = predict(m1, newdata = xgb_data_test)
# predict with no base_margin, different model
p2 = predict(m2, newdata = xgb_data_test)
# predict with no base_margin, same model as m1 but trained with an offset (base_margin = 0)
p3 = predict(m3, newdata = xgb_data_test)
# ALL results different!!!
expect_false(all(p1 == p2))
expect_false(all(p1 == p3)) # WHY??? same models here!!! and same test data?

# add base_margin = 0 to predictions
xgboost::setinfo(xgb_data_test, "base_margin", rep(0, length(test_ids)))
#> [1] TRUE
# predict with base_margin
p11 = predict(m1, newdata = xgb_data_test)
# predict with base_margin, different model
p22 = predict(m2, newdata = xgb_data_test)
# predict with base_margin, same model as m1 but trained with an offset (base_margin = 0)
p33 = predict(m3, newdata = xgb_data_test)

expect_equal(p1, p11)
expect_equal(p2, p22)
expect_equal(p3, p33)
expect_false(all(p11 == p33)) # WHY???

^{Created on 2025-12-12 with reprex v2.1.1}

[trivialfis]

Thank you for sharing, I will look into it.

[Sanidhyavijay24]

From my understanding (I might be wrong though) the trees are identical (xgb.dump passes) because during training, the algorithm looks at the "starting point" to calculate gradients.

• m1: Uses base_score=0.5 (which is a logit/raw score of 0).
• m3: Uses base_margin=0 (which is also a logit/raw score of 0).

Since the math starts from 0 in both cases, the trees learn the exact same residual values. That's why your expect_equal on the dump passes.

BUT during predictions (even though m3 ignored base_score=0.2 during training (because you provided a margin)), that 0.2 value was still saved inside the model's metadata. So when you call prdict() later:

• m1 calculates: Trees + 0.5 (The default score)
• m3 calculates: Trees + 0.2 (The parameter you passed)

The base_score acts as a global intercept. Unless you explicitly handle the margin correctly during prediction to strictly override this, that 0.2 drift causes the difference you are seeing.

This is from what I understand and as I mentioned I might have understood it wrongly.

[bblodfon]

I also think that the two models are the same after training. Your explanation is great because it gives a reason why $p1 \ne p3$ when we use a xgb.DMatrix test data matrix with no base_margin. But at the end of the code example I add the base_margin to the test data:

xgboost::setinfo(xgb_data_test, "base_margin", rep(0, length(test_ids)))

and the doc in the predict.xgb.Booster says that it is better to be given in the xgb.DMatrix newdata argument, ie as it has been done in the example. But then it is ignored in model m1 because of the base_score = 0.2, is that what is happening?

---

The other thing that I just saw in the doc of xgboost::xgb.train and xgb.params, is that if base_score (and base_margin) are not explicitly specified, at least during training, the base_score will be estimated. But from my example, I don't think this affects here?

---

In the below example, I tried to summarize the two observations above + show that the base_margin during prediction is probably not taken into account? (as I use exactly the same model):

library(xgboost)
library(testthat)

# Load data, split train/test
data(agaricus.train, package = "xgboost")
X = agaricus.train$data
y = agaricus.train$label

set.seed(42)
train_ids = sample(x = 1:nrow(X), size = 4000)
test_ids = setdiff(1:nrow(X), train_ids)

xgb_data_train = xgboost::xgb.DMatrix(data = X[train_ids, ], label = y[train_ids])
xgb_data_test  = xgboost::xgb.DMatrix(data = X[test_ids, ], label = y[test_ids])

# m1 model with no base_margin (default base_score = 0.5)
m1 = xgboost::xgb.train(
  data = xgb_data_train,
  params = list(base_score = 0.5, objective = "binary:logistic"),
  nrounds = 1)
# same as base_score = 0.5
xgboost::setinfo(xgb_data_train, "base_margin", rep(0, length(train_ids)))
#> [1] TRUE
m3 = xgboost::xgb.train(
  data = xgb_data_train,
  # base_score is estimated???
  params = list(objective = "binary:logistic"),
  nrounds = 1)
testthat::expect_equal(xgb.dump(m1), xgb.dump(m3)) # same models

# predictions different though => m3 has estimated base_margin != 0.5, this could explain this
p1 = predict(m1, newdata = xgb_data_test)
p3 = predict(m3, newdata = xgb_data_test)
head(p1-p3)
#> [1] 0.01767260 0.01722372 0.01722890 0.01722890 0.01722372 0.01767260

# now force base_margin = 0 for the test set, is it used?
xgboost::setinfo(xgb_data_test, "base_margin", rep(0, length(test_ids)))
#> [1] TRUE
xgb_data_test
#> xgb.DMatrix  dim: 2513 x 126  info: base_margin, label  colnames: yes
p1 = predict(m1, newdata = xgb_data_test)
p3 = predict(m3, newdata = xgb_data_test)
head(p1-p3) # maybe not used??????
#> [1] 0.01767260 0.01722372 0.01722890 0.01722890 0.01722372 0.01767260

# different base_margins from the same model, but same predictions ???
xgboost::setinfo(xgb_data_test, "base_margin", rep(0.5, length(test_ids)))
#> [1] TRUE
p11 = predict(m1, newdata = xgb_data_test)
xgboost::setinfo(xgb_data_test, "base_margin", rep(1.5, length(test_ids)))
#> [1] TRUE
p111 = predict(m1, newdata = xgb_data_test)
expect_equal(p1, p11) # I don't think these should be equal
expect_equal(p11, p111)

^{Created on 2025-12-14 with reprex v2.1.1}

[Sanidhyavijay24]

ohhh , looking at your example it seems that changing the base_margin in newdata from 0.5 to 1.5 results in the exact same predictions (expect_equal(p11, p111) passes), it proves that predict() is silently ignoring the base_margin provided in the xgb.DMatrix.

This explains some confusion:

• Since the test-time base_margin is being ignored by the code, the only difference left is the internal base_score of the models(between p1 and p3).
• m1 has a fixed base_score=0.5.
• m3 (where base_score wasn't set but training margin was used) likely defaulted to or estimated a slightly different global intercept.
• Because the test margin is ignored, you are seeing Trees + 0.5 vs Trees + EstimatedScore, hence the 0.017... difference.

This is probably a bug in the R wrapper's prediction logic (it seems to fails in passing the margin info to the C++ backend) (Though I haven't checked the code yet).

[Sanidhyavijay24]

@bblodfon Thanks for the last example. The fact that p11 == p111 (identical predictions despite different margins in the DMatrix) is definitive proof that the current workflow is broken.

@trivialfis I dug into R-package/R/xgb.Booster.R , and I found that predict.xgb.Booster explicitly forbids passing base_margin when the input is a xgb.DMatrix.

The Blocking Code (Lines 365 - 370):

if (is_dmatrix && !is.null(base_margin)) {
  stop(
    "'base_margin' is not supported when passing 'xgb.DMatrix' as input.",
    " Should be passed as argument to 'xgb.DMatrix' constructor."
  )
}

This forces users to rely on external setinfo calls, which—as the reproduction script shows—are seemingly not being picked up correctly during the prediction call sequence.

The Fix -> I propose we remove this restriction. If a user provides base_margin in predict(), we should simply apply it to the DMatrix internally using setinfo before calling the backend. This would make DMatrix behaviour consistent with dense matrices and fix the issue.

I can open a PR with this fix and a regression test if that sounds good!

[bblodfon]

Hi @Sanidhyavijay24, thanks for looking into this.

I don't know the internals of xgboost but I don't see how removing this restriction will solve the problem as the example code does not enter the if loop at all. The fix should address the case in which we have a xgb.DMatrix with base_margin already set inside and the code uses it for prediction correctly (and doesn't ignore it)!

[Sanidhyavijay24]

Ah, you are absolutely right @bblodfon! I realized I only explained the "unblocking" part in my previous comment, which caused the confusion.
My bad—that definitely wouldn't solve your specific case on its own.

Here is the full picture of the fix I'm thinking of:

• Remove the restriction: So users can manually pass/override the margin if they want to.
• Add Auto-detection (The part I missed explaining): I need to update predict() so that if base_margin is NULL (your case), the code explicitly checks inside the DMatrix using getinfo(newdata, "base_margin").

Currently, the wrapper defaults to NULL and blindly assumes "No Margin." My fix adds a step to "peek" inside the DMatrix. If it finds a margin there, it pulls it out and passes it to the backend.

# If user didn't provide a margin argument, check if the DMatrix has one hidden inside
if (is.null(base_margin) && inherits(newdata, "xgb.DMatrix")) {
   internal_margin <- getinfo(newdata, "base_margin")
   if (length(internal_margin) > 0) {
     base_margin <- internal_margin
   }
}

This ensures that the setinfo call you made earlier is actually respected during prediction, while still allowing other users to override it if they need to.

[bblodfon]

I can also confirm that using base_margin as a parameter in predict.xgb.Booster() with a simple newdata matrix (which is what we do in mlr3) works just fine - the issue is with the xgb.DMatrix format as shown above. Also tested the correspondence between base_score and base_margin, see tests in mlr3learners here.

[trivialfis]

Apologies, was heads deep in #11878 (which turns out to be a disappointment ...)

as the reproduction script shows—are seemingly not being picked up correctly during the prediction call sequence.

This seems like a bug. @david-cortes would you be able to help take a look when you are available (I can try to debug it as well if it's holiday for you).

the code explicitly checks inside the DMatrix using getinfo(newdata, "base_margin").

That's weird, I'm not sure why the base margin is not used. Let me double check.

[david-cortes]

I'm not understanding the issue here. You have two models where you are explicitly passing base_score=0.2 and base_score=0.5. It's expected that predictions from them should differ. If building two equivalent models where one uses base_score and the other a constant base_margin, their dumps from xgb.dump will not differ even if the models do, because they don't show all the information.

For a full comparison, you'd need to use xgb.save.raw. You can check the actual base_score like this:

xgboost::xgb.save.raw(m3, raw_format="json") |>
    rawToChar() |>
    jsonlite::fromJSON() |>
    _$learner$learner_model_param$base_score

From there, you can verify that m1 and m3 in your second example are not equivalent.

The only issue I'm seeing here is that setting the base_margin outside of the DMatrix constructor after a DMatrix has been used for predictions will result in subsequent calls igoring it, which I understood from some comments from @trivialfis might be due to caching results and likely also affects the Python interface.

But it does get used if you use a new DMatrix object:

> head(predict(m1, xgb.DMatrix(X)))
[1] 0.6412865 0.3557142 0.3557142 0.6412865 0.3554891 0.3557142
> head(predict(m1, xgb.DMatrix(X, base_margin=rep(1, nrow(X)))))
[1] 0.8293393 0.6001248 0.6001248 0.8293393 0.5998891 0.6001248

[trivialfis]

might be due to caching results and likely also affects the Python interface.

That would be a problem. Modifying an existing DMtraix violates the prediction cache.

I can add some checks to ensure this does not happen, but I'm not sure if there's a good use case for it.