diff --git a/econml/orf/_ortho_forest.py b/econml/orf/_ortho_forest.py
index 8a64692cb..154537541 100644
--- a/econml/orf/_ortho_forest.py
+++ b/econml/orf/_ortho_forest.py
@@ -885,8 +885,10 @@ class DROrthoForest(BaseOrthoForest):
     model_Y :  estimator, default sklearn.linear_model.LassoCV(cv=3)
         Estimator for learning potential outcomes at each leaf.
         Will be trained on features, controls and one hot encoded treatments (concatenated).
-        If different models per treatment arm are desired, see the :class:`.MultiModelWrapper`
-        helper class. The model(s) must implement `fit` and `predict` methods.
+        If different models per treatment arm are desired, wrap them in
+        :class:`.MultiModelWrapper` with ``encoding='full'`` (DROrthoForest
+        passes a full one-hot encoding rather than a drop-first one to
+        ``model_Y``). The model(s) must implement `fit` and `predict` methods.
 
     propensity_model_final : estimator, optional
         Model for estimating propensity of treatment at at prediction time.
@@ -896,8 +898,10 @@ class DROrthoForest(BaseOrthoForest):
     model_Y_final : estimator, optional
         Estimator for learning potential outcomes at prediction time.
         Will be trained on features, controls and one hot encoded treatments (concatenated).
-        If different models per treatment arm are desired, see the :class:`.MultiModelWrapper`
-        helper class. The model(s) must implement `fit` and `predict` methods.
+        If different models per treatment arm are desired, wrap them in
+        :class:`.MultiModelWrapper` with ``encoding='full'`` (DROrthoForest
+        passes a full one-hot encoding rather than a drop-first one to
+        ``model_Y``). The model(s) must implement `fit` and `predict` methods.
         If parameter is set to ``None``, it defaults to the value of `model_Y` parameter.
 
     categories: 'auto' or list
@@ -1082,7 +1086,11 @@ def const_marginal_ate(self, X=None):
     def nuisance_estimator_generator(propensity_model, model_Y, random_state=None, second_stage=False):
         """Generate nuissance estimator given model inputs from the class."""
         def nuisance_estimator(Y, T, X, W, sample_weight=None, split_indices=None):
-            # Expand one-hot encoding to include the zero treatment
+            # Expand one-hot encoding to include the zero treatment.
+            # TODO: consider a breaking change to harmonize with the rest of EconML's discrete-treatment
+            # estimators, which feed the drop-first one-hot encoding directly to downstream models rather
+            # than re-adding the dropped column here. That would let model_Y use the same
+            # ``encoding='drop_first'`` convention as MultiModelWrapper's default.
             ohe_T = np.hstack([np.all(1 - T, axis=1, keepdims=True), T])
             # Test that T contains all treatments. If not, return None
             T = ohe_T @ np.arange(ohe_T.shape[1])
diff --git a/econml/tests/test_model_selection.py b/econml/tests/test_model_selection.py
index f01e94f30..a08373a40 100644
--- a/econml/tests/test_model_selection.py
+++ b/econml/tests/test_model_selection.py
@@ -12,7 +12,7 @@
 from sklearn.preprocessing import PolynomialFeatures
 from econml.dml import LinearDML
 from econml.sklearn_extensions.linear_model import WeightedLassoCVWrapper
-from econml.utilities import SeparateModel
+from econml.utilities import MultiModelWrapper
 from econml.dr import LinearDRLearner
 
 
@@ -139,13 +139,13 @@ def test_sklearn_model_selection(self):
     def test_fixed_model_scoring(self):
         Y, T, X, W = self._simple_dgp(500, 2, 3, True)
 
