<jupyter_start><jupyter_text>IntroductionThe goal of this notebook is to show how the output of a particular transformer vector depends on earlier positions and lower layers. My initial understanding is that the attention sublayer directly depends on earlier positions in the same layer (and indirectly on earlier positions in lower layers) and that the NN sublayer depends directly on the same position in lower layers (and indirectly on the earlier positions in lower layers.I am taking the lowest layer to be the one taking input, and the highest layer generating output.   Setup  Packages<jupyter_code>%pip install transformer_utils
%pip install captum<jupyter_output>Collecting transformer_utils
  Downloading transformer_utils-0.1.0-py3-none-any.whl (17 kB)
Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (from transformer_utils) (1.9.0+cu102)
Collecting transformers
  Downloading transformers-4.9.2-py3-none-any.whl (2.6 MB)
[K     |████████████████████████████████| 2.6 MB 10.0 MB/s 
[?25hRequirement already satisfied: seaborn in /usr/local/lib/python3.7/dist-packages (from transformer_utils) (0.11.1)
Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from transformer_utils) (4.62.0)
Requirement already satisfied: scipy>=1.0 in /usr/local/lib/python3.7/dist-packages (from seaborn->transformer_utils) (1.4.1)
Requirement already satisfied: matplotlib>=2.2 in /usr/local/lib/python3.7/dist-packages (from seaborn->transformer_utils) (3.2.2)
Requirement already satisfied: pandas>=0.23 in /usr/local/lib/python3.7/dist-packages (from seaborn->transformer_utils) (1.1.5)
Requirement already satisfie[...]<jupyter_text>Load Model<jupyter_code>from transformer_utils.low_memory import enable_low_memory_load
import transformers

enable_low_memory_load()
tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")
model = transformers.AutoModelForCausalLM.from_pretrained('EleutherAI/gpt-neo-125M')
import torch

def text_to_input_ids(text):
    toks = tokenizer.encode(text)
    return torch.as_tensor(toks).view(1, -1).cuda()<jupyter_output><empty_output><jupyter_text>Exploration  Basics<jupyter_code>text = "The forecast said that thunderstorm activity with the risk of hail will lead to localised"
tokens = text_to_input_ids(text)
tokens, tokens.shape, tokenizer.batch_decode(tokens[0][-4:])
result = model.forward(tokens)
result.logits.shape<jupyter_output><empty_output><jupyter_text>For each of the 17 positions, the Laguage Model makes a prediction of likelihood of each of the 50257 possible tokens being the next one. This graph shows an apparent trend, probably explained by earlier tokens being the more commonly used.<jupyter_code>import matplotlib.pyplot as plt
plt.plot(result.logits[0,16].cpu().detach().numpy())<jupyter_output><empty_output><jupyter_text>Argmax gives the reasonable prediction of 'flooding' in the last position.<jupyter_code>predicted = torch.argmax(result.logits[0,16])
predicted, tokenizer.decode(predicted), result.logits[0,16,predicted]<jupyter_output><empty_output><jupyter_text>Embedding VectorsBypassing language modelling so we can work directly with the transformer vectors, first confirm we get the same result:<jupyter_code>embeds = model.transformer.wte(tokens)
transformer_outputs = model.transformer.forward(inputs_embeds=embeds)
type(transformer_outputs)
transformer_outputs.last_hidden_state.shape

def show_prediction(transformer_outputs,posn):
  output_vec = transformer_outputs.last_hidden_state[0,posn]
  t_logits = model.lm_head(output_vec)
  t_predicted = torch.argmax(t_logits)
  print(t_predicted, tokenizer.decode(t_predicted), t_logits[t_predicted])

show_prediction(transformer_outputs,16)<jupyter_output>tensor(17448, device='cuda:0')  flooding tensor(-0.7156, device='cuda:0', grad_fn=<SelectBackward>)<jupyter_text>Plot Embedding Vector WaveformsThis compares the embedding vector generated from the input token to that output by the top layer, predicting the next token.<jupyter_code>output_vec=transformer_outputs.last_hidden_state[0,posn]
plt.plot(embeds[0,16].cpu().detach().numpy())
plt.show()
plt.plot(output_vec.cpu().detach().numpy())
plt.show()<jupyter_output><empty_output><jupyter_text>Attribution  Attribute a specific logit output to embedding vector inputs Now we can define a custom modified forward function which we can use for attribution for a specific output word vector. The model is wrapped in a function to allow transformer.forward be called with the inputs_embeds argument rather than default first arg, and to return the relevant logit as a single record 1d-tensor.<jupyter_code>def make_forward_with_embeds(target_position, target_tokenid):
    
  def forward_with_embeds(embeds):
      
      transformer_outputs = model.transformer.forward(inputs_embeds=embeds)
      output_vec = transformer_outputs.last_hidden_state[0,target_position]
      lm_logits = model.lm_head(output_vec)
      return lm_logits[target_tokenid:target_tokenid+1] #captum expects
    
  return forward_with_embeds<jupyter_output><empty_output><jupyter_text>Test again (expect -0.7156)<jupyter_code>make_forward_with_embeds( 16, 17448 )(embeds)<jupyter_output><empty_output><jupyter_text>Using the [Integrated Gradients](https://arxiv.org/pdf/1703.01365.pdf) algorithm through [Captum](https://captum.ai/docs/extension/integrated_gradients) we use attribution to identify what part of the embedding inputs are most important to the logit for the expected output term (flooding).<jupyter_code>from captum.attr import IntegratedGradients
def attribute(attributor):
  embeds = model.transformer.wte(tokens)
  embeds.requires_grad_()
  return attributor.attribute(embeds)

attr = attribute(IntegratedGradients(make_forward_with_embeds( 16, 17448 )))<jupyter_output><empty_output><jupyter_text>Input Position SalienceUsing norm on each vector in the input and plotting against the corresponding label gives a plausible interpretation of the importance of each term.<jupyter_code>import numpy as np

def show_posn_salience(attribution):
  norms = np.linalg.norm(attribution.cpu().detach().numpy(), axis=1)
  labels = [tokenizer.decode(token) for token in tokens[0]]

  fig, ax = plt.subplots()
  ax.plot(norms)
  ax.set_xticks(np.arange(attribution.shape[0]))

  ax.set_xticklabels(labels,rotation='vertical')
  plt.show()

show_posn_salience(attr[0])<jupyter_output><empty_output><jupyter_text>Attribution to a single embedding vectorThis shows attribution against the last token vector.<jupyter_code>plt.plot(attr[0,16].cpu().detach().numpy())<jupyter_output><empty_output><jupyter_text>Attribute specific logit excess to embedding vector inputsThe previous wrapper of the forward method would identify inputs as significant even if they increased logits for other symbols more than the target one. The version below outputs the difference between the target logit and the mean of all logits. Further investigation is required to see if this is helpful or not.<jupyter_code>def make_logit_less_mean_embeds(target_position, target_tokenid):
    
  def logit_less_mean_embeds(embeds):
      
      transformer_outputs = model.transformer.forward(inputs_embeds=embeds)
      output_vec = transformer_outputs.last_hidden_state[0,target_position]
      lm_logits = model.lm_head(output_vec)
      logits_mean = torch.mean(lm_logits)
      return lm_logits[target_tokenid:target_tokenid+1] -logits_mean
    
  return logit_less_mean_embeds

attr = attribute(IntegratedGradients(make_logit_less_mean_embeds( 16, 17448 )))

plt.plot(attr[0,16].cpu().detach().numpy())

show_posn_salience(attr[0])<jupyter_output><empty_output><jupyter_text>Results are similar but fewer positions are higlighted.  Verify vector attributionAttempt to verify that the attribution makes some sense, first by zeroing out the elements with the lowest attribution.<jupyter_code>def drop_low_salience(level, abs_attrib=True, posn=16):
  embeds = model.transformer.wte(tokens)
  attr = ig.attribute(embeds)
  attribv = torch.abs(attr[0,posn]) if abs_attrib else attr[0,posn]
  mask = attribv.ge(level)
  drop_masked(embeds, mask)

def drop_masked(embeds, mask, posn=16):
  embeds = model.transformer.wte(tokens)
  original = torch.clone(embeds[0,posn])
  fake = torch.where(mask, original,torch.zeros_like(original))
  embeds[0,posn]=fake
  transformer_outputs = model.transformer.forward(inputs_embeds=embeds)
  output_vec = transformer_outputs.last_hidden_state[0,posn]
  t_logits = model.lm_head(output_vec)
  t_predicted = torch.argmax(t_logits)
  print(f'Kept {torch.count_nonzero(fake)} of {torch.count_nonzero(original)} nonzero elements')

  transformer_outputs = model.transformer.forward(inputs_embeds=embeds)
  output_vec = transformer_outputs.last_hidden_state[0,posn]
  t_logits = model.lm_head(output_vec)
  t_predicted = torch.argmax(t_logits)
  print(f'Predicted token {t_predicted} "{tokenizer.decode(t_predicted)}" logit value: {t_logits[t_predicted]}')<jupyter_output><empty_output><jupyter_text>Drop Low Salience Elements<jupyter_code>ig = IntegratedGradients(make_logit_less_mean_embeds( 16, 17448 ))
drop_low_salience(0.005)
drop_low_salience(0.003)
drop_low_salience(0.002)
drop_low_salience(0.001)


ig = IntegratedGradients(make_forward_with_embeds( 16, 17448 ))

drop_low_salience(0.01)
drop_low_salience(0.008)
drop_low_salience(0.005)
drop_low_salience(0.003)<jupyter_output>Kept 663 of 768 nonzero elements
Predicted token 68 "e" logit value: -4.121610164642334
Kept 683 of 768 nonzero elements
Predicted token 12 "-" logit value: -4.719791412353516
Kept 713 of 768 nonzero elements
Predicted token 6193 " weather" logit value: 0.20048046112060547
Kept 730 of 768 nonzero elements
Predicted token 6193 " weather" logit value: -1.3550100326538086<jupyter_text>Drop Random ElementsCompare this to randomly dropping a similar number of elements<jupyter_code>def drop_random(rate):
  embeds = model.transformer.wte(tokens)
  mask = torch.rand(embeds[0,0].shape, device=embeds.device).ge(rate) 
  drop_masked(embeds, mask)

drop_random(31/768)
drop_random(20/768)
drop_random(10/768)
drop_random(5/768)
drop_random(1/768)<jupyter_output>Kept 727 of 768 nonzero elements
Predicted token 12 "-" logit value: -0.5765995979309082
Kept 746 of 768 nonzero elements
Predicted token 9 "*" logit value: -2.5034475326538086
Kept 758 of 768 nonzero elements
Predicted token 2465 " damage" logit value: -0.4099869728088379
Kept 764 of 768 nonzero elements
Predicted token 17448 " flooding" logit value: 0.0221405029296875
Kept 768 of 768 nonzero elements
Predicted token 17448 " f<PASSWORD>ing" logit value: -0.7155776023864746<jupyter_text>Drop Highly Salient ElementsInverting this, how many high salience tokens do we need to drop to affect the predicted token.<jupyter_code>def drop_high_salience(level, abs_attrib=True):
  embeds = model.transformer.wte(tokens)
  attr = ig.attribute(embeds)
  attribv = torch.abs(attr[0,16]) if abs_attrib else attr[0,16]
  mask = attribv.le(level)
  drop_masked(embeds, mask)

drop_high_salience(0.1)
drop_high_salience(0.2)
drop_high_salience(0.4)
drop_high_salience(0.4, False)<jupyter_output>Kept 437 of 768 nonzero elements
Predicted token 447 "�" logit value: -18.424091339111328
Kept 585 of 768 nonzero elements
Predicted token 447 "�" logit value: -22.182804107666016
Kept 709 of 768 nonzero elements
Predicted token 357 " (" logit value: -25.75557518005371
Kept 734 of 768 nonzero elements
Predicted token 564 " �" logit value: -23.998882293701172<jupyter_text>Result 

Integrated Gradients seems to correctly identify important dimensions in the input embedding vector.

Dropping low salience vector elements results in conceptually similar output tokens, and approx 30 can be dropped. Dropping 1 high salient element changes the output.
 Next, repeat with other tokens. Check the input token in different positions...<jupyter_code>print(tokens[0,13], tokenizer.decode(tokens[0,13]))
print(tokens[0,14], tokenizer.decode(tokens[0,14]))
print(tokens[0,15], tokenizer.decode(tokens[0,15]))
print(tokens[0,16], tokenizer.decode(tokens[0,16]))<jupyter_output>tensor(1085, device='cuda:0')  lead
tensor(284, device='cuda:0')  to
tensor(1957, device='cuda:0')  local
tensor(1417, device='cuda:0') ised<jupyter_text>And the predicted tokens.<jupyter_code>embeds = model.transformer.wte(tokens)
transformer_outputs = model.transformer.forward(inputs_embeds=embeds)
show_prediction(transformer_outputs,13)
show_prediction(transformer_outputs,14)
show_prediction(transformer_outputs,15)<jupyter_output><empty_output><jupyter_text>First look at input 'lead' predicting token 'to'.
Surprisingly this appears to depend little on the previous token, possibly suggesting a mistake in the calculation or interpretation. A less likely possiblity is that 'to' is a grammatical word and so depends on the overall sentence structure or frame rather than the specific previous word. As expected masking prevents any influence from the last input tokens.<jupyter_code>attr = attribute(IntegratedGradients(make_forward_with_embeds(13,284)))
show_posn_salience(attr[0])<jupyter_output><empty_output><jupyter_text>The other tokens look similer.  I see no obvious explanation or interpretation of these.<jupyter_code>attr = attribute(IntegratedGradients(make_forward_with_embeds(14,1957)))
show_posn_salience(attr[0])
attr = attribute(IntegratedGradients(make_forward_with_embeds(15,1417)))
show_posn_salience(attr[0])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.092150671855496, 'fraction_numerical': 0.04838828107790587, 'mean_word_length': 6.442622950819672, 'pattern_counts': {'":': 0, '<': 81, '<?xml version=': 0, '>': 87, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '21', 'dir': 'jupyter-structured-clean-dedup', 'id': '53997', 'license': 'MIT', 'n_tokens_mistral': 4889, 'n_tokens_neox': 4414, 'n_words': 1098, 'path': 'attribution-gpt-neo-125.ipynb', 'repo_name': 'stprior/explore-transformer', 'size': '13613'}