Árboles de decisión: veinte preguntas con datos

¿Recuerdas el juego de las veinte preguntas? «¿Está vivo?», «¿Es más grande que un gato?». Un buen jugador no hace preguntas al azar; hace la que divide las posibilidades restantes más o menos por la mitad. Un árbol de decisión es un modelo que juega a esto con datos y, como sus preguntas están escritas en lenguaje llano, es uno de los pocos modelos que se pueden leer como un diagrama de flujo. Pruébalo tú mismo antes de leer la explicación.

Prueba esto

• Pon el control de profundidad máxima en 1. El árbol solo puede hacer una pregunta. Léela en la lista de divisiones: ¿qué característica consideró más informativa el algoritmo?
• Sube la profundidad máxima de uno en uno. Observa cómo sube la exactitud de entrenamiento y compárala con la de prueba en cada profundidad.
• Lleva la profundidad máxima al tope. ¿La exactitud de entrenamiento se acerca a la perfección? ¿Qué le pasó a la exactitud de prueba?
• Encuentra la profundidad con la mejor exactitud de prueba y lee las divisiones principales. ¿Podrías explicarle este árbol a alguien en dos frases?

La idea en pocas palabras

Un árbol de decisión predice haciendo una secuencia de preguntas de sí o no sobre las características: «¿La edad del pasajero es menor que cierto valor?», «¿Viajaba en primera clase?». Cada respuesta envía el ejemplo por una rama, y la siguiente pregunta depende de las respuestas anteriores. Al final de cada camino hay una hoja con una predicción, normalmente la etiqueta más común entre los ejemplos de entrenamiento que acabaron allí.

En esta lección usamos datos de 891 pasajeros del Titanic, de los que sobrevivieron 342. Las características incluyen la clase en que viajaban, el sexo, la edad, el número de familiares a bordo y la tarifa pagada. El árbol aprende qué preguntas sobre estas características separan mejor a quienes sobrevivieron de quienes no.

Piensa en una enfermera de triaje en urgencias. No calcula una suma ponderada; sigue una cadena corta de preguntas —¿respira?, ¿sangra?, ¿está consciente?—, cada una de las cuales acota qué hacer. Un árbol es esa cadena, aprendida de los datos en lugar de escrita por médicos.

Por qué funciona

El árbol se construye desde arriba, pregunta a pregunta. En cada paso, el algoritmo prueba cada característica y cada valor de corte posible y se pregunta: ¿qué pregunta divide este grupo en dos grupos lo más puros posible, es decir, que cada uno contenga sobre todo una etiqueta? Una medida popular de impureza es el índice de Gini de un grupo:

Gini = 1 − Σ pₖ²

donde pₖ es la proporción de la etiqueta k en el grupo. Un grupo en el que todos tienen la misma etiqueta tiene Gini 0; un grupo repartido mitad y mitad entre dos etiquetas tiene Gini 0,5, el peor valor posible con dos clases. El algoritmo elige la pregunta que más reduce la impureza media y luego repite el procedimiento por separado dentro de cada rama. Esta búsqueda voraz, paso a paso, no garantiza el mejor árbol posible, pero es rápida y funciona bien en la práctica.

Si lo dejas solo, el árbol sigue dividiendo hasta que cada hoja es pura, aunque una hoja contenga un solo pasajero. En ese punto ha memorizado los datos de entrenamiento, ruido incluido: el sobreajuste de la lección 7 con otro disfraz. El control principal es la profundidad máxima, el número de preguntas permitidas a lo largo de cualquier camino. Un árbol poco profundo es fácil de leer, pero puede pasar por alto patrones reales; uno profundo se ajusta casi a la perfección a los datos de entrenamiento y generaliza peor. La validación cruzada de la lección 8 es la forma honesta de elegir la profundidad.

Error frecuente

Una división no es una explicación causal. El árbol pregunta por la característica que mejor separa las etiquetas en estos datos, y esa puede ser un sustituto de otra cosa: la clase y la tarifa van juntas, y ambas dicen algo sobre en qué parte del barco dormía el pasajero. El segundo error es fiarse de un árbol profundo porque sus reglas parecen muy precisas. Una regla como «edad entre 31,5 y 32,5 y tarifa por encima de cierto valor» suele ser el árbol describiendo a uno o dos pasajeros, no un patrón. Pequeños cambios en los datos también pueden cambiar la primera división y, con ella, todo el árbol, lo que hace inestables a los árboles individuales.

Esa inestabilidad es una debilidad, pero puede convertirse en una fortaleza. En la próxima lección hacemos crecer muchos árboles distintos y los dejamos votar.

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

df = pd.read_csv("data/raw/titanic.csv")
X = pd.DataFrame({"female": (df["sex"] == "female").astype(int)}).join(df[["pclass", "age", "sibsp", "parch", "fare"]])
Xtr, Xte, ytr, yte = train_test_split(X, df["survived"], test_size=0.3, stratify=df["survived"], random_state=1)
Xtr, Xte = Xtr.fillna({"age": Xtr["age"].median()}), Xte.fillna({"age": Xtr["age"].median()})  # mediana del entrenamiento para ambos
print(f"passengers {len(df)}, survived {df['survived'].sum()}")
for d in (1, 3, 20):
    t = DecisionTreeClassifier(max_depth=d, random_state=0).fit(Xtr, ytr)
    print(f"depth {d:2d}: train {t.score(Xtr, ytr):.3f}, test {t.score(Xte, yte):.3f}, leaves {t.get_n_leaves()}")

passengers 891, survived 342
depth  1: train 0.780, test 0.802, leaves 2
depth  3: train 0.833, test 0.802, leaves 8
depth 20: train 0.986, test 0.780, leaves 141