-        # SeparatedModel doesn't support scoring; that should be fine when not compared to other models
-        mdl = LinearDRLearner(model_regression=SeparateModel(LassoCV(), LassoCV()),
+        # MultiModelWrapper doesn't support scoring; that should be fine when not compared to other models
+        mdl = LinearDRLearner(model_regression=MultiModelWrapper(LassoCV(), LassoCV()),
                               model_propensity=LogisticRegressionCV())
         mdl.fit(Y, T, X=X, W=W)
 
         # on the other hand, when we need to compare the score to other models, it should raise an error
         with self.assertRaises(Exception):
-            mdl = LinearDRLearner(model_regression=[SeparateModel(LassoCV(), LassoCV()), Lasso()],
+            mdl = LinearDRLearner(model_regression=[MultiModelWrapper(LassoCV(), LassoCV()), Lasso()],
                                   model_propensity=LogisticRegressionCV())
             mdl.fit(Y, T, X=X, W=W)
diff --git a/econml/tests/test_utilities.py b/econml/tests/test_utilities.py
index 2af0ca666..51cb706c4 100644
--- a/econml/tests/test_utilities.py
+++ b/econml/tests/test_utilities.py
@@ -7,11 +7,13 @@
 import warnings
 import numpy as np
 import sparse as sp
+import scipy.sparse
 import pytest
 from econml.utilities import (check_high_dimensional, einsum_sparse, todense, tocoo, transpose,
                               inverse_onehot, cross_product, transpose_dictionary, deprecated, _deprecate_positional,
-                              strata_from_discrete_arrays)
+                              strata_from_discrete_arrays, MultiModelWrapper, SeparateModel)
 from sklearn.preprocessing import OneHotEncoder, SplineTransformer
+from sklearn.linear_model import LinearRegression, LogisticRegressionCV, LassoCV
 
 
 class TestUtilities(unittest.TestCase):
@@ -197,3 +199,251 @@ def test_single_strata_from_discrete_array(self):
         assert set(strata_from_discrete_arrays([T, Z])) == set(np.arange(6))
         assert set(strata_from_discrete_arrays([T])) == set(np.arange(3))
         assert strata_from_discrete_arrays([]) is None
+
+
+class TestMultiModelWrapper(unittest.TestCase):
+
+    @staticmethod
+    def _encode_drop_first(T, K):
+        out = np.zeros((T.shape[0], K - 1))
+        for i in range(1, K):
+            out[T == i, i - 1] = 1
+        return out
+
+    @staticmethod
+    def _encode_full(T, K):
+        out = np.zeros((T.shape[0], K))
+        for i in range(K):
+            out[T == i, i] = 1
+        return out
+
+    def test_drop_first_routes_rows_to_models(self):
+        rng = np.random.default_rng(0)
+        n, d, K = 90, 3, 3
+        X = rng.normal(size=(n, d))
+        T = rng.integers(0, K, size=n)
+        # ground truth: arm k has slope k on the first feature
+        Y = T * X[:, 0]
+        Xt = np.hstack([X, self._encode_drop_first(T, K)])
+
+        w = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+        )
+        w.fit(Xt, Y)
+        for k in range(K):
+            self.assertAlmostEqual(float(w.models[k].coef_[0]), float(k), places=6)
+        np.testing.assert_allclose(w.predict(Xt), Y, atol=1e-8)
+
+    def test_full_encoding_routes_rows_to_models(self):
+        rng = np.random.default_rng(1)
+        n, d, K = 60, 2, 3
+        X = rng.normal(size=(n, d))
+        T = rng.integers(0, K, size=n)
+        Y = (T + 1) * X[:, 0]
+        Xt = np.hstack([X, self._encode_full(T, K)])
+
+        w = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            encoding='full',
+        )
+        w.fit(Xt, Y)
+        for k in range(K):
+            self.assertAlmostEqual(float(w.models[k].coef_[0]), float(k + 1), places=6)
+        np.testing.assert_allclose(w.predict(Xt), Y, atol=1e-8)
+
+    def test_label_encoding_matches_drop_first(self):
+        rng = np.random.default_rng(2)
+        n, d, K = 80, 2, 3
+        X = rng.normal(size=(n, d))
+        T = rng.integers(0, K, size=n)
+        Y = T * X[:, 0]
+        Xt_lbl = np.hstack([X, T.reshape(-1, 1)])
+        Xt_df = np.hstack([X, self._encode_drop_first(T, K)])
+
+        w_lbl = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            encoding='label',
+        )
+        w_df = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+        )
+        w_lbl.fit(Xt_lbl, Y)
+        w_df.fit(Xt_df, Y)
+        for k in range(K):
+            np.testing.assert_allclose(w_lbl.models[k].coef_, w_df.models[k].coef_)
+        np.testing.assert_allclose(w_lbl.predict(Xt_lbl), w_df.predict(Xt_df))
+
+    def test_single_model_with_n_categories_clones(self):
+        w = MultiModelWrapper(LinearRegression(fit_intercept=False), n_categories=4)
+        self.assertEqual(w.n_categories, 4)
+        self.assertEqual(len(w.models), 4)
+        ids = {id(m) for m in w.models}
+        self.assertEqual(len(ids), 4)  # genuine clones, not the same instance
+
+    def test_single_model_without_n_categories_raises(self):
+        with self.assertRaises(ValueError):
+            MultiModelWrapper(LinearRegression())
+
+    def test_mismatched_n_categories_raises(self):
+        with self.assertRaises(ValueError):
+            MultiModelWrapper(LinearRegression(), LinearRegression(), n_categories=3)
+
+    def test_zero_models_raises(self):
+        with self.assertRaises(ValueError):
+            MultiModelWrapper()
+
+    def test_invalid_encoding_raises(self):
+        with self.assertRaises(ValueError):
+            MultiModelWrapper(LinearRegression(), LinearRegression(), encoding='bogus')
+
+    def test_too_few_columns_raises(self):
+        w = MultiModelWrapper(LinearRegression(), LinearRegression(), LinearRegression())
+        with self.assertRaises(ValueError):
+            # K=3 with default drop_first needs >= 2 trailing one-hot columns
+            w.fit(np.array([[1.0]]), np.array([0.0]))
+
+    def test_sample_weight_is_forwarded(self):
+        # With one heavily down-weighted point in arm 1, the fitted slope
+        # should be determined by the other arm-1 point alone.
+        X = np.array([[1.0], [2.0], [3.0], [4.0]])
+        T = np.array([0, 0, 1, 1])
+        Y = np.array([0.0, 0.0, 6.0, 100.0])  # (3, 6) -> slope 2; (4, 100) is noise
+        Xt = np.hstack([X, T.reshape(-1, 1).astype(float)])
+        sw = np.array([1.0, 1.0, 1.0, 1e-12])
+
+        w = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            encoding='label',
+        )
+        w.fit(Xt, Y, sample_weight=sw)
+        self.assertAlmostEqual(float(w.models[1].coef_[0]), 2.0, places=3)
+
+    def test_integration_with_linear_drlearner_multinary(self):
+        # End-to-end smoke test: drop-first MultiModelWrapper fed to LinearDRLearner
+        # with 3 treatment categories (the case the old MultiModelWrapper couldn't handle).
+        from econml.dr import LinearDRLearner
+        rng = np.random.default_rng(3)
+        n = 300
+        X = rng.normal(size=(n, 2))
+        T = rng.integers(0, 3, size=n)
+        Y = X[:, 0] + T * (1 + X[:, 1]) + rng.normal(size=n)
+
+        mdl = LinearDRLearner(
+            model_regression=MultiModelWrapper(LassoCV(), n_categories=3),
+            model_propensity=LogisticRegressionCV(max_iter=200),
+        )
+        mdl.fit(Y, T, X=X)
+        effects = mdl.effect(X[:5], T0=0, T1=1)
+        self.assertEqual(effects.shape, (5,))
+
+    def test_sparse_input_drop_first(self):
+        # csr_matrix with K=3, drop-first encoding: full matrix stays sparse,
+        # only the trailing 2-column treatment block is densified internally.
+        rng = np.random.default_rng(10)
+        n, d, K = 60, 4, 3
+        X = rng.normal(size=(n, d))
+        T = rng.integers(0, K, size=n)
+        Y = T * X[:, 0]
+        Xt = np.hstack([X, self._encode_drop_first(T, K)])
+        Xt_sparse = scipy.sparse.csr_matrix(Xt)
+
+        w = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+        )
+        w.fit(Xt_sparse, Y)
+        for k in range(K):
+            self.assertAlmostEqual(float(w.models[k].coef_[0]), float(k), places=6)
+        np.testing.assert_allclose(w.predict(Xt_sparse), Y, atol=1e-8)
+
+    def test_sparse_input_full_encoding(self):
+        rng = np.random.default_rng(11)
+        n, d, K = 50, 3, 2
+        X = rng.normal(size=(n, d))
+        T = rng.integers(0, K, size=n)
+        Y = (T + 1) * X[:, 0]
+        Xt = np.hstack([X, self._encode_full(T, K)])
+        Xt_sparse = scipy.sparse.csr_matrix(Xt)
+
+        w = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            encoding='full',
+        )
+        w.fit(Xt_sparse, Y)
+        for k in range(K):
+            self.assertAlmostEqual(float(w.models[k].coef_[0]), float(k + 1), places=6)
+        np.testing.assert_allclose(w.predict(Xt_sparse), Y, atol=1e-8)
+
+    def test_sparse_input_label_encoding(self):
+        rng = np.random.default_rng(12)
+        n, d, K = 70, 2, 3
+        X = rng.normal(size=(n, d))
+        T = rng.integers(0, K, size=n)
+        Y = T * X[:, 0]
+        Xt = np.hstack([X, T.reshape(-1, 1).astype(float)])
+        Xt_sparse = scipy.sparse.csr_matrix(Xt)
+
+        w = MultiModelWrapper(
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            LinearRegression(fit_intercept=False),
+            encoding='label',
+        )
+        w.fit(Xt_sparse, Y)
+        for k in range(K):
+            self.assertAlmostEqual(float(w.models[k].coef_[0]), float(k), places=6)
+        np.testing.assert_allclose(w.predict(Xt_sparse), Y, atol=1e-8)
+
+    def test_model_list_kwarg_is_deprecated(self):
+        # Old API: MultiModelWrapper(model_list=[...]) should still work but
+        # emit a FutureWarning and produce a wrapper equivalent to the new
+        # positional-args form.
+        with self.assertWarnsRegex(FutureWarning, "model_list"):
+            w = MultiModelWrapper(
+                model_list=[LinearRegression(fit_intercept=False),
+                            LinearRegression(fit_intercept=False)],
+            )
+        self.assertEqual(w.n_categories, 2)
+        # Smoke-fit to make sure the wrapper is fully functional. Default
+        # encoding is 'drop_first', so we one-hot-drop-first the binary T.
+        X = np.array([[1.0], [2.0], [3.0], [4.0]])
+        T = np.array([0, 0, 1, 1])
+        Y = np.array([0.0, 0.0, 6.0, 8.0])
+        Xt = np.hstack([X, (T == 1).reshape(-1, 1).astype(float)])
+        w.fit(Xt, Y)
+        np.testing.assert_allclose(w.predict(Xt), Y, atol=1e-8)
+
+    def test_positional_list_is_deprecated(self):
+        # Old API: MultiModelWrapper([m1, m2, ...]) should warn and unpack.
+        with self.assertWarnsRegex(FutureWarning, "single positional argument"):
+            w = MultiModelWrapper(
+                [LinearRegression(fit_intercept=False),
+                 LinearRegression(fit_intercept=False),
+                 LinearRegression(fit_intercept=False)],
+            )
+        self.assertEqual(w.n_categories, 3)
+
+    def test_model_list_and_positional_models_together_raises(self):
+        with self.assertRaises(ValueError):
+            MultiModelWrapper(
+                LinearRegression(),
+                model_list=[LinearRegression(), LinearRegression()],
+            )
+
+
+class TestSeparateModelDeprecation(unittest.TestCase):
+
+    def test_separate_model_emits_future_warning(self):
+        with self.assertWarnsRegex(FutureWarning, "SeparateModel is deprecated"):
+            SeparateModel(LinearRegression(), LinearRegression())
diff --git a/econml/utilities.py b/econml/utilities.py
index 5673771c8..636884ea3 100644
--- a/econml/utilities.py
+++ b/econml/utilities.py
@@ -1044,62 +1044,189 @@ def _sampled_inputs(self, X, y, sample_weight):
 
 
 class MultiModelWrapper:
