Steering vectors: Una mirada al interior.
Sobre lo ultimo en alignment y el futuro para controlar LLMs.
¡Hola a todos! Soy Luis Ibáñez-Lissen, un apasionado de la ciberseguridad y el deep learning, y estoy terminando mi doctorado en este tema. Junto a mi colega Iván, hemos decidido lanzar una serie de publicaciones en formato de notas personales para compartir nuestras lecturas, investigaciones y cualquier tema que nos haya despertado curiosidad, porque, sinceramente, ¡somos muy inquietos!
En este primer post quiero profundizar en conceptos que me han captado mucho la atención. Actualmente, en los laboratorios de Anthropic y OpenAI se están desarrollando avances importantes que, creo, están marcando el rumbo de la IA controlable, o al menos, de las LLMs.
Para este post, voy a asumir el conocimiento básico en LLMs, cómo se pre-entrenan y los distintos elementos que conforman estos modelos. Si no, te recomiendo echarle un ojo al paper original que lo cambió todo en el mundo del procesamiento del lenguaje natural, “Attention is all you need“ .
Aun así, si no tienes tiempo para leer el paper o bien no te interesa saber tan en profundidad, voy a tratar de simplificar la idea principal sobre la que se está asentado el futuro del alignment. Pero ¿Qué es alignment?
Alignment es un concepto, a priori, muy sencillo, se trata de asegurar o conseguir que los modelos de Deep-learning, principalmente generativos, sean capaces de generar salidas alineadas con una serie de características predefinidas.
Por ejemplo, queremos hacer un modelo que genera imágenes, pero no nos gustan los perros, por lo tanto, quiero limitar, o asegurarme que nunca o casi nunca, se generen situaciones en las cuales, a través de la interacción del usuario con el modelo, por ejemplo, en un chat a través de prompts, el modelo las genere.
En el caso de los modelos de lenguaje, es mucho más interesante, porque la cantidad de prompts o situaciones que pueden hacer que un modelo “salga“ de una linea de actuación (a.k.a Jailbreak) son potencialmente infinitas (Idea loca: mira este paper donde hace Jailbreaks con ASCII art).
Como me encanta la ciberseguridad, voy a exponer dos puntos de vista, una de los “buenos” y otro de los “malos“:
A favor, limitar lenguaje potencialmente nocivo.
En contra, inducir al modelo cierto sesgo para “convencer“ o hacer dudar a una víctima, es decir maltrato psicológico.
Dada esta sencilla introducción, con esta idea en la cabeza, comenzaremos a adentrarnos en la parte más técnica.
Como probablemente sabes, los modelos de redes neuronales están hechos de varias capas. Las entradas pasan por estas capas, y en cada una, hay "neuronas" que se activan dependiendo de ciertos valores. Estos valores están controlados por "pesos", que se ajustan durante el entrenamiento usando un método llamado backpropagation. Además, cada neurona usa una función de activación que decide si esa neurona se "enciende" o no, en función del valor que recibe.
Pues bien, las últimas investigaciones apuntan que, hay cierta linealidad, es decir, que las LLMs codifican conceptos internamente que son “linearly separable” en su espacio de muestra interno. Esto se refiere a que algunos conceptos pueden aislarse o diferenciarse usando direcciones o vectores específicos y por tanto, si fuéramos capaces de encontrarlos, podríamos inducir un sesgo hacia esa dirección o concepto. ¿Interesante no?
Os pongo un ejemplo que me gustó mucho. En este paper, la hipótesis principal se basa en que, según estas ideas, uno puede llegar a estudiar si un modelo está bajo los efectos de jailbreak o no, siempre y cuando existiese el concepto “You are under a jailbreak“ dentro del modelo.
Y aquí va la pregunta, ¿Cómo podemos sacar ventaja de todo esto?
Pues la respuesta es sencilla, si asumimos que las representaciones internas “apuntan“ en una dirección, se podría llegar a calcular una dirección “contraria” a ese concepto para así maximizar o reducir su influencia. Y aquí es donde entran los steering vectors.
En el caso anterior, podemos llegar a analizar los casos en los que el modelo estaba bajo el backdoor, hacer una media de los valores de las activaciones en estos casos, intervenir el modelo durante la inferencia y añadir estos valores, de esta forma, podemos introducir un sesgo positivo que minimize potencialmente posibles backdoors.
El concepto “estar bajo un backdoor“ es en sí mismo bastante abstracto, y por mis pruebas, funcionan mejor con otros conceptos tipo: “No mentir“, “Estar contento“, “Hablar en español“. Jugando con este tipo de conceptos, se puede alterar en tiempo real el comportamiento de un modelo e incluso, parar ejecuciones de forma condicional.
A nivel de arquitectura, normalmente se intervienen modelos tras cada una de las multilayer perceptron que componen los bloques de un transformer.
Y ahora, ¡Manos a la obra!
A continuación voy a dejar un poco de código de ejemplo de cómo se puede llevar a cabo. Existen distintos niveles de abstracción, que facilitan la tarea de la intervención, desde usar wrappers y usar directamente pytorch, jugando con las activaciones del último token o usar otras librerías como nnsight o baukit.
En mi caso usaré baukit, porque es bastante sencillo de entender pero lo suficiente “vanilla“ para ocultar lo que estamos haciendo realmente.
Lo primero, sería extraer el vector del concepto que queremos introducir en el modelo, en este caso, queremos que el modelo hable de forma contento.
with TraceDict(model, layers=hook_layers, retain_input=True, retain_output=True) as rep:
for i in range(len( model.model.layers)):
module = model.model.layers[i]
inputs = tokenizer('happy', return_tensors="pt").to(device)
with Trace(module) as cache:
_ = model(**inputs)
act_happy = cache.output[0]
act_happy = act_happy.detach().cpu().numpy()
modules.append(act_happy[:,-1:,:])Para almacenar el vector de “happy“, lo guardaremos en la variable act_happy.
Luego, aplicamos un wrapper a la capa que queremos modificar y definimos un coeficiente, que actuará como la “intensidad” de nuestro vector “happy”. Con esto, mandamos un prompt al modelo para ver los efectos.
Además, he añadido la variable top_neurons_to_affect, que nos permite controlar el número de activaciones que queremos modificar. La idea es intervenir solo en las k activaciones más importantes.
chat = [
{ "role": "user", "content": "Hello, Tell me what you think of madrid?" },
]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
coeff = 0.9
for i in range(2):
steering_vecF = torch.tensor(modules[i]).to('cuda:1')
module = model.model.layers[i]
with Trace(module, edit_output=act_add(coeff*steering_vecF,top_neurons_to_affect)) as _:
outputs = model.generate(input_ids=inputs.to('cuda:1'), max_new_tokens=50)
print(tokenizer.decode(outputs[0]))Los resultados dependerán de factores como el coeficiente que usamos y la capa afectada, así que es importante experimentar con el código y ajustar el modelo según nuestras necesidades.
Por ejemplo, usando el vector de la palabra “happy” en la capa 1, obtenemos una respuesta como:
“¡Qué bueno saber que estás disfrutando de Madrid! 🎨😊”
En cambio, si aplicamos un vector de “sad”, obtenemos:
“No tengo la capacidad de formar opiniones personales sobre Madrid.”
Como ves, esto impacta en la salida del modelo, pero aún hay que ajustar la granularidad para intervenir solo en las activaciones más significativas. ¡Ahora os toca a vosotros afinar estos detalles!
¡Os dejamos el link del repo! → Codigo
Espero que os haya gustado, si tenéis dudas estaremos encantados de responder los comentarios, ¡nos vemos en el siguiente post!