-    """Helper class for training different models for each treatment.
+    """Train a separate underlying model for each treatment category.
+
+    Wraps a collection of base estimators, dispatching each row to the
+    estimator for its treatment category, based on a one-hot encoding of the
+    treatment that is concatenated onto the right of the feature matrix.
+
+    Most EconML estimators internally one-hot-encode discrete treatments using
+    :class:`sklearn.preprocessing.OneHotEncoder` with ``drop='first'`` (the
+    control category becomes a row of all zeros and is not represented by an
+    explicit column) and concatenate that encoding to the right of the feature
+    matrix before handing it to first-stage models. ``MultiModelWrapper``'s
+    default ``encoding='drop_first'`` matches this convention, making it
+    suitable as a wrapper for first-stage models passed to e.g.
+    :class:`~econml.dr.DRLearner` and its subclasses. The alternative
+    ``encoding='full'`` expects a full ``K``-column one-hot encoding instead
+    (no dropped category), which is the convention used internally by
+    :class:`~econml.orf.DROrthoForest`'s ``model_Y``.
 
     Parameters
     ----------
-    model_list : array_like, shape (n_T, )
-        List of models to be trained separately for each treatment group.
-    """
-
-    def __init__(self, model_list=[]):
-        self.model_list = model_list
-        self.n_T = len(model_list)
+    *models : estimators
+        Either ``K`` estimators (one per treatment category, the first one
+        used for the control category), or a single estimator that will be
+        cloned ``n_categories`` times.
+
+    n_categories : int, optional
+        The number of treatment categories ``K``. Required when exactly one
+        positional model is passed (so it can be cloned ``K`` times). When
+        multiple positional models are passed, this must be either ``None``
+        or equal to the number of models.
+
+    encoding : {'drop_first', 'full', 'label'}, default 'drop_first'
+        The expected encoding of the treatment block at the right of the
+        input matrix.
+
+        - ``'drop_first'``: the last ``K - 1`` columns are a drop-first
+          one-hot encoding (the all-zero row is the control category).
+        - ``'full'``: the last ``K`` columns are a full one-hot encoding
+          (column ``0`` is the control category).
+        - ``'label'``: the last column is an integer category index in
+          ``{0, ..., K - 1}``.
+    """
+
+    def __init__(self, *models, n_categories=None, encoding='drop_first', model_list=None):
+        # Backward-compat shims for the pre-rewrite API
+        # (``MultiModelWrapper(model_list=[...])`` and the equivalent
+        # ``MultiModelWrapper([...])`` positional-list form). Both are
+        # accepted with a deprecation warning, and unpacked into ``models``
+        # so the rest of the constructor is unchanged.
+        if model_list is not None:
+            if models:
+                raise ValueError(
+                    "Cannot specify both positional models and the deprecated "
+                    "`model_list=` kwarg. Pass models positionally, e.g. "
+                    "MultiModelWrapper(*model_list).")
+            warnings.warn(
+                "The `model_list=` kwarg of MultiModelWrapper is deprecated and "
+                "will be removed in a future release; pass models as positional "
+                "arguments instead, e.g. MultiModelWrapper(*model_list).",
+                FutureWarning, stacklevel=2)
+            models = tuple(model_list)
+        elif len(models) == 1 and isinstance(models[0], (list, tuple)):
+            warnings.warn(
+                "Passing a list or tuple as a single positional argument to "
+                "MultiModelWrapper is deprecated and will be removed in a future "
+                "release; unpack with `*` instead, e.g. "
+                "MultiModelWrapper(*model_list).",
+                FutureWarning, stacklevel=2)
+            models = tuple(models[0])
+        if encoding not in ('drop_first', 'full', 'label'):
+            raise ValueError(
+                f"encoding must be 'drop_first', 'full', or 'label', got {encoding!r}.")
+        if len(models) == 0:
+            raise ValueError("MultiModelWrapper requires at least one model.")
+        if len(models) == 1:
+            if n_categories is None:
+                raise ValueError(
+                    "When a single model is passed, n_categories must be specified "
+                    "so the model can be cloned once per treatment category.")
+            if n_categories < 2:
+                raise ValueError("n_categories must be at least 2.")
+            self.models = [clone(models[0]) for _ in range(n_categories)]
+        else:
+            if n_categories is not None and n_categories != len(models):
+                raise ValueError(
+                    f"n_categories ({n_categories}) does not match the number of "
+                    f"positional models passed ({len(models)}).")
+            self.models = [clone(m) for m in models]
+        self.n_categories = len(self.models)
+        self.encoding = encoding
+
+    def _encoded_width(self):
+        if self.encoding == 'full':
+            return self.n_categories
+        if self.encoding == 'label':
+            return 1
+        return self.n_categories - 1
+
+    def _split(self, Xt):
+        """Return (X, labels) where labels is the per-row category index in [0, n_categories)."""
+        width = self._encoded_width()
+        if Xt.shape[1] < width:
+            raise ValueError(
+                f"Input has only {Xt.shape[1]} columns, but the trailing "
+                f"{width} columns are required to encode "
+                f"{self.n_categories} treatment categories with "
+                f"encoding={self.encoding!r}.")
+        if self.encoding == 'label':
+            X = Xt[:, :-1]
+            # The trailing column is a per-row category index; densify only
+            # that single column (cheap even for huge sparse inputs).
+            labels = todense(Xt[:, -1]).ravel().astype(int)
+            return X, labels
+        # Both 'drop_first' and 'full' reduce to a (K-1)-column drop-first
+        # one-hot encoding: for 'full' the leading control column drops out,
+        # leaving the trailing K-1 columns as exactly the drop-first form.
+        # ``inverse_onehot`` handles both dense and sparse inputs via matmul.
+        X = Xt[:, :-width]
+        labels = inverse_onehot(Xt[:, -(self.n_categories - 1):])
+        return X, labels
 
     def fit(self, Xt, y, sample_weight=None):
-        """Fit underlying list of models with weighted inputs.
+        """Fit each underlying model on its category's rows.
 
         Parameters
         ----------
-        X : array_like, shape (n_samples, n_features + n_treatments)
-            Training data. The last n_T columns should be a one-hot encoding of the treatment assignment.
-
-        y : array_like, shape (n_samples, )
+        Xt : array_like, shape (n_samples, n_features + width)
+            Training data. ``width`` depends on ``encoding``:
+            ``n_categories - 1`` for ``'drop_first'`` (the default),
+            ``n_categories`` for ``'full'``, or ``1`` for ``'label'``.
+            The trailing block must be a treatment encoding of the
+            corresponding form.
+
+        y : array_like
             Target values.
 
+        sample_weight : array_like, optional
+            Per-sample weights, forwarded to each underlying model.
+
         Returns
         -------
-        self: an instance of the class
+        self : MultiModelWrapper
         """
-        X = Xt[:, :-self.n_T]
-        t = Xt[:, -self.n_T:]
-        if sample_weight is None:
-            for i in range(self.n_T):
-                mask = (t[:, i] == 1)
-                self.model_list[i].fit(X[mask], y[mask])
-        else:
-            for i in range(self.n_T):
-                mask = (t[:, i] == 1)
-                self.model_list[i].fit(X[mask], y[mask], sample_weight[mask])
+        X, labels = self._split(Xt)
+        y = np.asarray(y)
+        if sample_weight is not None:
+            sample_weight = np.asarray(sample_weight)
+        for i, mdl in enumerate(self.models):
+            mask = labels == i
+            if sample_weight is None:
+                mdl.fit(X[mask], y[mask])
+            else:
+                mdl.fit(X[mask], y[mask], sample_weight=sample_weight[mask])
         return self
 
     def predict(self, Xt):
-        """Predict using the linear model.
+        """Predict by dispatching each row to its category's model.
 
         Parameters
         ----------
-        X : array_like, shape (n_samples, n_features + n_treatments)
-            Samples. The last n_T columns should be a one-hot encoding of the treatment assignment.
+        Xt : array_like, shape (n_samples, n_features + width)
+            Samples to predict for. See :meth:`fit` for the expected layout.
 
         Returns
         -------
-        C : array, shape (n_samples, )
-            Returns predicted values.
+        C : array, shape (n_samples, ...)
+            Predictions, with category-specific predictions reassembled in
+            the original row order.
         """
-        X = Xt[:, :-self.n_T]
-        t = Xt[:, -self.n_T:]
-        predictions = [self.model_list[np.nonzero(t[i])[0][0]].predict(X[[i]]) for i in range(len(X))]
-        return np.concatenate(predictions)
+        X, labels = self._split(Xt)
+        out = None
+        for i, mdl in enumerate(self.models):
+            mask = labels == i
+            if not np.any(mask):
+                continue
+            preds = mdl.predict(X[mask])
+            if out is None:
+                shape = (X.shape[0],) + np.shape(preds)[1:]
+                out = np.empty(shape, dtype=np.asarray(preds).dtype)
+            out[mask] = preds
+        if out is None:
+            return np.empty(0)
+        return out
 
 
 def _safe_norm_ppf(q, loc=0, scale=1):
@@ -1224,35 +1351,6 @@ def as_html(self):
         return html
 
 
-class SeparateModel:
-    """
-    Splits the data based on the last feature and trains a separate model for each subsample.
-
-    At predict time, it uses the last feature to choose which model to use to predict.
-    """
-
-    def __init__(self, *models):
-        self.models = [clone(model) for model in models]
-
-    def fit(self, XZ, T):
-        for (i, m) in enumerate(self.models):
-            inds = (XZ[:, -1] == i)
-            m.fit(XZ[inds, :-1], T[inds])
-        return self
-
-    def predict(self, XZ):
-        t_pred = np.zeros(XZ.shape[0])
-        for (i, m) in enumerate(self.models):
-            inds = (XZ[:, -1] == i)
-            if np.any(inds):
-                t_pred[inds] = m.predict(XZ[inds, :-1])
-        return t_pred
-
-    @property
-    def coef_(self):
-        return np.concatenate((model.coef_ for model in self.models))
-
-
 def deprecated(message, category=FutureWarning):
     """
     Enable decorating a method or class to providing a warning when it is used.
@@ -1319,6 +1417,47 @@ def m(*args, **kwargs):
     return decorator
 
 
+@deprecated(
+    "SeparateModel is deprecated and will be removed in a future release; "
+    "use MultiModelWrapper(..., encoding='label') instead, which provides "
+    "the same trailing-integer-column dispatch behavior."
+)
+class SeparateModel:
+    """
+    Splits the data based on the last feature and trains a separate model for each subsample.
+
+    At predict time, it uses the last feature to choose which model to use to predict.
+
+    .. deprecated::
+        Use :class:`MultiModelWrapper` with ``encoding='label'`` instead,
+        which has the same trailing-integer-column dispatch behavior.
+        ``MultiModelWrapper`` additionally supports the one-hot encodings
+        produced by EconML's discrete-treatment estimators (see its
+        ``encoding`` parameter).
+    """
+
+    def __init__(self, *models):
+        self.models = [clone(model) for model in models]
+
+    def fit(self, XZ, T):
+        for (i, m) in enumerate(self.models):
+            inds = (XZ[:, -1] == i)
+            m.fit(XZ[inds, :-1], T[inds])
+        return self
+
+    def predict(self, XZ):
+        t_pred = np.zeros(XZ.shape[0])
+        for (i, m) in enumerate(self.models):
+            inds = (XZ[:, -1] == i)
+            if np.any(inds):
+                t_pred[inds] = m.predict(XZ[inds, :-1])
+        return t_pred
+
+    @property
+    def coef_(self):
+        return np.concatenate((model.coef_ for model in self.models))
+
+
 class MissingModule:
     """
     Placeholder to stand in for a module that couldn't be imported, delaying ImportErrors until use.